Technical Summary: Step-by-Step Optimization-like Reasoning in LLMs over Expanding Search Spaces

1. Problem Statement

While Large Language Models (LLMs) have demonstrated success in mathematical and coding reasoning via verifiable reward training, they remain brittle in constrained step-by-step decision-making tasks. In these scenarios, a locally plausible step may render the remaining problem infeasible or force a low-quality completion. Unlike standard answer-correctness settings where a single correct answer exists, many real-world tasks (e.g., resource allocation, scheduling, routing) involve many feasible solutions, where the objective is to construct a high-value solution under constraints.

Current scaling methods face two obstacles:

1. Generating diverse tasks that reliably require non-trivial reasoning is difficult.
2. Obtaining supervision for intermediate steps is expensive, as it requires identifying which partial actions lead to optimal outcomes.

The paper addresses the challenge of training LLMs to perform optimization-like reasoning over expanding search spaces, where the difficulty is controlled by a complexity parameter $\alpha$ without requiring new human labels.

2. Methodology: The OPT⋆ Framework

The authors introduce OPT⋆ (OPTimization-based Scalable Tasks for Auto-verifiable Reasoning), a family of tasks designed to train and evaluate step-by-step reasoning.

2.1 Task Definition

An OPT⋆ task instance $t$ is a constrained decision process defined by:

• Scalable Generation: Instances are generated by a procedure $Build_\alpha(\cdot)$ controlled by a complexity parameter $\alpha$ (e.g., number of cities in TSP, number of items in Knapsack). Increasing $\alpha$ expands the search space $M(t)$ but does not increase annotation cost.
• Cheap Verification: Each task exposes two inexpensive routines:
  • Feasibility Checker ($chkt$): Validates partial actions (e.g., ensuring load limits are not exceeded).
  • Outcome Evaluator ($evalt$): Scores complete terminal solutions (e.g., total tour length).
• Auto-verifiability: Both $chkt$ and $evalt$ are computable in polynomial time, even if finding the optimal solution is hard.

2.2 Training Regimes

The paper proposes two complementary training approaches based on solver availability:

A. Solver-Guided Online Policy Optimization (Online RL)

When a fast solver is available to compute the best completion value from a partial state:

• Mechanism: The solver acts as a value oracle. For a given state, candidate next steps are sampled, and each is scored by the best reachable terminal value under a fixed solver budget.
• Reward Shaping: The policy is updated using rank-based rewards (e.g., GRPO, GSPO, PPO) where actions leading to higher-value completions receive higher ranks.
• Goal: Reinforce steps that lead toward better final outcomes, effectively distilling solver heuristics into the policy.

B. Search-Based Offline RL

When solvers are unavailable or too expensive for partial states:

• Mechanism: The model generates candidate trajectories. A search algorithm (e.g., MCTS, Beam Search) explores the space, using $chkt$ to prune invalid paths and $evalt$ to score terminal nodes.
• Key Search Components:
  1. Feasibility Pruning: Discarding children where $chkt(s, a) = 0$.
  2. Duplicate Merging: Grouping distinct text outputs that map to the same structured action to avoid wasting search budget.
• Training: The model is fine-tuned (SFT) on the best trajectories discovered by the search process, effectively learning to mimic high-value reasoning traces.

2.3 Theoretical Foundation

The paper establishes an information-theoretic bound (Theorem A.3) relating success in large search spaces to the information extracted per unit of search budget. As the complexity $\alpha$ increases, the number of effective optima grows. To maintain a fixed success probability, a reasoner must either increase the search budget or extract more information per step. The effective branching factor ($b_{eff}$) is introduced as a metric to measure search efficiency; lower $b_{eff}$ indicates that a search budget places more mass on high-value terminal states.

3. Key Contributions

1. Conceptual (OPT⋆): Formalization of a scalable, auto-verifiable task family for optimization-like reasoning. This framework decouples reasoning difficulty from annotation effort by using complexity parameters to expand search spaces while keeping verification cheap.
2. Algorithmic:
   • Development of plug-and-play search components (feasibility checking and duplicate merging) that significantly improve search efficiency in offline settings.
   • Introduction of $\nu$PO (Solver-Guided Policy Optimization), a method that leverages solvers to provide step-level rewards for online RL, enabling training on partial states.
3. Empirical:
   • Demonstration that training on OPT⋆ improves step-by-step reasoning in optimization tasks.
   • Evidence that curriculum learning (training from easy $\alpha_1$ to hard $\alpha_4$) yields better generalization than training on hard tasks alone.
   • Showcasing that optimization training transfers to spatial reasoning capabilities (e.g., rotation, reflection, coverage) and improves performance on mathematical benchmarks (AMC, AIME).

4. Experimental Results

The authors evaluated the framework on tasks including Role Assignment, MaxSAT, Single-machine Scheduling, TSP, 2D Packing, and Quadratic Assignment Problem (QAP).

• Search Efficiency (Offline RL):
  • Enabling feasibility pruning and duplicate merging reduced the effective branching factor ($b_{eff}$) significantly. For example, in QAP, $\nu$MCTS achieved $b_{eff} \approx 24$ with these components, compared to $b_{eff} \approx 728$ for a solver-reference baseline that does not use LLM proposals.
  • These components were the primary drivers of search efficiency, increasing pass@16 from ~13% to ~85% in QAP.
• Model Improvement:
  • Models trained on OPT⋆ tasks (Role Assignment, MaxSAT) showed substantial improvements over base models. For instance, Qwen2.5-3B improved pass@1 on Role Assignment from 32.5% to 44.5% (Level 1) and 17.8% to 31.1% (Level 4).
  • Solver-guided online RL consistently outperformed domain-transfer baselines (e.g., training on Math or Code tasks did not transfer well to optimization tasks).
• Generalization:
  • Training on optimization tasks improved performance on out-of-distribution (OOD) spatial QA tasks (e.g., rotation, symmetry, overlap) and standard math benchmarks (AMC'23, AIME'24/25).
  • Curriculum effects: Mixed curriculum training ($\alpha_1 \to \alpha_4$) outperformed both "easy-only" and "hard-only" training strategies, particularly on harder levels and OOD tasks.

5. Significance and Claims

The paper claims that OPT⋆ provides a principled testbed for studying step-by-step reasoning under constraints, addressing the gap between verifiable reward training (effective for math/code) and real-world optimization tasks.

• Scientific Value: It offers a controlled environment to study how expanding search spaces affect reasoning capabilities, supported by an information-theoretic analysis showing that success requires extracting more information per budget unit as complexity grows.
• Practical Value: The proposed methods (plug-and-play search components and solver-guided RL) offer scalable ways to train models on complex workflows (scheduling, routing) where exact solvers may be unavailable or too slow for repeated calls.
• Modest Scope: The authors note that while solvers can verify or solve specific instances, the goal is to distill reusable heuristics into general-purpose models. The work focuses on synthetic optimization tasks motivated by traditional problems, acknowledging that future work could scale task generation using LLM priors to create even more diverse and complex scenarios.

In summary, the paper demonstrates that by leveraging auto-verifiable optimization tasks and complementary offline/online RL strategies, LLMs can learn to navigate expanding search spaces more effectively, improving both their optimization performance and their underlying spatial and logical reasoning skills.