Stream of Search (SoS): Learning to Search in Language

ArXiv: 2404.03683

🎯 Pitch

Training language models on messy search trajectories—including backtracking and dead ends—boosts problem-solving accuracy by 25 percentage points over models that only see perfect solutions. When these models then self-improve through reinforcement learning, they solve over a third of problems that stumped the original symbolic solvers used to generate the training data, effectively discovering novel search strategies.

---

1. Executive Summary

This paper introduces the Stream of Search (SoS) framework, which teaches language models to search and backtrack by training them on the full serialized process of exploration — including mistakes, recovery, and suboptimal paths — rather than on clean optimal trajectories alone. Using the game of Countdown (a high-branching-factor arithmetic search problem) and a 250M-parameter GPT-Neo model trained from scratch, the authors demonstrate that SoS pretraining boosts accuracy by roughly 25 percentage points over a model trained only on optimal paths (51.27% vs. 25.73% on held-out inputs). When further self-improved with Advantage-Induced Policy Alignment (APA, an actor-critic RL method that uses a separate value network with periodic reference-policy resets) and Self-Taught Reasoner (STaR, expert iteration that filters model-generated trajectories for correctness and retrains on them), the SoS models solve 36% of problems that were unsolvable by any symbolic heuristic solver in the training set, establishing that LMs can discover novel search strategies when exposed to the messy process of search — but only when that process is represented explicitly in training data.

2. Context and Motivation

The Core Problem: Language Models Don't Learn from Mistakes

The fundamental problem this paper tackles is that language models, as typically trained, are never shown what recovery from error looks like. Training data for reasoning tasks — math problem solutions, code generation, logical derivations — almost universally consists of clean, correct final outputs. A student who only ever saw polished proofs and never witnessed a mathematician getting stuck, trying a dead-end approach, or backtracking from a false lemma would develop a deeply skewed understanding of how problem-solving actually works. This paper argues that LMs suffer from exactly this skew.

The consequence is not merely that LMs fail to produce interesting intermediate mistakes (though they do). The consequences are structural and severe:

• Snowballing errors (Ross et al., 2011): When an autoregressive model produces a single incorrect token in a multi-step reasoning chain, that error compounds. The model has no mechanism for recognizing the mistake, nor any learned behavior for backing up and trying an alternative. Each subsequent token is conditioned on the incorrect prefix, driving the trajectory further from correctness. The paper frames this as an exposure bias problem: at training time, the model always sees ground-truth prefixes; at test time, it sees its own generated prefixes, which drift from the training distribution.

• Lookahead failure (LeCun, 2023; Bachmann & Nagarajan, 2024): Many reasoning problems require considering the downstream consequences of a local decision before committing — what the RL literature calls credit assignment (Sutton & Barto, 2018). Standard next-token prediction offers no mechanism for this. The model predicts one token at a time with no explicit representation of future states or alternative branches. The paper argues that search — the ability to mentally explore multiple paths before selecting one — is the missing capability, and that standard LM training never cultivates it.

Both of these failures trace to the same root: the training data reflects only the outcome of decision-making, not the process. The paper's core insight is that search is not just a test-time strategy to be bolted onto a pretrained model, but a capability that can and should be learned during training — provided the training data includes search itself.

Why This Matters: Beyond Inference-Time Scaffolding

The practical importance of this problem has grown with the scale of LM deployment. If LMs can only produce correct reasoning in one forward pass, their effective reasoning horizon is bounded by what can be predicted in a single autoregressive trajectory. Real-world problems — planning, scientific discovery, mathematical proof, complex code debugging — routinely exceed this horizon. The question is how to extend it.

Existing approaches have largely converged on one answer: keep the LM as a next-token predictor, but wrap it in an external search system at inference time. This includes methods like Tree of Thoughts (Yao et al., 2024), where a symbolic search algorithm (BFS/DFS) calls the LM to generate candidate successor states and evaluate them, and Graph of Thoughts (Besta et al., 2023), which generalizes this to graph structures. These "extrinsic" methods improve performance but leave the LM itself unchanged — it never learns to search. The search strategy is fixed by the external algorithm, the computational cost scales poorly (each step requires multiple LM calls, as Sel et al., 2023 note), and the LM cannot discover new strategies or adapt its search behavior based on problem characteristics.

The paper identifies a more ambitious goal: making search an intrinsic capability of the LM, learned during training and executed autonomously at inference time. If achievable, this would mean:

• Self-improvement during training (Silver et al., 2018): An LM that can search during training can generate better training data for itself — exploring solution spaces, finding novel reasoning paths, and distilling those discoveries back into improved parameters. This is the loop that produced AlphaZero's superhuman performance, and the paper explicitly draws this analogy in positioning SoS.

• Computational efficiency at inference: An intrinsic search policy avoids the LM-call overhead of extrinsic methods. Rather than the LM being called dozens of times per step by an external search algorithm, the LM executes the entire search in a single autoregressive generation, interleaving exploration, backtracking, and pruning as it writes.

• Strategy discovery: Unlike extrinsic methods constrained by their hardcoded search algorithms, a model trained on diverse search trajectories could combine, adapt, or even transcend the strategies in its training data. The paper's finding that SoS models solve problems unsolvable by any of their training heuristics is direct evidence for this possibility.

Where Prior Approaches Fall Short

The paper positions itself against three lines of prior work, each with specific limitations:

1. Extrinsic search (Tree of Thoughts, Graph of Thoughts, and related methods). These approaches use the LM as a generative module and state evaluator within a search algorithm that exists outside the model. The search strategy — BFS, DFS, beam search — is fixed by the system designer. Shortcomings:

• Inference cost: Each node expansion requires at least one (often multiple) LM forward passes. For problems with high branching factors like Countdown (where each state has $\binom{N}{2} \times 4$ possible children for $N$ remaining numbers), the number of LM calls grows combinatorially with search depth.
• No learning transfer: The LM does not improve at search through this process. Its parameters remain frozen. Any efficiency gains from better search strategies are lost on the next problem.
• Rigidity: The LM cannot adapt the search strategy to the problem. It cannot decide, mid-search, that a depth-first approach is stalling and breadth-first would be better. The fixed algorithm makes these meta-decisions blindly.

2. In-context demonstrations of search (Gandhi et al., 2023; Sel et al., 2023). Instead of an external algorithm, the search procedure is demonstrated through examples in the prompt — for instance, showing the LM a BFS trace and asking it to continue in the same style. While this internalizes search to the LM's forward pass (avoiding the extrinsic call overhead), the approach has critical limits:

• Prompt-bound strategies: The demonstrated search procedure is fixed in the prompt. The LM cannot generalize beyond the specific strategy shown, nor can it discover better ones. If the prompt shows BFS with a particular heuristic, the LM mimics BFS with that heuristic.
• Context-length constraints: Complex search traces run into prompt-length limits. The paper notes that Countdown trajectories for 5-input problems can reach 60,000 tokens — far exceeding standard context windows.
• No parametric learning: The search capability lives entirely in the prompt, not in the model's weights. It does not compound across problems or improve with experience.

3. Process supervision (Lightman et al., 2023). Training separate verifier models to score intermediate reasoning steps, then using those scores to guide generation. The limitation here is practical rather than conceptual:

• Annotation cost: Lightman et al. required human-generated labels for each intermediate step of each training solution — a massive annotation effort. The paper notes this "may not scale well for other, more complex domains."
• Separation of search and generation: The verifier evaluates steps; the LM still generates them one at a time with no explicit search representation. The model never learns to backtrack — it only learns to produce steps that the verifier approves.

How This Paper Positions Itself

The SoS framework attempts to unify the strengths of these approaches while addressing their weaknesses. The key moves are:

From extrinsic to intrinsic search. Like in-context methods, SoS represents the search process entirely within the LM's autoregressive generation — a single forward pass produces the full search trajectory, including exploration, heuristic evaluation, backtracking, and goal-checking. But unlike in-context methods, this capability is trained into the model's parameters, allowing it to transfer across problems and improve through policy optimization.

From optimal paths to search trajectories. The training data is not the cleaned-up solution but the raw search trace — including states that were explored but didn't lead to a solution, heuristic evaluations that were incorrect, and the specific backtracking operations that moved between branches. This directly addresses the snowballing problem: the model sees, repeatedly during training, what it looks like to recognize a dead-end and back out of it. It learns that generating a mistake is not fatal — it is recoverable through search.

From fixed heuristics to learned, improvable ones. The training data includes trajectories from multiple search strategies (BFS and DFS with two different heuristics, across various hyperparameter settings). The model is not trained to mimic any single strategy. As the alignment analyses in Section 5 show (Figure 3c), the trained SoS model does not correlate strongly with any one symbolic strategy in its state-visitation patterns. It appears to learn a more general search capability that can flexibly utilize or combine strategies. Furthermore, through APA and STaR fine-tuning (Section 6), the model can improve beyond the strategies in its training data — solving problems that none of the symbolic strategies could solve (Figure 5c).

Internal world model. A subtle but important positioning: unlike extrinsic methods that rely on a ground-truth environment model (the symbolic search algorithm knows exactly what states are reachable and evaluates transitions exactly), the SoS LM must simulate state transitions itself. When the model writes "3 + 5 = 8" as a state transition, it is both proposing the action and computing its result. Errors in this computation (arithmetic mistakes, invalid operations) are part of the search trajectory and the model learns to recover from them. The paper frames this as learning an "internal world model" for the search domain (Section 7), connecting to work on learned models in RL (Schrittwieser et al., 2020).

Language as the representation for search. The paper devotes Section 3 to formalizing "a language for search" — a vocabulary of primitive operations (current state, goal state, state expansion, exploration choice, pruning, backtracking, goal check, heuristic evaluation) that can be composed into diverse search strategies. This is not merely a notational convenience. By representing search operations in the same token space as the problem content, the paper makes search itself a sequence modeling problem. The LM can learn to interleave content generation (arithmetic operations on numbers) with meta-cognitive operations (deciding which frontier node to explore next, checking whether the goal has been reached, backtracking to a previous branch point). This unification is what enables the full search process to be captured in a single autoregressive generation.

Self-improvement as the target, not just assisted reasoning. The paper explicitly cites Silver et al. (2018) on AlphaZero as the aspirational reference point: the most consequential outcome of learning to search is not better test-time performance, but the ability to use search during training to generate improved data for further training. The APA and STaR experiments (Section 6) are a first step in this direction — the model generates trajectories, filters for correctness, and retrains on the improved data, closing the self-improvement loop. The fact that this loop solves problems unsolvable by the original training heuristics (Figure 5b-c) is the paper's strongest evidence that search-trained LMs can surpass their teachers.

The Countdown Game as a Testbed

The choice of Countdown is not arbitrary. The paper needs a domain where:

• Search is genuinely necessary, not just helpful. Countdown's branching factor of $\binom{N}{2} \times 4$ per depth means exhaustive search is infeasible — for 4 input numbers, the first step has $\binom{4}{2} \times 4 = 24$ possible children. Heuristic guidance is required.
• Problems have variable difficulty, so that search traces range from short (a few steps of exploration) to very long (thousands of steps), testing the model's ability to maintain coherence over extended search.
• The solution space is structured, so that search quality can be evaluated objectively (a solution either reaches the target or doesn't) and heuristics can be defined cleanly (the sum heuristic, the multiply/factors heuristic).
• Training data can be generated automatically, via symbolic search algorithms that produce both correct and incorrect trajectories without human annotation.

Countdown satisfies all of these. It is a deliberately simplified domain — the paper's goal is not to achieve state-of-the-art on a benchmark but to demonstrate a capability (learned search) in a controlled setting where the mechanisms can be analyzed. The extension to more complex, real-world tasks is left explicitly to future work (Section 7).

3. Technical Approach

This is primarily a representation learning and policy optimization paper whose core idea is that language models can learn to perform search intrinsically — as an autoregressive sequence generation task — if the training data explicitly includes the full process of exploration, including mistakes, backtracking, and heuristic evaluation, serialized into a common textual format called a Stream of Search (SoS).

3.1 Reader Orientation

The paper builds a system for teaching a transformer-based language model to solve combinatorial search problems by generating the entire search process as a string, rather than generating only the final answer or using an external search engine to guide it. The system solves the problem of LMs' inability to recover from errors during reasoning by exposing them, during training, to trajectories that contain exploration of dead ends, explicit backtracking operations, and multiple search strategies — essentially teaching the model that "to err is fine, to backtrack is learnable." The solution takes the shape of a two-stage pipeline: first, pretraining from scratch on a synthetic dataset of serialized search traces generated by diverse symbolic solvers; second, self-improvement through reinforcement learning techniques that optimize the model's own generated trajectories for correctness and efficiency.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components:

1. Countdown Problem Generator — produces problem instances (input numbers + target number) and serves as the environment against which all search is performed.
2. Symbolic Search Strategy Suite — a collection of 12 heuristic-guided search algorithms (combinations of BFS/DFS with two heuristic functions and various hyperparameters) that generate the initial training data by actually searching the problem space and recording their trajectories as serialized text.
3. SoS Serialization Format — a unified language for representing search operations (state exploration, backtracking, goal checking, heuristic evaluation, frontier management) as a linear sequence of tokens, enabling any search algorithm's execution trace to be flattened into a single string.
4. GPT-Neo Base Model (250M parameters) — a transformer trained from scratch with a causal language modeling objective on the serialized search trajectories. It learns to simulate both the arithmetic operations of the environment and the meta-cognitive operations of search (deciding which state to explore next, when to backtrack, when to declare a goal reached) in a single autoregressive pass.
5. Policy Improvement Module (APA or STaR) — a reinforcement learning stage that takes the pretrained SoS model, generates new trajectories on training problems, filters or reweights them based on correctness and efficiency, and fine-tunes the model on this improved data, closing a self-improvement loop.

Information flows as follows: the problem generator creates Countdown problems → symbolic strategies execute search and serialize traces → the SoS model pretrains on these traces → the model generates new trajectories on training problems → a filtering/reweighting step (STaR: keep only correct trajectories; APA: use advantage-weighted regression with a learned value function) selects high-quality data → the model fine-tunes on this data → the cycle repeats until validation accuracy converges.

3.3 Roadmap for the Deep Dive

• First, the formalization of search as a Markov Decision Process (Section 3 of the paper), because every subsequent component — the serialization language, the training data generation, the policy improvement objectives — builds on this abstraction.
• Second, the language for search (Section 3), which defines the vocabulary of primitive operations and specifies which are made explicit versus implicit in the serialized format. This is the representational innovation that makes the entire approach possible.
• Third, the Countdown problem environment (Section 4), including the branching factor, the heuristic functions, and the dataset construction — since the dataset's composition and the specific heuristics used determine what the model can learn.
• Fourth, the base SoS model training setup (Section 5), including the comparison with optimal-path training and the alignment metrics used to assess what strategies the model learned.
• Fifth, the policy improvement methods (Section 6): STaR (expert iteration) and APA (advantage-induced policy alignment), with detailed algorithm descriptions, reward functions, and training procedures.

3.4 Detailed, Sentence-Based Technical Breakdown

Formalizing Search as a Markov Decision Process (MDP)

The paper models any search problem as a Markov Decision Process with four components. The state space $\mathcal{S}$ contains all possible configurations of the problem — in Countdown, a state is a set of remaining numbers and the operations applied so far. The action space $\mathcal{A}$ contains the operations the solver can perform on states or to transition between them — in Countdown, an action is selecting two numbers and an arithmetic operation to combine them. The transition function $T: \mathcal{S} \times \mathcal{A} \to \mathcal{S}$ defines the deterministic outcome of applying an action to a state — for example, taking $\{3, 5, 2\}$ and applying the action "add 3 and 5" produces the new state $\{8, 2\}$. The reward function $R: \mathcal{S} \to \mathbb{R}$ assigns a positive reward only upon reaching the goal state $s_g$.

A search problem instance consists of an initial state $s_0$ (the input numbers and target) and a goal state $s_g$ (a state containing only the target number). The search tree is implicitly defined by recursively applying all valid actions from $s_0$ to generate child states, then from each child to generate grandchildren, continuing until leaf states (states with no valid actions remaining). A correct solution path $\mathcal{P}$ is a sequence of state-action pairs:

$(s_0, a_0, s_1, a_1, \dots, s_{g-1}, a_{g-1}, s_g)$

where each successive state $s_{i+1}$ is obtained by applying the valid action $a_i$ to state $s_i$ — that is, $s_{i+1} = T(s_i, a_i)$ — and the final state $s_g$ equals the goal state.

What it computes: the formal structure of what "searching" means in this domain. The initial state is the problem as given. The transition function defines what moves are legal. The solution path is the specific sequence of legal moves that transforms the initial state into the goal state. The search problem is to find this path within the exponentially large implicit tree.

Why this form: casting search as an MDP lets the paper draw on the standard RL vocabulary of states, actions, transitions, and rewards. This matters because the language for search (next section) will need operations that correspond to navigating an MDP — exploring child states, backtracking to parent states, checking whether the current state is the goal. The MDP formulation also makes the connection to policy improvement methods (APA) natural, since those methods are designed for MDPs. The path itself, $\mathcal{P}$, is what the "optimal paths" baseline trains on — it is the clean, mistake-free sequence that traditional LM training would use.

The Language for Search: A Vocabulary of Primitive Operations

This is the paper's core representational contribution. Instead of treating search as an external procedure that calls the LM, the paper defines a token-level vocabulary that can express the internal state of any search algorithm as a linear string. Each operation corresponds to something a search algorithm does, but serialized into natural language tokens that a transformer can process autoregressively.

The vocabulary consists of nine operations, each of which can be either explicit (written out as text in the trajectory) or implicit (performed by the model's internal computation without being verbalized). The paper makes deliberate choices about which to make explicit:

Explicit operations (verbalized in the trajectory):

• Current State $s_c$: The state currently being explored. Written as the set of remaining numbers and the operations applied so far. This is made explicit because the LM needs to track which state it is currently examining — it cannot maintain this purely in hidden states over long trajectories without textual grounding.
• Goal State $s_g$: The target to reach. Written as the target number. Made explicit so the model can compare current state to goal.
• Backtracking: An operation that moves from the current node to a previously explored node. Written as a statement indicating "moving back to state X." Made explicit because backtracking is precisely the behavior the model needs to learn — recognizing dead ends and returning to earlier branch points.
• Goal Check: Comparing $s_c$ to $s_g$. Written as a verification statement. Made explicit so the model learns to check whether it has solved the problem.
• Exploration Choice: Deciding which state to expand next from the frontier. Written by selecting and writing out the next state to explore. Made explicit because different search strategies differ primarily in their exploration order (BFS vs. DFS vs. heuristic-guided).

Implicit operations (not verbalized, performed internally by the model):

• State Queue $S_q$: The frontier of unexplored states. The model must track which states have been generated but not yet expanded. Kept implicit — the model learns to maintain this in its internal representations rather than listing the full queue, which would consume excessive context length.
• State Expansion Function $SE: \mathcal{S} \to \mathcal{S}^n$: The operation that generates all valid child states of the current state by applying all legal arithmetic combinations. Kept implicit — the model generates children one at a time as it writes them, rather than calling an external simulator.
• Pruning: Discarding states or subtrees unlikely to lead to a solution. Kept implicit — the model learns to avoid exploring certain branches without explicitly stating that it is pruning them.
• Heuristic: A function $h: \mathcal{S} \times \mathcal{S} \to \mathbb{R}$ that estimates the distance from the current state to the goal. Kept implicit — the model internalizes heuristic evaluation as part of its learned representations rather than writing out numerical scores.

The design rationale for explicit vs. implicit is practical. Operations that are made explicit become reasoning moves the LM produces as text — the model can inspect them, and more importantly, the policy improvement methods can optimize over them. Operations that are kept implicit become abstract representations the model learns internally, which makes training more flexible (the model can discover better heuristics than the ones in the training data) but also makes them harder to analyze or directly optimize. The paper keeps the heuristic, pruning, and state queue implicit specifically to allow the model to develop its own strategies beyond the training heuristics — a choice validated by the alignment analyses (Section 5, Figure 3c) showing the trained model does not simply mimic any one symbolic strategy.

How the language serializes a search tree (Figure 2). The paper illustrates the serialization with a diagram. On the left is a tree representation: a root node $s_0$, several explored child nodes (some colored to show they were visited, some leading to a solution shown in green), and many unexplored nodes (black circles). The search trajectory $\mathcal{T}$ is the order in which the search algorithm visited these nodes — it includes the branch that leads to the solution but also the dead-end branches that were explored and abandoned, and the backtracking operations that moved between them. In the center, this tree is flattened into a linear string: the text writes out the state at node A, then its child B, then recognizes B is a dead end and writes a backtracking statement, then moves to node C, explores its children, and so on until the goal is found. On the right is the optimal path $\mathcal{P}$ — only the sequence of states that leads directly from $s_0$ to $s_g$, with no exploration, no dead ends, no backtracking. This is what the OP baseline trains on. The distinction between $\mathcal{T}$ (full search trace) and $\mathcal{P}$ (clean solution path) is the crucial training data difference between SoS and OP.

The Countdown Problem Environment and Heuristic Functions

Countdown is a search problem where a set of input numbers must be combined using the four basic arithmetic operations (addition, subtraction, multiplication, division) to reach a target number. The paper constrains problems to 4 input numbers to keep search traces within the 4096-token context window of the GPT-Neo model — the authors note that 5-input problems can produce traces up to 60,000 tokens long. Targets range from 10 to 100.

Branching factor. At each state with $N$ remaining numbers, the number of possible actions is:

$\binom{N}{2} \times 4$

where $\binom{N}{2}$ is the number of ways to choose two numbers from the $N$ available, and $\times 4$ accounts for the four arithmetic operations (+, -, ×, ÷). For the initial state with 4 numbers, this gives $\binom{4}{2} \times 4 = 6 \times 4 = 24$ possible child states. After one operation, 3 numbers remain, giving $\binom{3}{2} \times 4 = 3 \times 4 = 12$ children. After two operations, 2 numbers remain, giving $\binom{2}{2} \times 4 = 1 \times 4 = 4$ children. The total search space is therefore approximately $24 \times 12 \times 4 = 1152$ leaf nodes, small enough to exhaustively enumerate but large enough that naive random search is inefficient — heuristic guidance matters.

Why this matters: the branching factor determines how much search is necessary. If the branching factor were 2-3 at each step, depth-first search with backtracking would be overkill — you could just sample a few random paths. The combinatorial growth makes search strategies genuinely necessary, creating the conditions where learning to search provides real benefit over guessing.

Heuristic functions used for dataset generation. The paper defines two simple, interpretable heuristic functions to guide the symbolic search strategies that generate training data. Both use the sum of the remaining input numbers as a proxy for the current state's value:

Let $I = [n_1, n_2, \dots, n_k]$ be the input numbers at a given state, and let $T$ be the target number.

Sum heuristic $h_{sum}$: The absolute difference between the sum of all remaining inputs and the target:

$h_{sum}(I, T) = \left|T - \sum_{i=1}^{i=k} n_i\right|$

where $k$ is the number of remaining input numbers, $T$ is the target, and $n_i$ are the individual remaining numbers.

What it computes: the absolute distance between the sum of remaining numbers and the target. If the inputs are $\{20, 5, 3\}$ and the target is 28, then $h_{sum} = |28 - (20+5+3)| = |28 - 28| = 0$. This would be a perfect score, indicating that the target can be reached by combining all remaining numbers through addition. If the inputs sum to 50 and the target is 30, $h_{sum} = 20$, indicating that substantial subtraction or division is needed.

Why this form: the sum heuristic is a simple lower-bound proxy. If $h_{sum}$ is zero, the numbers can potentially be combined into the target (though not necessarily — the specific values matter for non-commutative operations). Larger values indicate that the current state is numerically far from the target, suggesting it may be a poor branch. The absolute value is used because both overshooting and undershooting are equally problematic. This heuristic ignores which arithmetic operations would be needed and whether they're valid — it's purely a numerical distance measure.

Multiply/factors heuristic $h_{multiply}$: The minimum distance between the sum of remaining inputs and any factor of the target:

$h_{multiply}(I, T) = \min\left(\left\{\left|f_j - \sum_{i=1}^{i=k} n_i\right| \quad \forall j \in [1, m]\right\}\right)$

where $f_1, f_2, \dots, f_m$ are the factors of the target $T$, $m$ is the number of factors, and the minimum is taken over all factors $f_j$.

What it computes: the smallest distance between the sum of remaining numbers and any number that divides the target evenly. For example, if the target is 24 (factors: 1, 2, 3, 4, 6, 8, 12, 24) and the remaining numbers sum to 11, then the closest factor is 12 (distance 1) or 8 (distance 3), so $h_{multiply} = 1$. This heuristic captures the intuition that if you can combine numbers to reach a factor of the target, the remaining operation is simple multiplication or division.

Why this form: many Countdown solutions involve reaching an intermediate number that is a factor of the target, then multiplying by the appropriate value. This heuristic explicitly rewards states whose sum is close to a factor, offering a complementary signal to the sum heuristic. The minimum is taken because any factor could potentially be the right intermediate target — closeness to the nearest factor is what matters. The paper uses both heuristics because they encode different problem-solving intuitions (additive vs. multiplicative reasoning), and training on trajectories guided by both exposes the model to diverse search biases.

Dataset generation procedure. The paper constructs 12 distinct search strategies by combining:

• Two search algorithms: DFS (Algorithm 3) and BFS (Algorithm 4)
• Two heuristic functions: $h_{sum}$ and $h_{multiply}$
• For DFS: a threshold $h_{th} = T$ (the target value), meaning states with heuristic value above $T$ are pruned
• For BFS: five breadth limits $b \in \{1, 2, 3, 4, 5\}$, meaning only the top $b$ child states (by heuristic score) are kept at each expansion

This yields $2_{algorithms} \times 2_{heuristics} \times (1_{DFS\_setting} + 5_{BFS\_settings}) = 2 \times 2 \times 6 = 24$ possible strategy configurations, but the paper mentions "12 search strategies" (Section 4) — the precise combination is that DFS uses one threshold per heuristic (2 strategies total for DFS), and BFS uses 5 breadth limits × 2 heuristics (10 strategies for BFS), summing to 12.

Each strategy is run on each Countdown problem. The algorithm records its search trajectory as a serialized string following the SoS language format. The trajectory includes every state visited, every backtracking operation, every goal check, and the order of exploration — whether the search eventually finds the solution or not.

Final dataset composition. The SoS dataset contains 500,000 search trajectories. Of these, only 285,501 (57.1%) lead to a correct solution. The remaining ~43% are trajectories where the search algorithm explored the tree but failed to find the correct path. These incorrect trajectories are included in the training data — the model sees what a failed search looks like as well as a successful one. The optimal paths (OP) dataset, by contrast, contains only the 500,000 correct solution paths $\mathcal{P}$, stripped of all exploration, dead ends, and backtracking.

Why include incorrect trajectories: the paper's hypothesis is that seeing failed searches is essential for learning to search. If the model only sees correct paths, it never learns what a dead end looks like, nor what action to take when it encounters one. By training on both successful and unsuccessful searches, the model learns three things simultaneously: what good states look like (from successful trajectories), what bad states look like and that they lead to dead ends (from unsuccessful trajectories), and most importantly, that backtracking is the appropriate response to a dead end — because every trajectory, successful or not, includes backtracking operations when a branch fails.

Base Model Training: SoS Pretraining vs. Optimal Paths Baseline

Model architecture. The paper trains a GPT-Neo model (Gao et al., 2020) with 250 million parameters from scratch — not fine-tuning a pretrained LM. The architecture specifications (Appendix C.1):

• 16 transformer layers, all using global (not local) attention
• 16 attention heads
• Hidden size of 1024
• Context length of 4096 tokens
• Dropout rate of 0.1
• bf16 precision
• Flash Attention 2 (Dao et al., 2022) for memory efficiency

Why train from scratch rather than fine-tune: the authors want to isolate the effect of the training data (SoS vs. OP) from any search-relevant capabilities the model might have acquired during general-purpose pretraining. Training from scratch on only Countdown search traces ensures that any search behavior observed is learned entirely from the SoS format, not from analogical transfer from natural language data.

Training objective. Both the SoS model and the OP baseline are trained with a standard causal language modeling objective — next-token prediction with cross-entropy loss. There is no special search loss, no auxiliary heuristic prediction head, and no reinforcement learning at this stage. The model simply learns to maximize the probability of the next token in the training sequences.

Training hyperparameters (Appendix C.2 and C.3, identical for both models):

• Training examples: 500,000
• Batch size: 24 per GPU, gradient accumulation steps: 1 (effective batch size depends on the 4 GPUs: $24 \times 4 = 96$ if using 4 GPUs)
• Training steps: 50,000
• Optimizer: AdamW with weight decay 0.01
• Learning rate: $1 \times 10^{-5}$, with cosine scheduler and 1 warmup step
• Evaluation steps: every 500 steps, model selection based on accuracy on a validation set of 1,000 problems

The critical difference between conditions. The SoS model is trained on search trajectories $\mathcal{T}$, which include exploration, dead ends, backtracking, and potentially no solution. The OP model is trained on optimal paths $\mathcal{P}$, which contain only the correct sequence of state transitions from start to goal. Crucially, the OP model sees 500,000 correct solutions (100% correct), while the SoS model sees only ~285,000 correct solutions (57% correct) embedded within longer trajectories that include errors. The SoS model thus has fewer correct examples but more total information about the search process.

Why this comparison is fair despite unequal correct examples: the paper argues explicitly that both models are trained for the same number of gradient steps. If the SoS trajectories are longer (they include dead ends), this means the SoS model processes more tokens per training step, giving it more total training signal. The paper runs both models for exactly 50,000 steps, controlling for optimization budget rather than example count. This is a deliberate choice: in practice, you can always train longer, but the question is whether the SoS format yields better performance per unit of compute. The fact that SoS achieves 51.27% vs. 25.73% despite having fewer correct examples in its training data is the key result — it demonstrates that the information content of search traces (learning from mistakes) is more valuable than the correctness density of optimal traces (learning only from successes).

Generation at test time. At inference, the model is given a Countdown problem as a prompt (the initial state) and generates autoregressively until it produces a "Goal Reached" token or hits the maximum context length. The generated text is a complete search trajectory — the model decides which states to explore, when to backtrack, when to check for the goal. The paper does not use any external verifier or search algorithm at test time. The entire search is executed by the LM's autoregressive generation.

Alignment metrics. To quantitatively assess what search strategies the trained models use, the paper defines two alignment measures (Section 4):

Alignment of Correctness: the Pearson correlation between two strategies' binary solve/fail patterns across problems.

$r(\text{strategy}_A, \text{strategy}_B) = \text{Pearson}(\mathbf{c}_A, \mathbf{c}_B)$

where $\mathbf{c}_A$ and $\mathbf{c}_B$ are binary vectors of length equal to the number of test problems, with 1 indicating the strategy solved that problem and 0 indicating it did not.

What it computes: whether two strategies tend to succeed and fail on the same problems. High correlation means the strategies have similar strengths and weaknesses; low correlation means they are complementary — one solves problems the other cannot.

Why this form: Pearson correlation captures linear association between binary vectors and is standard for measuring agreement. It does not require the strategies to visit the same states — only to agree on which problems are solvable.

Alignment of States Visited: the normalized overlap in states visited by two search trajectories on the same problem.

$\text{State Alignment}(\mathcal{T}_1, \mathcal{T}_2) = \frac{|\mathcal{T}_1 \cap \mathcal{T}_2|}{\max(|\mathcal{T}_1|, |\mathcal{T}_2|)}$

where $\mathcal{T}_1$ and $\mathcal{T}_2$ are the sets of states visited in two trajectories, $|\mathcal{T}_1 \cap \mathcal{T}_2|$ is the number of states common to both, and $\max(|\mathcal{T}_1|, |\mathcal{T}_2|)$ is the length of the longer trajectory (used for normalization).

What it computes: the fraction of states in the longer trajectory that were also visited by the shorter trajectory. A score of 1.0 means one trajectory's visited states are a subset of the other's. A score of 0.0 means they explored completely disjoint parts of the search tree.

Why this form: Jaccard similarity ($|A \cap B| / |A \cup B|$) would penalize trajectories of different lengths, since the denominator grows with the size of the union. The max-normalization used here is asymmetric — it measures overlap relative to the more exploratory trajectory, which is appropriate for comparing search strategies where one may be more thorough than another. The authors compute this per-problem and then average across problems.

These alignment metrics are not used for training — they are diagnostic tools to assess whether the trained model is copying a single symbolic strategy or combining elements from multiple ones.

Policy Improvement Method 1: Expert Iteration with STaR

The first self-improvement method is Self-Taught Reasoner (STaR, Zelikman et al., 2022), an expert iteration approach. The core idea is: use the current model to generate trajectories on training problems, filter to keep only the correct trajectories, and fine-tune the model on this filtered data. Iterating this process creates a positive feedback loop where the model generates better trajectories, which become better training data, which produces a better model.

Algorithm (Algorithm 1 in the paper):

1. Initialize: $M_0 \leftarrow \text{SoS}$ — start from the pretrained SoS model.
2. For each iteration $n = 1, \dots, N$:
   • Generate: For each training problem $x_i$, generate a trajectory $\mathcal{T}_i \leftarrow M_{n-1}(x_i)$ using the current model with temperature 0.8 (to encourage exploration and diverse trajectories).
   • Filter: Construct dataset $D_n = \{(x_i, \mathcal{T}_i) \mid i \in [1, m] \text{ s.t. } \mathcal{P}_i \in \mathcal{T}_i\}$ — keep only the $(x_i, \mathcal{T}_i)$ pairs where the correct solution path $\mathcal{P}_i$ is present in the generated trajectory $\mathcal{T}_i$. This discards all trajectories that failed to solve the problem.
   • Fine-tune: $M_n \leftarrow \text{train}(M_0, D_n)$ — fine-tune from the original SoS checkpoint $M_0$ (not from $M_{n-1}$) on the filtered dataset $D_n$. This reset to the base model is crucial, following the original STaR procedure — it prevents the model from drifting too far from its initial distribution and overfitting to idiosyncrasies of its own earlier generations.
3. Stop when validation accuracy converges (3 iterations in the paper).

Hyperparameters for STaR fine-tuning (Appendix C.4):

• Training examples: 100,000 (correct trajectories filtered from a larger set of generations)
• Temperature for sampling: 0.8
• Batch size: 24, gradient accumulation: 1
• Training steps: 20,000 per iteration
• Learning rate: $1 \times 10^{-5}$ with cosine scheduler and 100 warmup steps
• Weight decay: 0.01
• Model reset: each iteration starts from $M_0$, the base SoS model

Why reset to $M_0$ each iteration: fine-tuning on model-generated data creates a distributional shift. The model's generated trajectories differ from the training data, and if you iteratively fine-tune without resetting, the model compounds its own biases — it learns to generate trajectories that look like what it already generates, rather than trajectories that are correct. Resetting to the original checkpoint provides a stable anchor and ensures that each fine-tuning iteration learns from the base model's capabilities improved by fresh, diverse search experiences rather than from a progressively narrowing distribution.

Why temperature 0.8: deterministic generation (temperature 0) would produce the same trajectory every time, providing no diversity for the filter step. Temperature 1.0 would maximize diversity but reduce the proportion of correct trajectories (making the filter step discard more data). Temperature 0.8 is an intermediate choice that balances exploration with quality. The paper does not report ablations on this value.

What the filter step achieves: it implements a simple form of outcome-based reward: trajectories are kept if they contain the solution, discarded otherwise. There is no intermediate reward for efficient search or penalization for long trajectories — only binary correctness matters. This means STaR optimizes purely for solving more problems, not for solving them faster or with fewer explored states (though the paper reports in Figure 6 that efficiency improves as a side effect).

Policy Improvement Method 2: Advantage-Induced Policy Alignment (APA)

The second self-improvement method is Advantage-Induced Policy Alignment (APA, Zhu et al., 2023), an actor-critic reinforcement learning technique adapted for language model fine-tuning. Unlike STaR, which uses binary filtering, APA uses a learned value function to assign continuous advantage weights to actions, enabling the model to learn from both successful and unsuccessful trajectories — good actions in failed trajectories can still receive positive reinforcement.

Algorithm (Algorithm 2 in the paper):

1. Initialize three copies of the SoS model:

   • $\pi_{init}$: the base pretrained SoS model
   • $\pi_{ref} \leftarrow \pi_{init}$: a frozen reference policy used to prevent the policy from drifting too far from its initial behavior (KL regularization)
   • $\pi_{value} \leftarrow \pi_{init}$: a value network that learns to predict expected future reward from each state, trained alongside the policy
   • $\pi_0 \leftarrow \pi_{init}$: the active policy being optimized

2. For each training step $t = 1, \dots, T$:

   • Roll out: Generate trajectories from the current policy $\pi_{\theta_{t-1}}$ using temperature 1.0, collecting 32 rollouts to form a dataset $D_t = \{(s_1, a_1, r_1), \dots, (s_n, a_n, r_n)\}$ of state-action-reward tuples. Each trajectory receives a reward based on correctness and length.
   • Update policy: Fine-tune $\pi_\theta$ by maximizing the APA objective:

$\mathcal{L}_{APA}(\theta; D) = \frac{1}{|D|} \sum_{(s,a) \in D} \left(\log \pi_\theta(a|s) - \frac{\text{Adv}^{\pi_{\theta_{t-1}}}(s, a)}{\lambda} - \log \pi_{ref}(a|s)\right)^2$

where $D$ is the dataset of state-action pairs from rollouts, $\log \pi_\theta(a|s)$ is the log-probability the current policy assigns to action $a$ in state $s$, $\text{Adv}^{\pi_{\theta_{t-1}}}(s, a)$ is the advantage of action $a$ in state $s$ estimated using the value network and the rollout rewards, $\lambda = 2.0$ is the advantage coefficient (scaling the influence of the advantage term), and $\log \pi_{ref}(a|s)$ is the log-probability the reference policy assigns to the same action.

• Update value network simultaneously (the critic loss is omitted from the algorithm pseudocode for simplicity but is trained alongside the policy).
• Periodically reset reference: When validation reward converges, update $\pi_{ref} \leftarrow \pi_{\theta_t}$, replacing the old reference policy with the current improved policy. The paper resets the reference 3 times.

3. Stop when resetting the reference no longer yields validation improvement.

What the APA objective computes: it is a squared-error regression target for the policy's log-probabilities, not a standard likelihood objective. For each state-action pair, the target log-probability is:

$\text{target\_logprob} = \frac{\text{Adv}(s, a)}{\lambda} + \log \pi_{ref}(a|s)$

The policy is trained so that $\log \pi_\theta(a|s)$ matches this target. Actions with positive advantage get a target log-probability higher than the reference policy's log-probability — the model is encouraged to increase their probability. Actions with negative advantage get a target lower than the reference — the model is discouraged from taking them. The squared error penalizes deviations in both directions equally.

Why squared error rather than a likelihood ratio (as in PPO): Zhu et al. (2023) designed APA as a simpler, more stable alternative to PPO for language model fine-tuning. PPO uses a clipped importance sampling ratio $\pi_\theta(a|s) / \pi_{old}(a|s)$ multiplied by the advantage, which can be high-variance when the new and old policies diverge. APA replaces this with direct regression on log-probabilities, which is mathematically equivalent to fitting a Gaussian in log-probability space and avoids importance sampling entirely. The trade-off is that APA requires the reference policy for regularization (to keep the target log-probabilities in a reasonable range), while PPO's clipping provides regularization implicitly. The paper chose APA "due to its stability and robustness to changes in hyperparameters" (Section 6).

The reward function. The paper defines a "straightforward reward function that takes into account the correctness and length of the generated trajectory" (Section 6). The exact formula is not given in the main text, but the reward structure means:

• Correct trajectories receive positive reward, with longer trajectories receiving less reward than shorter ones (efficiency incentive).
• Incorrect trajectories receive zero or negative reward.

The advantage $\text{Adv}(s, a)$ is then computed as the difference between the observed return (sum of future rewards) from taking action $a$ in state $s$ and the value network's estimate of the expected return from state $s$. This is standard advantage estimation: positive advantage means the action led to better-than-expected outcomes; negative advantage means it led to worse-than-expected.

The role of the value network. The value network $\pi_{value}$ is trained to predict $V(s)$, the expected future reward from state $s$. It is initialized from the same SoS checkpoint and updated alongside the policy. Its loss (value loss coefficient: 10) is added to the APA objective. The value network enables credit assignment: even if a trajectory ultimately fails, individual actions early in the trajectory that moved toward the goal might have positive advantage if the value network recognizes them as progress. This is the key difference from STaR, which discards all actions from failed trajectories.

APA hyperparameters (Appendix C.5):

• Number of rollouts per step: 32
• Temperature for sampling: 1.0 (higher than STaR's 0.8 — APA can afford more exploration because even failed trajectories provide learning signal via the advantage)
• Batch size: 8, gradient accumulation: 1
• Online epochs: 2 (each batch of rollout data is trained on for 2 epochs before being discarded)
• $\gamma$ (discount factor): 1.0 (no discounting — all future rewards are equally valuable, appropriate because the goal is to solve the problem regardless of how many steps it takes)
• $\lambda$ (advantage coefficient): 2.0
• Learning rate: $1 \times 10^{-6}$ (an order of magnitude lower than STaR, typical for RL fine-tuning where stability is paramount)
• No learning rate scheduler
• Value loss coefficient: 10 (the value network's loss is weighted 10 times higher than the policy loss in the combined objective, ensuring the critic learns quickly)

Reference policy resets (Figure 4b). The paper observes that validation accuracy plateaus after about 4,000 training steps, at which point they reset $\pi_{ref} \leftarrow \pi_{\theta_t}$. This changes the regularizer: previously, the policy was penalized for deviating from the original SoS model; after reset, it is penalized for deviating from its current improved self. This allows progressive relaxation of the KL constraint — the policy is initially anchored close to SoS, then close to SoS+APA-iteration-1, then close to SoS+APA-iteration-2. The paper explicitly states this shifting reference can be interpreted as "a means to reduce the weight assigned to staying close to the reference distribution (the $\lambda$ parameter in the APA objective)" (Section 6). In practice, the authors found this strategy of resetting the reference to be "more stable for training when compared to designing a schedule for reducing $\lambda$ over training."

Why APA completes the picture. STaR is simple and effective but only learns from successes — it reinforces whatever the model already does that happens to produce correct answers. APA, through the value network, can learn from failures too: an action that moves closer to the goal but is followed by a mistake can still be recognized as valuable. The paper uses both methods not because one is strictly better but because they optimize different aspects of search: STaR for correctness, APA for correctness-plus-efficiency (via the length penalty in the reward). The empirical results (Figures 4, 5, 6) show that both methods improve over the base SoS model, with slightly different characteristics — APA appears to diverge more from the symbolic training strategies (Figure 5a) and suppresses arithmetic errors more effectively (Appendix Table 2).

Summary of Key Design Decisions and Their Justifications

• Explicit vs. implicit operations in the SoS language: explicit operations (current state, goal state, backtracking, goal check, exploration choice) become text that the model can inspect and optimize over; implicit operations (heuristic evaluation, pruning, state queue, state expansion) become internal representations the model can improve beyond the training heuristics. This split balances learnability with flexibility.
• Two heuristic functions (sum and multiply) rather than one: exposes the model to complementary search biases — additive reasoning (combine numbers to match the target's magnitude) and multiplicative reasoning (combine numbers to match the target's factors). The model can learn to flexibly deploy either or invent new heuristics.
• BFS with five breadth limits: varying the breadth from 1 (greedy beam search) to 5 (wider exploration) creates trajectories with different exploration-exploitation trade-offs. The model sees both narrow, efficient searches and broad, thorough searches.
• DFS with a single threshold: provides trajectories where the search goes deep before backtracking, complementing BFS's breadth-first approach. The threshold $h_{th} = T$ prunes states whose heuristic value exceeds the target — a simple way to prevent exploring states that are clearly numerically far off.
• Training on incorrect trajectories (~43% of SoS data): the central hypothesis — that seeing recovery from mistakes is essential for learning search. Empirically validated by the 25 percentage point improvement over OP training.
• GPT-Neo from scratch rather than fine-tuning a pretrained LM: isolates the effect of SoS training data from any search capabilities the model might have incidentally acquired during general pretraining.
• STaR with reset to $M_0$ each iteration: prevents distributional collapse where the model's training data becomes progressively narrower and self-reinforcing.
• APA with periodic reference resets: provides a practical stability mechanism that avoids the need to tune a complex $\lambda$ annealing schedule, while still allowing progressive relaxation of the KL constraint as the policy improves.
• Pairing APA (actor-critic with value function) with STaR (binary filtering): APA learns from both successes and failures; STaR only from successes. The combination (in the paper's analysis, not as a single algorithm) shows that both mechanisms yield improvements, suggesting the model benefits from both binary outcome feedback and continuous advantage feedback.

4. Key Insights and Innovations

Innovation 1: Search as a Trainable Sequence Modeling Problem, Not an External Scaffold

The dominant paradigm when this paper was written treated search as something that happens outside the language model—a symbolic algorithm (BFS, DFS, Monte Carlo Tree Search) calls the LM as a subroutine to propose and evaluate states, but the LM itself never learns to search. Tree of Thoughts (Yao et al., 2024), Graph of Thoughts (Besta et al., 2023), and related methods exemplify this "extrinsic" approach: the search strategy is hardcoded, the LM's parameters are frozen, and any capability improvements from search evaporate when the external algorithm is removed.

The SoS framework makes a fundamentally different move: search itself is recast as an autoregressive sequence generation task. Rather than treating exploration, backtracking, and heuristic evaluation as operations performed by an external controller, the paper serializes them into a linear token stream and trains the LM to generate that stream directly. This is not merely an engineering convenience—it is a conceptual reframing that shifts search from an inference-time scaffolding problem to a training-time representation learning problem.

The significance of this reframing extends beyond the immediate performance gains. If search is just sequence generation, then every capability that sequence models have demonstrated—transfer learning, in-context adaptation, self-improvement through fine-tuning—becomes available for search itself. The model can learn to adapt its search strategy to problem characteristics (something extrinsic methods cannot do without explicit meta-reasoning modules). It can improve its search through policy optimization (STaR and APA in Section 6). And critically, it can discover search strategies that were not present in its training data, because the model is learning a distribution over search trajectories, not executing a fixed algorithm. Figure 5c provides the empirical anchor for this claim: the fine-tuned SoS models solve problems that none of the 12 symbolic strategies could solve—direct evidence that the model has gone beyond mimicking its teachers.

The contrast with in-context search demonstrations (Gandhi et al., 2023; Sel et al., 2023) sharpens the distinction. Those methods also represent search in language, but the representation lives in the prompt, not the parameters. The model executes the demonstrated strategy but cannot improve it. SoS moves the search capability from the context window into the weights, making it parametric, improvable, and transferable across problems without consuming prompt length.

This is a fundamental shift rather than an incremental improvement. Prior work asked: "given a frozen LM, how do we search effectively?" SoS asks: "can we train an LM so that search is what it does?" The answer—a 25 percentage point accuracy improvement over optimal-path training (51.27% vs. 25.73%, Figure 3a), achieved despite the SoS model seeing fewer correct examples—suggests the question was worth asking.

Innovation 2: Mistakes and Recovery as Essential Training Data, Not Noise to Be Filtered

The standard practice in training LMs for reasoning is to curate datasets of correct solutions and filter out errors. The underlying assumption is straightforward: models learn what they see, so show them only correct reasoning and they will learn to produce correct reasoning. This paper demonstrates that this assumption is not just oversimplified—it is actively harmful for learning to search.

The optimal paths (OP) baseline sees 500,000 perfectly clean solution trajectories. Every token it observes is part of a correct solution. Yet this model achieves only 25.73% accuracy on held-out Countdown problems. The SoS model sees only 285,501 correct solutions (~57% of its training data), embedded within longer trajectories full of dead ends, incorrect arithmetic, backtracking, and outright failures. Yet it achieves 51.27%—nearly double the OP model's performance (Figure 3a).

This is a counterintuitive and conceptually significant finding: the presence of errors in training data, far from corrupting the model, is what enables it to reason robustly. The paper's explanation—that seeing recovery from mistakes teaches the model that errors are recoverable and backtracking is the appropriate response—has deep implications for how we think about training data curation for reasoning tasks. It suggests that the field's instinct to clean datasets by removing incorrect solutions may be systematically removing the very information that models need to learn robust reasoning strategies.

The significance is not merely empirical (though the 2× performance gap is striking). It is diagnostic: the paper identifies a specific mechanism—exposure bias in autoregressive generation—and shows that training on search trajectories directly addresses it. When the model generates at test time, it sees its own potentially erroneous outputs as context, which drift from the clean prefixes it saw during training. The OP model has never seen an error in context and therefore has no learned behavior for what to do when one appears. The SoS model has seen thousands of examples where an error in context was followed by backtracking and recovery—it has learned that errors are transitional states, not terminal ones. This also explains the snowballing problem: errors compound not because the model cannot reason, but because it has never been taught that reasoning includes error correction.

This finding connects to broader questions in machine learning about the value of negative examples and the role of process versus outcome supervision (cf. Lightman et al., 2023), but the mechanism is different: SoS does not require labeled correctness annotations for intermediate steps. The model learns from the structure of the trajectory itself—the sequence of explore, fail, backtrack, succeed—without needing explicit verifier scores.

Innovation 3: The Language for Search as a Unifying Abstraction for Strategy Diversity

Section 3 of the paper defines a vocabulary of nine primitive search operations (current state, goal state, state queue, state expansion, exploration choice, pruning, backtracking, goal check, heuristic) and specifies which are made explicit versus implicit in the serialized format. This is easy to read as an implementation detail—a domain-specific language for a specific game. But it represents a conceptual innovation with implications beyond Countdown: the idea that diverse search strategies (BFS, DFS, heuristic-guided variants) can be expressed in a common token-level representation, enabling a single sequence model to learn a distribution over strategies rather than any single one.

Prior work on learned search typically trains models to imitate a specific algorithm. Yang et al. (2022) train transformers to mimic MCTS with a fixed heuristic. Lehnert et al. (2024) train transformers to imitate A* with a fixed evaluation function. In both cases, the model learns one strategy. The SoS language, by contrast, is designed to be algorithm-agnostic—the same vocabulary can express BFS with breadth 5 and sum heuristic, DFS with multiply heuristic, or any combination thereof. The training data includes trajectories from 12 different strategy configurations, and the model is trained on all of them without labels indicating which strategy produced which trajectory.

The result, shown in Figure 3c, is that the trained SoS model does not strongly align with any single symbolic strategy in its state-visitation patterns. The highest correlation is 0.57 with DFS+sum, and the lowest is 0.27 with BFS-breadth-5+sum. The model is not copying; it is composing. This is a qualitatively different outcome from imitation learning—the model has learned something more abstract than any individual strategy, a kind of meta-search capability that flexibly utilizes different approaches.

The design choice of making heuristics implicit while making backtracking and exploration choice explicit is the key enabler here. If heuristics were explicit in the trajectory (e.g., numerical scores written out), the model would learn to compute those specific heuristics. By keeping them implicit, the model is forced to develop its own internal representations for state evaluation—representations that can improve through policy optimization (Figure 5a shows that after APA fine-tuning, the model's state-visitation patterns shift toward the multiply heuristic, suggesting it has learned to value that signal more). This split between explicit meta-cognitive operations (which provide the structure of search) and implicit evaluation (which provides flexibility) is a design pattern that could transfer to other domains.

This is an incremental-but-important advance in representation design for learned search. The specific vocabulary is tailored to Countdown, but the principle—express diverse strategies in a common action space with explicit structural moves and implicit evaluation—is general. It points toward a future where training data for reasoning includes not just solutions but structured traces of the reasoning process itself, serialized in a format that a transformer can model autoregressively.

Innovation 4: Self-Improvement That Transcends the Training Distribution

The self-improvement results in Section 6 are not just performance gains (though the ~6% improvement from APA and ~5% from STaR over the base SoS model, per Figure 4, are notable). The genuinely significant finding is qualitative: the fine-tuned models solve problems that none of the 12 symbolic strategies could solve (Figure 5c, ~4% of the "difficult" problem set) and solve 36% of problems that were unsolvable by the specific symbolic strategies that generated the training data (Figure 5b).

This is a stronger claim than "the model improved." It is evidence that the model has discovered search behavior that goes beyond interpolation between its training strategies. The symbolic strategies form a discrete set of fixed algorithms—BFS with this breadth, DFS with that heuristic. The model, trained on a mixture of their outputs, has learned something that is not simply a weighted combination of those algorithms. It has discovered a search policy that succeeds where all of its teachers fail.

This connects to a central question in machine learning: can models transcend their training distribution through self-play or iterative improvement? AlphaZero (Silver et al., 2018) demonstrated this for games with clean reward signals and perfect environment models. The SoS results suggest a path toward analogous capabilities for language models in reasoning domains, but with an important difference: the model must learn both the environment dynamics (arithmetic) and the search strategy simultaneously, without access to a ground-truth simulator. The fact that it can do so—that the same autoregressive generation produces both the arithmetic operations and the meta-cognitive decisions about what to explore next—is evidence for the viability of learned world models in reasoning.

The mechanism matters here. STaR improves the model by filtering for correctness: generate many trajectories, keep the ones that succeed, retrain on them. This is a form of outcome-based reinforcement that can amplify existing capabilities but cannot easily discover qualitatively new strategies—it reinforces whatever the model already does that works. APA, with its learned value function, can assign credit to actions that move toward the goal even in trajectories that ultimately fail. This enables the model to learn from near-misses and partial progress, which is more likely to surface novel solution patterns. The paper's finding that the APA model diverges more from the symbolic strategies than the STaR model (Figure 5a: APA shows larger shifts in state-visitation alignment) is consistent with this interpretation—APA is exploring more broadly in strategy space, while STaR is refining what already works.

The ~4% figure on the "difficult" problems may seem modest, but its significance is amplified by context: these are problems that 12 different heuristic-guided search algorithms, run to completion with access to a ground-truth environment model, could not solve. A 250M-parameter transformer, generating autoregressively with no external environment access, discovered solutions they missed. This is an existence proof that learned search can outperform designed search on some problems, which is the conceptual threshold the paper needed to cross.

Limitation note: The paper does not characterize what these newly solved problems look like or what novel strategies the model discovered for them. The alignment metrics (Figure 5a) show shifts in state-visitation patterns but do not reveal the actual search behavior at a mechanistic level. This leaves open the question of whether the model found genuinely novel heuristics or simply got lucky through stochastic exploration that the deterministic symbolic strategies did not attempt. Future work that extracts and analyzes the model's implicit heuristic would strengthen this contribution.

Innovation 5: Internal World Models as an Emergent Property of Search Training

The paper makes a claim in Section 7 that deserves recognition as a distinct conceptual contribution, even though it is more suggestive than fully demonstrated: that training on search trajectories causes the LM to develop an internal world model—a learned representation of the environment's transition dynamics that enables the model to simulate state transitions without access to a ground-truth simulator.

This is significant because the absence of world models is a standard criticism of autoregressive LMs (LeCun, 2023; Bachmann & Nagarajan, 2024). The argument goes: next-token prediction trains models to produce plausible continuations, not to understand the causal structure of the domain they're operating in. They cannot plan because they cannot simulate the consequences of their actions. SoS addresses this by forcing the model to simulate consequences as part of the generation: when the model writes "3 + 5 = 8" as part of a search trajectory, it must compute the result of that arithmetic operation correctly or the subsequent search will be incoherent. The model cannot rely on memorized patterns—the specific numbers vary across problems—so it must learn the operation itself.

The evidence for this claim comes from multiple angles. The low arithmetic error rate of the base SoS model (approximately 2 errors per trajectory on average, Appendix Table 2) shows that the model has learned to perform arithmetic reliably within its search trajectories. The fact that these error rates decrease with policy improvement (Figure 6, left) shows that optimizing for search correctness also improves environment simulation accuracy—the two capabilities are coupled. And the fact that the model generates valid trajectories with only 0.8% exploration errors (visiting states that don't follow from previous operations, Appendix Table 2) shows that it has learned the transition structure of the domain.

This is the paper's answer to the "snowballing errors" and "lookahead failure" criticisms. The model can backtrack because it has learned to recognize when its internal simulation has led to a dead end. It can look ahead because its autoregressive generation is the simulation of a forward trajectory—it writes out states, evaluates them against an implicit heuristic, and decides whether to continue deeper or backtrack. The search process is the mechanism by which the world model is exercised and refined.

This contribution is more suggestive than proven—the paper does not isolate the world model from the search policy or measure its accuracy independently. But the framing is conceptually valuable because it connects learned search in LMs to the broader literature on model-based reinforcement learning (Schrittwieser et al., 2020) and suggests that the path to better reasoning may not require explicit world model architectures but can emerge from the right training data structure. It also explains why SoS outperforms OP: the OP model never needs to simulate transitions (it only needs to reproduce the correct ones), so it never develops a robust world model. The SoS model must simulate both correct and incorrect transitions to generate coherent search trajectories, forcing it to learn the domain's causal structure.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper uses the Countdown game, a generalization of the 24 Game (Yang et al., 2022; Countdown, 2024). A problem consists of a set of input numbers and a target number (range: 10–100); the goal is to combine the inputs using the four basic arithmetic operations (+, -, ×, ÷) to reach the target. The paper constrains problems to 4 input numbers to keep search trajectories within the model's 4096-token context window — the authors note that 5-input problems can produce traces exceeding 60,000 tokens. The training dataset contains 500,000 problems generated synthetically, with 10% of target numbers held out for an out-of-distribution evaluation split. Two generalization test sets are constructed: (1) seen targets with unseen input number combinations, and (2) unseen targets with unseen input combinations.

• Base model(s). All experiments use a GPT-Neo model (Gao et al., 2020) with 250 million parameters, trained from scratch — not fine-tuned from a pretrained checkpoint. The architecture consists of 16 transformer layers with 16 attention heads, a hidden size of 1024, a context length of 4096 tokens, and Flash Attention 2 (Dao et al., 2022) for memory efficiency (Appendix C.1). The model is intentionally trained from scratch to isolate the effect of SoS training data from any search-relevant capabilities the model might have acquired during general-purpose pretraining.

• Metrics. The primary metric is accuracy: the percentage of problems for which the model generates a trajectory containing a correct solution path. Formally, correctness is a binary function $\mathbf{1}$ indicating whether the correct path $\mathcal{P}$ is present in the generated trajectory $\mathcal{T}$. The paper also reports two alignment metrics (defined in Section 4 of the paper, described in Section 3.4 of this analysis): (1) Alignment of Correctness, the Pearson correlation between two strategies' binary solve/fail vectors across problems, measuring whether strategies succeed and fail on the same problems; and (2) Alignment of States Visited, the normalized overlap in states visited by two search trajectories on the same problem, computed as $|\mathcal{T}_1 \cap \mathcal{T}_2| / \max(|\mathcal{T}_1|, |\mathcal{T}_2|)$ and averaged across problems. Additional diagnostic metrics include the average number of arithmetic errors per trajectory, the average number of states explored per correct trajectory (efficiency), and error rates by type (arithmetic, formatting, exploration, and other; Appendix Table 2).

• Baselines. The primary baseline is the Optimal Paths (OP) model: the same GPT-Neo architecture trained on 500,000 clean solution paths $\mathcal{P}$ — the correct sequence of state transitions from start to goal, stripped of all exploration, dead ends, and backtracking. The OP model sees 100% correct trajectories compared to the SoS model's ~57% correct trajectories (285,501 out of 500,000). Both models are trained for 50,000 gradient steps with identical hyperparameters (Appendix C.2–C.3). The paper also compares against the average accuracy of the symbolic search strategies used to construct the SoS dataset — 12 algorithm configurations combining DFS and BFS with two heuristic functions across five breadth limits (Section 4, Appendix B). These symbolic strategies have access to a ground-truth environment model (they can compute state transitions exactly) and serve as an upper-bound reference for what hand-designed search can achieve.

• Generation budget / compute accounting. For the base model comparison, both SoS and OP models are trained for the same number of gradient steps (50,000) with the same batch size (24) on the same hardware (4 × 80 GB A100 GPUs). The SoS trajectories are longer than OP trajectories (since they include dead ends and backtracking), so the SoS model processes more total tokens, but the paper controls for optimization budget rather than token count — the question is whether SoS yields better performance per unit of training compute. At inference, the model generates autoregressively until it produces a "Goal Reached" token or hits the context limit, with generation length varying per problem. For policy improvement (STaR and APA), the paper reports the number of training steps (20,000 per STaR iteration; ~4,000 per APA reference-reset cycle) and rollout counts (32 rollouts per APA step). No wall-clock latency comparisons are reported, which means the sequential nature of autoregressive search generation — where the model must generate one token at a time for the entire search trace — is not compared against the parallel sampling possible with extrinsic search methods.

• Cross-validation / statistical protocol. Model selection for all methods is based on accuracy on a validation set of 1,000 problems (Appendix C). For the base model comparison (Figure 3a), error bars represent 95% binomial confidence intervals computed over the 10,000 test problems for each evaluation set (held-out inputs and held-out targets). For policy improvement experiments, the paper reports convergence of validation accuracy during training (Figure 4b shows the validation reward curve for APA, with reference policy resets at convergence points; Figure 4c shows validation accuracy per STaR iteration). The alignment metrics (Figures 3b–c, 5a) are point estimates without reported confidence intervals. For the previously unsolved and difficult problem analyses (Figures 5b–c), error bars are 95% binomial confidence intervals computed over the 10,000 problems in each set. The paper does not report multi-seed training runs, so there is no variance estimate for the effect of random initialization or data ordering — all results are from single training runs.

Main Quantitative Results

Base Model Comparison: SoS vs. Optimal Paths (Figure 3a)

The headline result of Section 5 is that the SoS model substantially outperforms the OP model on held-out Countdown problems despite seeing fewer correct examples during training. On held-out inputs (unseen number combinations, seen targets): the SoS model achieves 51.27% accuracy, compared to 25.73% for the OP model — a gap of approximately 25.5 percentage points, or roughly a 2× improvement. On held-out targets (unseen targets, unseen number combinations): the pattern is "similar" according to the paper, though exact numbers for this split are not given in the main text — Figure 3a displays both bars, with the held-out targets showing a comparable gap.

The SoS model's 51.27% accuracy is slightly below the average accuracy of the symbolic search strategies used to construct the training dataset (roughly 57–60%, based on the fact that 57.1% of the generated trajectories were correct and the symbolic strategies had varying individual accuracies). However, the paper emphasizes that the SoS model is solving a harder problem: the symbolic strategies have access to a ground-truth environment model that computes state transitions exactly, while the SoS LM must simulate those transitions itself (performing arithmetic operations in language) and learn to search with an internal world model. Despite this disadvantage, the SoS model closes most of the gap while the OP model, trained on perfect solutions alone, achieves less than half the symbolic average.

Figure 3a error bars show 95% binomial confidence intervals. The gap between SoS and OP is well outside the confidence intervals for both bars, confirming statistical significance. The SoS bar's confidence interval appears to be roughly ±1 percentage point (given ~10,000 test problems), and the OP bar's interval is similar — the 25-point gap is therefore highly robust.

Alignment Analysis: What Strategy Does SoS Learn? (Figures 3b–c)

The paper measures whether the trained SoS model copies any single symbolic strategy or learns a composite approach. Alignment of Correctness (Figure 3b) shows that the SoS model's pattern of which problems it solves correlates most strongly with DFS using the sum heuristic — but the correlations are moderate, not near 1.0, and vary across strategies. The SoS model does not perfectly mirror any one training strategy in its success/failure pattern.

Alignment of States Visited (Figure 3c) provides more granular evidence. The highest state-visitation correlation is 0.57 with DFS using the sum heuristic. The lowest is 0.27 with BFS using breadth 5 and the sum heuristic. All correlations are well below 1.0, and there is meaningful variation across strategies. This is the key diagnostic: the model is not simply executing one memorized strategy. If it were, one bar in Figure 3c would be near 1.0 and the rest near 0. The observed pattern of moderate correlations across multiple strategies is consistent with the model having learned to flexibly deploy or combine elements from different strategies depending on problem characteristics.

The paper notes that the SoS model has "higher scores for alignment with strategies that use the sum heuristic" (Section 5), suggesting the model has learned to weight this heuristic more heavily — consistent with the sum heuristic being more generally useful than the multiply heuristic. However, the alignment scores are not high enough to claim the model is simply doing "sum-heuristic-guided search." The model has learned something that overlaps with but is not identical to any training strategy.

Policy Improvement: STaR Fine-Tuning (Figures 4a, 4c)

The SoS model is fine-tuned using Self-Taught Reasoner (STaR), an expert iteration method where the model generates trajectories on training problems, only correct trajectories are kept, and the model is fine-tuned on the filtered data, resetting to the original SoS checkpoint at each iteration. The paper reports convergence after 3 iterations of STaR fine-tuning (Figure 4c).

On held-out inputs: after 3 STaR iterations, the SoS+STaR model solves an additional ~5% of problems beyond the base SoS model (Figure 4a and Figure 4c). The paper states: "after three iterations, the finetuned SoS+STaR model solves an additional 5% of the held-out inputs test set beyond the base SoS model" — bringing accuracy from ~51% to ~56% (exact numbers: base SoS at 51.27%, SoS+STaR at approximately 56.27% based on the "additional 5%" figure, though the precise value is read from Figure 4a/c bar heights). On held-out targets: "a similar pattern is seen," indicating comparable improvement on the out-of-distribution target split.

Figure 4c shows the accuracy progression across the three STaR iterations, with error bars representing 95% binomial confidence intervals. The improvement is monotonic but diminishing — the largest gain comes in iteration 1, with smaller increments in iterations 2 and 3 before convergence.

Policy Improvement: APA Fine-Tuning (Figures 4a, 4b)

The SoS model is also fine-tuned using Advantage-Induced Policy Alignment (APA), an actor-critic method that trains a value network alongside the policy and uses advantage-weighted regression to update the policy, with periodic reference policy resets. On held-out inputs: APA yields an improvement of ~6% over the base SoS model (Figure 4a). This is comparable to or slightly better than STaR's ~5% improvement.

Figure 4b shows the training dynamics. Validation accuracy improves over approximately 4,000 training steps before plateauing, at which point the reference policy is reset (the vertical dashed lines in Figure 4b mark these resets). The paper performs 3 reference policy resets total, with each reset producing a jump in validation accuracy followed by further improvement and eventual plateau. The paper states: "We stop seeing improvements after 3 such resets." The three colored segments in Figure 4b correspond to training between reference resets (first segment: initial training from SoS; second segment: after first reset; third segment: after second reset; training stops after the third reset shows no further improvement).

The slightly larger improvement from APA (~6%) compared to STaR (~5%) is notable but within the error bars visible in Figure 4a — the paper does not claim a significant difference between the two methods. The more meaningful distinction is qualitative: as shown in Figure 5a (discussed below), APA diverges more from the symbolic training strategies than STaR does, suggesting it explores more broadly in strategy space.

Strategy Shift After Policy Improvement (Figure 5a)

To understand how the fine-tuned models differ from the base SoS model, the paper measures the difference in state-visitation alignment compared to the symbolic strategies. Figure 5a displays this as a bar chart: for each symbolic strategy (e.g., DFS+sum, BFS-breadth-3+multiply, etc.), the bars show the change in alignment between the fine-tuned model and the base SoS model.

Both STaR and APA models visit more states associated with the multiply heuristic after fine-tuning. The bars for multiply-heuristic strategies increase relative to the base model, indicating that the models have learned to weight multiplicative reasoning more heavily — consistent with the idea that the models are discovering that reaching factors of the target is an effective intermediate goal, a strategy that may have been underexploited by the base model.

The APA model diverges more from the symbolic strategies than the STaR model does. The APA bars in Figure 5a show larger shifts (both positive and negative) compared to the STaR bars. This is interpreted by the paper as evidence that APA is exploring more broadly in strategy space and discovering less symbolically-aligned approaches, while STaR is refining strategies already present in the base model. This aligns with the mechanistic difference between the two methods: STaR reinforces whatever the model already does that succeeds, while APA's value function can assign credit to novel action sequences that happen to work.

Solving Previously Unsolved and Difficult Problems (Figures 5b–c)

The most striking result of Section 6 is the models' ability to solve problems that were unsolvable by the training strategies. The paper selects two special evaluation sets from the training data:

Previously unsolved problems (Figure 5b): 10,000 problems from the SoS training set that were not solved by the symbolic strategies during dataset generation. These are problems where the search strategies explored the tree but failed to find a correct path within their search budget. The base SoS model, SoS+STaR, and SoS+APA are all evaluated on this set.

• The models solve approximately 36% of these previously unsolved problems. The bar chart in Figure 5b shows all three models (SoS, SoS+STaR, SoS+APA) at roughly the 36% level, with overlapping error bars. The paper states: "Remarkably, the models are able to solve approximately 36% of the previously unsolved problems." The error bars (95% binomial confidence intervals) are roughly ±1 percentage point given the 10,000-problem sample size.

Difficult problems (Figure 5c): 10,000 problems from the SoS training set that none of the 12 symbolic search strategies could solve. This is a harder subset — these problems defeated all the heuristic-guided search algorithms, not just a subset of them.

• The models solve approximately 4% of these difficult problems. The bar chart in Figure 5c shows the base SoS model, SoS+STaR, and SoS+APA all achieving roughly 4% accuracy on this set. The error bars again indicate statistical reliability at this sample size.

The paper frames the 4% figure as particularly significant: these are problems that no combination of BFS/DFS with sum/multiply heuristics, run to completion with access to a ground-truth environment model, could solve. A 250M-parameter transformer, generating autoregressively with no external environment access, found solutions to ~400 of these 10,000 problems. This is the paper's strongest evidence for strategy discovery beyond the training distribution — the model has learned something that is not simply an interpolation of its teachers' strategies.

Caveat on these results: The paper does not analyze how the models solved these problems — what novel search patterns, heuristics, or strategies they employed. The state-visitation alignment shifts in Figure 5a suggest movement toward multiply-based heuristics, but these are aggregate statistics; they do not reveal the mechanism for solving the previously-unsolvable problems. It remains possible that the models succeeded through stochastic exploration (temperature sampling at test time) that the deterministic symbolic strategies never attempted, rather than through genuinely novel heuristic reasoning. The paper does not rule this out.

Efficiency and Error Analysis (Figure 6, Appendix Tables 1–2)

Beyond accuracy, the paper reports two efficiency metrics that improve with policy optimization:

Average number of states explored per correct trajectory (Figure 6, right). Policy improvement leads to more efficient search — the models find the solution after exploring fewer states on average. The bar chart in Figure 6 (right panel) shows that the base SoS model explores more states per correct solution than either SoS+STaR or SoS+APA. Exact numbers are given in Appendix Table 1: the base SoS model explores approximately 51.2 states per correct trajectory (with a 95% confidence interval of roughly ±1.5), while SoS+STaR and SoS+APA explore fewer (exact values not quoted in the main text but visible in the table). This efficiency gain is a side effect of optimizing for correctness — the models learn that shorter, more direct searches are more likely to succeed, or that efficient search patterns are reinforced during policy improvement.

Average number of arithmetic errors per trajectory (Figure 6, left). Policy improvement reduces the error rate. The bar chart in Figure 6 (left panel) and Appendix Table 2 show that the base SoS model averages approximately 2 arithmetic errors per trajectory. After STaR fine-tuning, this drops slightly. After APA fine-tuning, the drop is larger — the paper notes that "APA training method suppressed arithmetic errors significantly" (Appendix D). This is consistent with APA's advantage-weighted learning: actions that produce arithmetic errors likely lead to trajectory failure or negative advantage, so the model learns to avoid them. The exact numbers from Appendix Table 2 are: base SoS averages ~2.02 arithmetic errors per trajectory; SoS+STaR averages slightly lower; SoS+APA averages substantially lower (the paper does not quote exact values in the main text).

Exploration error rate. The base SoS model generates valid trajectories with a low rate of exploration errors — approximately 0.8% of visited states are exploration anomalies (jumping to a node with no parent, etc.; Appendix Table 2). This shows the model has learned the transition structure of Countdown. Notably, the STaR fine-tuned model shows an increased rate of exploration errors compared to the base model (exact numbers in Appendix Table 2), despite its higher accuracy — suggesting that STaR's filtering-based optimization may encourage slightly riskier exploration that sometimes produces invalid states. APA, by contrast, maintains a low exploration error rate similar to the base model.

Formatting and other errors. Appendix Table 2 also reports formatting errors (malformed output that doesn't parse) and "other" errors (child nodes with incorrect number sets). Both are low across all models (well under 1 per trajectory on average) and show different patterns across training methods, confirming that the models learn the syntactic format of search trajectories reliably.

Ablation Studies and Robustness Checks

The paper does not contain a dedicated ablation section with systematic removal of components. However, several comparisons in the main experiments serve as implicit ablations, and Appendix D provides diagnostic analyses:

• Optimal paths (OP) vs. Stream of Search (SoS) training data (Figure 3a). This is the key ablation: training on search trajectories versus training on clean solution paths. The 25-percentage-point gap (51.27% vs. 25.73%) is the paper's central demonstration that the search process itself — including mistakes and backtracking — is the active ingredient in the training data. This implicitly ablates the presence of dead ends, backtracking operations, and incorrect exploration from the training distribution, showing they are necessary for the model to learn search.

• Base SoS vs. SoS+STaR vs. SoS+APA (Figures 4, 5, 6). This comparison ablates the choice of policy improvement method. Both methods improve over the base model, with APA showing slightly larger gains (~6% vs. ~5%) and greater divergence from symbolic strategies (Figure 5a). The fact that both methods work — one using binary filtering (STaR), one using continuous advantage weighting (APA) — suggests the improvement is robust to the specific RL algorithm and driven by the self-improvement loop itself.

• State-visitation alignment analysis (Figures 3c, 5a). While not a causal ablation, this analysis serves as a robustness check on the claim that the model is not simply copying a single training strategy. The moderate correlations (0.27–0.57 in Figure 3c) and the shifts after fine-tuning (Figure 5a) provide evidence that the model learns a flexible, composite search capability rather than overfitting to one heuristic.

• Error-type breakdown (Appendix Table 2). This analysis checks whether improved accuracy comes at the cost of increased errors in other categories (e.g., more exploration anomalies because the model is exploring more aggressively). The results are mixed: APA suppresses arithmetic errors substantially without increasing exploration errors, while STaR increases exploration errors slightly — a trade-off that the paper notes but does not explore in depth.

• Reference policy resets in APA (Figure 4b). The three resets of the reference policy during APA training provide an implicit ablation of the KL regularization schedule. The paper states that resetting the reference is "more stable for training when compared to designing a schedule for reducing $\lambda$ over training" (Section 6). Figure 4b shows that each reset enables further improvement after a plateau, validating the strategy — without resets, the model would likely converge to a lower accuracy.

What is missing in terms of ablations:

• No ablation on the number of training strategies. The SoS model is trained on 12 strategy configurations. Would training on only 1 strategy (e.g., only DFS+sum) produce comparable results? If so, the strategy diversity might be less important than the presence of search traces per se. The paper does not isolate this variable.
• No ablation on the explicit vs. implicit split in SoS language. The paper makes deliberate choices about which search operations are verbalized and which are kept implicit (Section 3), but no experiment varies this split. Would making heuristics explicit (writing out numerical scores) improve or degrade the model's ability to discover new strategies? The paper claims that implicit heuristics enable flexibility but provides no empirical test.
• No ablation on the proportion of incorrect trajectories. The SoS dataset is 57% correct. What if the proportion were 80%? 30%? The paper does not vary this ratio, leaving open the question of how much "mess" is optimal.
• No ablation on model scale. All experiments use a 250M-parameter GPT-Neo model. How does the SoS vs. OP gap scale with model size? Does a larger model benefit more or less from search trajectory training?
• No temperature ablation for STaR. The paper uses temperature 0.8 for STaR trajectory generation and 1.0 for APA. How sensitive are the results to this choice? Would deterministic generation (temperature 0) for STaR be more effective?
• No single-seed replication. All results are from single training runs. The paper does not report variance across random seeds, so the reliability of the exact numerical improvements (~5% for STaR, ~6% for APA) cannot be assessed.

Critical Assessment

Claim 1: "SoS pretraining increases search accuracy by 25% over models trained to predict only the optimal search trajectory."

What the experiments demonstrate: The SoS model achieves 51.27% on held-out inputs compared to 25.73% for the OP model (Figure 3a) — a gap of approximately 25.5 percentage points. The paper reports a "similar pattern" for held-out targets. These numbers come from a single training run of each model on 500,000 Countdown problems, evaluated on 10,000 test problems per split.

Does this support the claim? Yes, directly and strongly, for the specific setting tested. The gap is large, well outside the reported confidence intervals, and exists for both test splits. The claim as stated is factual.

What limits the strength of this support: The comparison controls for training steps (50,000 each) but not for total tokens processed. The SoS trajectories are longer (they include dead ends) so the SoS model likely sees more total tokens — the paper does not report token counts. If the OP model were trained for more steps to match total tokens, would it catch up? The paper's framing is that what matters is the information content of the data, not the raw token count, but the experiment does not fully rule out that token volume contributes to the gap. A fairer comparison would match total tokens and report whether SoS still outperforms OP — or match total correct examples (both models trained on 285,501 correct trajectories, with SoS having additional incorrect ones) to isolate the value of seeing errors. Neither comparison is reported.

Additionally, Countdown is a specific arithmetic search domain. The claim that "search accuracy improves" is in-domain — the paper does not test whether SoS training transfers to other search problems or yields generalizable search capabilities. The model is trained and evaluated on Countdown only, so the learned search ability is domain-specific by construction.

Claim 2: "The finetuned SoS models solve 36% of previously unsolved problems, including problems that cannot be solved by any of the heuristic solvers."

What the experiments demonstrate: On a set of 10,000 previously unsolved problems (from the training distribution, not held-out), the models solve ~36% (Figure 5b). On a set of 10,000 difficult problems that defeated all 12 symbolic strategies, the models solve ~4% (Figure 5c).

Does this support the claim? The numbers support the claim as stated, but several qualifications are necessary:

• These are in-distribution problems. Both sets are drawn from the training data, not from a held-out test set. The models were trained on these exact problems during the STaR/APA fine-tuning (which uses training-set problems for trajectory generation). Solving them after fine-tuning is a measure of self-improvement on the training distribution, not generalization to novel problems. The paper does not report whether the model also solves held-out problems that the symbolic strategies cannot — that would be a stronger test of strategy discovery.

• The 36% figure is for "previously unsolved" — but by which solvers? The paper states these are "problems from the SoS training set that were unsolved by symbolic strategies when the dataset was generated" (Section 6). This means the specific heuristic-guided BFS/DFS runs that generated the training data failed on these problems. It does not necessarily mean no symbolic strategy could solve them — a different heuristic, a larger search budget, or a different algorithm might succeed. The "difficult" set (Figure 5c), which includes only problems that all 12 strategies failed on, is a cleaner test. The 4% figure there is more meaningful but also more modest.

• The mechanism is not explained. The paper speculates about strategy discovery but does not analyze how the model solved these problems. The alignment shift toward multiply heuristics (Figure 5a) is suggestive but does not reveal what novel search patterns emerged. Without trajectory-level analysis of the solved difficult problems, the "discovery of new strategies" claim remains an interpretation rather than a demonstrated fact. The model could be succeeding through variants of the training strategies with better stochastic exploration (temperature sampling at test time vs. deterministic symbolic search), which would be an improvement in implementation rather than a discovery of new algorithmic principles.

• The 4% figure is small in absolute terms. Solving 4% of the hardest problems (400 out of 10,000) is an existence proof that learned search can exceed designed search on some instances, but it does not demonstrate that the model has broadly transcended its training strategies. Most difficult problems remain unsolved. The claim that the model "discover[s] new search strategies" should be tempered by the observation that this discovery, if it occurred, applies to a small fraction of the hardest cases.

Claim 3: "Language models can learn to solve problems via search, self-improve to flexibly use different search strategies, and potentially discover new ones."

What the experiments demonstrate: The base SoS model solves 51.27% of problems vs. 25.73% for OP (Figure 3a), showing that training on search traces enables search-based problem solving. The alignment analyses (Figures 3b–c) show the model does not copy any single strategy. The policy improvement methods increase accuracy by 5–6% (Figure 4a) and shift strategy alignment (Figure 5a). The previously unsolved and difficult problem results (Figures 5b–c) suggest capabilities beyond the training strategies.

Does this support the claim? The claim has three parts, each with different levels of support:

• "Learn to solve problems via search": Strongly supported by Figure 3a. The SoS model generates valid search trajectories (low error rates per Appendix Table 2) that contain correct solution paths for over half of the test problems. The search is intrinsic — the model autonomously explores, backtracks, and goal-checks within a single autoregressive generation.

• "Self-improve to flexibly use different search strategies": Partially supported. The alignment analyses (Figures 3b–c, 5a) show the model uses a mix of strategies and that this mix shifts with policy improvement. However, "flexibly use" implies the model adapts its strategy to problem characteristics — for example, using DFS-like search on problems with deep narrow solutions and BFS-like search on problems with broad shallow solutions. The paper provides no evidence for such per-problem adaptation. The alignment metrics are aggregate correlations; they do not show that individual trajectories vary systematically with problem features. The model could be using a single hybrid strategy (a blend of DFS and BFS) uniformly across all problems, which would not demonstrate "flexible use."

• "Potentially discover new ones": Weakly supported, with qualifications. The 4% difficult-problem solve rate (Figure 5c) is consistent with strategy discovery but does not prove it. Alternative explanations: (a) stochastic sampling succeeds where deterministic symbolic search failed — the symbolic strategies may have been run deterministically or with limited exploration, while the LM samples with temperature; (b) the model benefits from the implicit heuristic it learned, which could be more effective than the explicit sum/multiply heuristics even though it was derived from training on trajectories that used those heuristics — this is learning a better heuristic, not necessarily a new strategy; (c) the model benefits from simply having trained on more diverse trajectories than any single symbolic strategy saw during its single run. To demonstrate discovery of new strategies, the paper would need to show that the model's trajectories on the solved difficult problems exhibit qualitatively different exploration patterns (e.g., different branching behavior, different backtracking criteria) than any of the 12 training strategies — and it does not provide this analysis.

Additional Strengths

• The SoS vs. OP comparison is clean and well-controlled. Both models have identical architecture, training steps, and optimizer settings. The only difference is the training data format. This is a strong experimental design that isolates the variable of interest.
• The error analysis (Appendix Tables 1–2) adds credibility. The paper does not just report accuracy — it shows that the model's generated trajectories are structurally valid (low exploration error rate) and that the model can perform arithmetic (moderate error rate that improves with training). This rules out the possibility that the SoS model is succeeding through some degenerate strategy like generating random operations until the target accidentally matches.
• The use of two policy improvement methods (STaR and APA) with similar results strengthens the self-improvement finding. If only one method worked, the result could be an artifact of that specific algorithm. Both methods improving over the base model, with slightly different characteristics, suggests the self-improvement loop itself is the active ingredient.

Genuine Weaknesses and Missing Experiments

• Single benchmark, single model scale, single training run. All results are on Countdown with a 250M-parameter GPT-Neo model, trained once per condition. There is no evidence that SoS generalizes to other search problems, other model architectures, or other model scales. The paper acknowledges this limitation in Section 7 ("Our empirical results were restricted to the game of Countdown") but the claims in the abstract and introduction are stated without this qualification.

• No comparison against extrinsic search methods (Tree of Thoughts, etc.). The paper positions SoS against extrinsic search in Section 2 but never runs a head-to-head comparison. Would Tree of Thoughts with the same base LM and comparable inference budget outperform SoS? If so, the advantage of intrinsic search over extrinsic scaffolding remains theoretical. The paper speculates about inference efficiency gains but provides no empirical cost comparison.

• No comparison against process supervision. The paper discusses Lightman et al. (2023) as an alternative approach but does not compare against it. An experiment training a verifier model on the Countdown domain and using it for guided search, versus the SoS approach, would help locate SoS in the solution space.

• The difficulty estimation problem is not addressed. Unlike the test-time compute scaling paper (Snell et al., 2024), this paper does not attempt to adaptively allocate compute based on problem difficulty. All problems get the same autoregressive generation budget. This means the model wastes compute on easy problems (generating long search traces when the answer is obvious) and may run out of context length on hard ones.

• Context length as a hard ceiling. The 4096-token context window limits the model's search depth. For 5-input Countdown problems, the paper notes that trajectories can reach 60,000 tokens — well beyond the model's capacity. This means SoS, as implemented, does not scale to harder instances of the same problem. An extrinsic search system with a symbolic state tracker could search arbitrarily deep. The paper acknowledges this limitation but does not propose a solution (like dynamic context management or hierarchical search).

• No analysis of what the model learned about heuristics. The paper keeps heuristics implicit to allow the model to develop its own. But it never probes what implicit heuristic the model actually learned. An experiment where the model's state evaluations are extracted (e.g., by examining attention patterns or probing hidden states) could reveal whether the model internalized the sum heuristic, the multiply heuristic, a combination, or something novel. Without this analysis, the claim of "strategy discovery" remains speculative.

• The STaR and APA improvements are small in absolute terms (5–6%). While statistically significant, these gains are modest relative to the 25-point gap between OP and SoS. The paper's strongest signal comes from the SoS representation itself, not from the policy improvement methods. The self-improvement loop is demonstrated but not transformative in this setting — it adds incrementally to what SoS pretraining already achieves.

• The "difficult" problem results conflate two interpretations. The fact that the model solves 4% of problems that no symbolic strategy solved could mean the model discovered novel strategies, or it could mean the symbolic strategies were suboptimally configured (e.g., insufficient search budget, poorly chosen heuristic thresholds) and the LM's learned heuristic is simply better tuned to the problem distribution. The paper does not rule out the possibility that a symbolic strategy with optimized hyperparameters could match or exceed the LM's performance on these problems.

6. Limitations and Trade-offs

Limitation 1: Single Benchmark, Single Model Family, Single Training Run

The assumption or constraint. The paper explicitly states in Section 7: "Our empirical results were restricted to the game of Countdown." All experiments use a single model architecture (250M-parameter GPT-Neo trained from scratch), a single task domain (arithmetic search), a single training dataset size (500,000 problems), and — critically — a single training run per experimental condition. There is no multi-seed replication, no evaluation on other search domains (code generation, mathematical proof, planning), and no test of whether the SoS representation transfers to other model families or scales.

The consequence. Three distinct uncertainties arise. First, domain specificity: Countdown is a clean, formally specifiable search problem with a well-defined state space, deterministic transitions, and objectively verifiable solutions. Whether SoS training benefits transfer to "messier" reasoning domains — where the search space is not enumerable, intermediate states are ambiguous, or correctness signals are noisy — is entirely unknown. The paper's findings could be specific to domains that look like the training data: arithmetic operations on small sets of numbers with a clear goal state.

Second, model-specificity: The 250M GPT-Neo model is relatively small by contemporary standards. The paper presents no evidence on whether the SoS vs. OP gap (25 percentage points, Figure 3a) widens, narrows, or vanishes at different model scales. A larger model might benefit more from SoS training (because it has greater capacity to internalize an implicit heuristic and world model), or it might benefit less (because larger models may already develop rudimentary search-like behaviors from diverse pretraining, reducing the marginal value of explicit search traces). The single-model design makes the scaling behavior completely uncharacterized.

Third, statistical reliability of the exact numbers: Without multi-seed training, the reported improvements — 51.27% vs. 25.73% for SoS vs. OP, ~5% gain from STaR, ~6% from APA (Figure 4a) — are point estimates from single runs. The 95% binomial confidence intervals on test accuracy reflect sampling variance over the 10,000 test problems but do NOT capture variance from random initialization, data ordering, or other training-run-level stochastic factors. The stability of the SoS training procedure and the reproducibility of the exact performance gap are unmeasured.

What evidence exists in the paper. The paper acknowledges the domain limitation transparently: "We are optimistic that SoS extends to more challenging, real-world tasks" but provides no empirical support for this optimism (Section 7). The model scale limitation is not discussed. The single-run limitation is not acknowledged — Appendix C reports training details for single runs per condition, with no mention of variance across seeds.

Mitigation status. Not mitigated. The paper suggests future work on extending SoS to other domains (Section 7: "external-structured search methods... are likely to be more efficient for these tasks; in the longer run the increased flexibility and learnability of internally-structured search... may prevail") but conducts no such experiments. The single-domain, single-model, single-run design is a deliberate scoping choice for a paper introducing a new framework, but it means the generality of the findings remains a hypothesis rather than an established fact.

---

Limitation 2: The Symbolic Strategy Baseline Encoding Is a Weak Comparison, Not a Strong Upper Bound

The assumption or constraint. The paper's central narrative — that SoS models "discover new search strategies" and transcend their training distribution — rests on a specific comparison: the SoS models solve ~4% of problems that none of the 12 symbolic search strategies could solve (Figure 5c). The symbolic strategies are BFS and DFS guided by two simple heuristics (sum distance and factor distance) with fixed hyperparameters (DFS threshold = target value; BFS breadth limits = 1–5). The paper treats this suite as the envelope of what heuristic-guided symbolic search can achieve in Countdown.

However, the symbolic strategies have known limitations that make them a weak upper bound. They use fixed heuristic thresholds and breadth limits that were not optimized for maximal solve rate — they were chosen to be "simple and interpretable" and to "provide a simple and intuitive way to design symbolic algorithms to generate the SoS dataset" (Appendix B). The symbolic strategies are run once per problem with deterministic expansion — if a strategy fails on a problem, it is never retried with different hyperparameters or a different random seed. The strategies never combine heuristics (e.g., using sum heuristic for breadth and multiply heuristic for depth pruning within the same search). Most importantly, the symbolic strategies never get additional compute budget — no best-of-N, no beam search over heuristic values, no iterative deepening.

The consequence. The "solved problems that no symbolic strategy could solve" result (Figure 5c, ~4%) is not evidence that learned search outperforms designed search in general. It is evidence that a learned policy — which combines elements of multiple heuristics, benefits from stochastic sampling (temperature > 0 at inference), and has been optimized through multiple rounds of self-improvement — can solve some problems that a specific, fixed set of non-optimized symbolic strategies could not solve on a single attempt. This is a much weaker claim.

A more appropriate baseline would be: symbolic search strategies with optimized hyperparameters (tuned breadth limits, tuned heuristic thresholds), run with stochastic tie-breaking or multiple random restarts, potentially combining heuristics in an ensemble, given a compute budget comparable to the SoS model's generation cost. Such a baseline might match or exceed the SoS model's performance on the difficult problem set — the paper cannot rule this out because the experiment was not conducted. The paper's conclusion that SoS enables "discovery of new strategies" would be on firmer ground if the symbolic baseline were properly tuned to represent the best achievable performance from designed search with comparable computational resources.

What evidence exists in the paper. The symbolic strategies are described in Appendix B and Algorithms 3–4. Their hyperparameters are listed: DFS uses a single threshold h_th = T; BFS uses breadth limits b ∈ [1, 2, 3, 4, 5]. There is no hyperparameter optimization reported for these strategies — no sweep over different thresholds, no combination of heuristics, no ensemble. The paper reports the average accuracy of the symbolic strategies (roughly 57–60%, inferred from the 57.1% correct trajectory rate in the SoS dataset, Section 5) but does not report the maximum accuracy achievable by the best single symbolic strategy or by an optimized ensemble. The alignment analyses (Figures 3b–c) show that different symbolic strategies have different strengths — their solve/fail patterns are imperfectly correlated — but the paper never reports what accuracy would be achieved by running all 12 strategies and taking a majority vote or oracle selection, which would be the natural ensemble baseline.

Mitigation status. Not addressed. The paper treats the 12 symbolic strategies as representative of the space of possible search algorithms without justifying why a fixed set of non-optimized strategies constitutes a strong baseline. The "strategy discovery" interpretation of Figure 5c relies on this baseline being strong, but the paper provides no evidence or argument that it is.

---

Limitation 3: Context Length as a Hard Ceiling on Search Depth

The assumption or constraint. The SoS model generates the entire search trajectory — exploration, dead ends, backtracking, and all — within a single autoregressive generation limited by the model's context window of 4096 tokens (Appendix C.1). The training data is constructed to fit within this window: "We consider problems with 4 input numbers since these problems are challenging enough to have long search traces without the search traces exceeding a standard LM context window" (Section 4). The paper explicitly acknowledges that 5-input problems can produce trajectories up to 60,000 tokens — roughly 15× the available context.

The consequence. There is a fundamental tension between problem difficulty and representation capacity. Harder problems require more search — more states explored, more backtracking, deeper trees — but also require longer trajectories that consume more context. The 4096-token window imposes a ceiling on how much search the model can perform. If a problem requires exploring more states than can be serialized in 4096 tokens, the model will either (a) truncate its search and fail to find a solution despite it being discoverable with more exploration, or (b) generate past the context limit and produce incoherent output. The paper does not report what fraction of failures on the test set are due to the model hitting the context limit before finding a solution versus due to genuinely incorrect search.

This limitation is structural, not incidental. An extrinsic search system (Tree of Thoughts or similar) can search arbitrarily deep because the state is tracked symbolically and the LM is called as a stateless subroutine at each step — context length is only consumed by the current reasoning step, not by the entire search history. SoS trades this depth scalability for the efficiency and learnability of intrinsic search, but the trade-off means SoS cannot handle problems whose search depth exceeds what fits in the model's context window. As problem complexity increases, this limitation becomes binding. The paper's own example — 60,000-token trajectories for 5-input problems — demonstrates that even a modest increase in problem size pushes well beyond the feasible range.

Furthermore, this limit interacts with the policy improvement methods. STaR and APA optimize for correctness and efficiency (Figures 4–6), and the paper reports that these methods reduce the average number of states explored per correct trajectory (Figure 6, right). This could partially mitigate the context-length constraint — more efficient search fits more exploration within the same window — but the paper does not quantify how much headroom this efficiency gain provides or whether it would be sufficient for 5-input problems.

What evidence exists in the paper. The paper explicitly states the 4-input constraint in Section 4 and acknowledges the 60,000-token trajectory length for 5-input problems. Appendix Table 1 reports the average number of states explored per trajectory for the base SoS model (~51.2 states per correct trajectory) and the improved models (fewer). The paper does NOT report: what fraction of test trajectories hit the 4096-token limit, what fraction of failures are attributable to context truncation, or how accuracy would change if the context window were doubled or quadrupled. There is no experiment varying context length to measure the search-depth scalability curve.

Mitigation status. Partially acknowledged but not addressed. The paper notes the limitation in discussing domain choice (Section 4) and briefly mentions "limits, summarization, cycle checks, and subgoal setting" as future operations that "could enhance the SoS framework" (Section 7) — these could in principle compress search trajectories to fit more exploration in fewer tokens, but no such mechanisms are implemented or tested. The limitation is intrinsic to the autoregressive-generation design of SoS and would require architectural changes (e.g., hierarchical search representations, dynamic context management, or retrieval-augmented state tracking) to overcome.

---

Limitation 4: Policy Improvement Gains Are Small Relative to the SoS Representation Gain, and the Mechanism Is Unexplained

The assumption or constraint. The paper presents STaR and APA as methods that enable the SoS model to "self-improve" and "discover new search strategies" (Section 6, Abstract). The empirical improvements are approximately 5% (STaR) and 6% (APA) over the base SoS model on the held-out test set (Figure 4a). These are statistically significant but modest in absolute terms — especially compared to the ~25 percentage point gain from moving from OP to SoS training (Figure 3a). The paper's strongest result is the representation (SoS training data), not the optimization (policy improvement methods applied on top of it).

The consequence. There are two distinct concerns. First, the small absolute gain raises a practical question: do the policy improvement methods justify their computational cost? STaR requires three iterations of generating trajectories on training problems, filtering for correctness, and fine-tuning the model (Algorithm 1). APA requires training a separate value network, running 32 rollouts per training step, and managing reference policy resets (Algorithm 2). The paper does not report the wall-clock time or FLOP cost of these procedures relative to base SoS training. If STaR and APA each multiply the total training cost by a factor of 2–3× but yield only 5–6% accuracy gains, a practitioner might reasonably choose to skip policy improvement and accept the base SoS model's performance — or invest the same compute budget into training a larger base SoS model, which the paper does not explore as an alternative.

Second, and more fundamentally, the paper does not explain the mechanism by which policy improvement produces gains. The alignment analysis (Figure 5a) shows that STaR and APA shift state-visitation patterns (more multiply-heuristic states visited after fine-tuning), but this is a description of what changed, not a causal explanation of why the change improved performance. Did the models learn that the multiply heuristic is more reliable than the sum heuristic on certain problem types? Did they learn to dynamically switch between heuristics based on the problem? Did they simply become more efficient at executing the same strategies (shorter trajectories = fewer opportunities for error)? The paper provides diagnostic metrics (error rates, trajectory lengths in Figure 6 and Appendix Tables 1–2) but does not synthesize them into a coherent account of the improvement mechanism.

The "36% of previously unsolved problems" result (Figure 5b) is particularly opaque. The base SoS model already achieves ~36% on this set — the policy improvement methods do not increase this number meaningfully (all three bars in Figure 5b overlap within error bars). This means the ability to solve previously unsolved problems comes from SoS pretraining alone, not from STaR or APA. The policy improvement methods improve aggregate test accuracy (Figure 4a) presumably by getting better at problems the base SoS model could already sometimes solve, rather than by solving qualitatively new categories of problems. The paper does not make this distinction, leaving the impression that STaR and APA contribute to the "discovery" narrative when the evidence suggests the base SoS representation is the primary driver.

What evidence exists in the paper. The accuracy improvements are in Figure 4a and 4c. The state-visitation shifts are in Figure 5a. The previously unsolved and difficult problem results are in Figures 5b–c, where the base SoS model matches or nearly matches the fine-tuned models. The error analysis is in Appendix Table 2 and Figure 6. The paper does not report computational cost for policy improvement, does not ablate whether the same total compute budget invested in a larger base SoS model would outperform SoS+APA, and does not provide trajectory-level analysis of how the fine-tuned models' search behavior differs from the base model on problems where both succeed or where only the fine-tuned model succeeds.

Mitigation status. Not addressed. The paper presents policy improvement as a secondary contribution that demonstrates self-improvement is possible, but does not characterize its efficiency, its mechanism, or its necessity. The claim that policy improvement enables "discovery of new strategies" (Abstract, Section 6) is not supported by the experimental evidence at the level of granularity provided — the improvement appears to be quantitative refinement rather than qualitative discovery, and the paper does not disentangle these possibilities.

---

Limitation 5: Difficulty Estimation Overhead Is Absent — All Problems Receive the Same Generation Budget

The assumption or constraint. The SoS model generates complete search trajectories for every problem with no adaptive allocation of compute based on problem difficulty. Easy problems (where the target is trivially reachable, e.g., all input numbers sum exactly to the target) receive the same autoregressive generation budget as hard problems (where extensive backtracking and exploration are needed). The model generates until it produces a "Goal Reached" token or hits the context limit, with no mechanism to stop early on solved problems or allocate extra context to promising but incomplete explorations.

This contrasts with work on test-time compute scaling (e.g., Snell et al., 2024) that explicitly conditions inference strategy on estimated problem difficulty. The SoS framework, as implemented, provides no difficulty signal to the model before or during generation — the model must discover through its own search whether a problem is easy or hard, consuming tokens either way.

The consequence. There is a computational inefficiency in the inference budget. On easy problems, the model likely generates unnecessarily long trajectories — the paper does not report the average trajectory length for correct solutions stratified by problem difficulty, but the efficiency numbers in Appendix Table 1 (~51 states per correct trajectory for base SoS, fewer for fine-tuned models) are averages that likely include many easy problems where a few states would suffice. The model wastes compute exploring alternatives when the first path is already correct, and the paper's generation procedure (generate until "Goal Reached") provides no early-stopping mechanism based on confidence.

On hard problems, the fixed context window may be insufficient — the model may need more exploration than fits in 4096 tokens but has no mechanism to request additional budget. The paper does not report what fraction of failures on hard problems are due to the search being incomplete (the model was still exploring promising branches when it hit the context limit) versus the search being exhausted (the model explored all plausible branches and concluded none work). Without this distinction, it is unclear whether the accuracy ceiling is due to the model's search capability or its context-length budget.

Additionally, the uniform budget means the model's generation cost is unpredictable — easy problems might generate short trajectories, hard problems generate up to the full 4096 tokens. For deployment, this variance in per-problem cost is challenging to budget for, unlike best-of-N or beam search methods where the compute cost is fixed by the chosen hyperparameters.

What evidence exists in the paper. The paper does not provide difficulty-stratified analysis of trajectory lengths, solve rates by problem difficulty, or the fraction of failures attributable to context truncation. The efficiency metrics in Appendix Table 1 and Figure 6 are aggregate averages. The paper does not discuss difficulty estimation or adaptive compute allocation as a design consideration. The contrast with test-time compute scaling approaches is not addressed.

Mitigation status. Not addressed. The paper does not propose difficulty estimation, early stopping, or adaptive budget allocation. This is a limitation of the current implementation rather than the SoS framework itself — in principle, the SoS language could include explicit difficulty-assessment operations or confidence estimates that would enable adaptive stopping, but these are not implemented. Section 7 mentions "subgoal setting" and "reflection and self-evaluation" as potential enhancements but does not connect these to adaptive compute allocation.

---

Limitation 6: The Model's Implicit Heuristic and World Model Are Uncharacterized

The assumption or constraint. The SoS language design makes a deliberate choice: heuristic evaluation is kept implicit — the model does not write out numerical scores or explicit value judgments for states. The paper argues this enables "the model to internalize abstract representations for them that can be improved with training" (Section 3). The state expansion function (performing arithmetic operations to generate child states) is also implicit — the model generates the result of an operation as text without calling an external calculator. The consequence is that the model learns an internal world model (for arithmetic transitions) and an internal heuristic (for evaluating state quality) that are never directly observed, measured, or validated.

The consequence. The paper makes claims about what the model has learned — that it develops an "internal world model" (Section 7), that it "discover[s] new search strategies" (Abstract, Section 6), that it shifts toward "multiply" heuristics after fine-tuning (Figure 5a) — but provides only indirect, aggregate evidence for these claims. The specific nature of the model's learned heuristic is unknown. Does the model compute something analogous to the sum heuristic? The multiply heuristic? A learned combination? Something entirely different that correlates with both? Without probing the model's internal representations or analyzing its state-visitation patterns at the level of individual trajectory decisions, the claim that the model has "discovered new strategies" is an interpretation of aggregate behavior rather than a demonstrated mechanism.

This limitation matters for both scientific understanding and practical deployment. Scientifically, the paper's contribution would be substantially stronger if it could characterize the learned heuristic — showing, for example, that the model's internal state evaluations (extractable via linear probes on hidden states) correlate more strongly with solution proximity than either the sum or multiply heuristic, demonstrating genuine discovery. Practically, a deployment that relies on an uncharacterized heuristic is harder to trust, debug, or improve — if the model makes a search error, there is no interpretable heuristic score to inspect for diagnosing why a particular branch was explored or pruned.

The same applies to the world model. The paper reports low arithmetic error rates (~2 per trajectory, Appendix Table 2) and low exploration error rates (0.8% per trajectory) as evidence that the model has learned the domain's transition dynamics. But this only shows the model's behavioral accuracy, not the structure of its internal world model. Does the model compute arithmetic by memorizing operation tables? By learning algorithmic procedures? By pattern-matching against training examples? The answer has implications for how well the world model would generalize to numbers outside the training range (the paper tests held-out targets from 10–100, but not out-of-range targets) and whether the world model is robust to adversarial inputs.

What evidence exists in the paper. The indirect evidence includes: low error rates (Appendix Table 2), alignment shifts after fine-tuning (Figure 5a), and the fact that the model solves some problems unsolvable by symbolic strategies (Figure 5c). The paper does NOT include: linear probes of hidden states to extract heuristic values, analysis of attention patterns during state evaluation decisions, comparison of model-generated state orderings against ground-truth distance-to-goal, or any other method for characterizing the learned heuristic. The paper does NOT test arithmetic generalization to out-of-distribution numbers. The paper does NOT analyze individual trajectories on the solved "difficult" problems to identify what novel search patterns emerged.

Mitigation status. Not addressed. The paper's design choice to make heuristics implicit is well-motivated for flexibility but leaves a gap in mechanistic understanding. Section 7 suggests that "explicitly representing state evaluations... could enhance the SoS framework," which would address this limitation in future work — making the heuristic explicit would sacrifice some flexibility but gain interpretability. The current work provides no tools for characterizing what the model learned, relying entirely on behavioral outcome metrics.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a conceptual reframing rather than an incremental improvement: it recasts search from an inference-time scaffolding problem into a training-time representation learning problem. This distinction matters because it redirects where the field looks for gains. Prior to this work, the dominant paradigm for improving LM reasoning was to keep the model fixed and wrap it in increasingly sophisticated extrinsic search systems (Tree of Thoughts, Graph of Thoughts, and their variants). Those approaches treat the LM as a stateless subroutine — a generator and evaluator of local moves — while the search strategy itself lives outside the model, hardcoded by the system designer. Stream of Search inverts this: the LM itself learns to be the search algorithm, internalizing exploration, backtracking, and heuristic evaluation into its autoregressive generation.

The magnitude of this shift is moderate but real. It does not render extrinsic methods obsolete — the paper explicitly says they are "likely to be more efficient for [more challenging] tasks" in the short term (Section 7). But it opens a research program that previously lacked empirical grounding: can language models learn search in a way that compounds through training, as opposed to search being applied only at inference? The 25-percentage-point gap between SoS (51.27%) and optimal-path training (25.73%) on Countdown (Figure 3a) establishes that the answer is yes, at least in a domain where search is structurally necessary. This is the paper's foundational result, and it changes the default question from "how should we scaffold search around this LM?" to "how should we train this LM so that search is what it does?"

A further shift concerns what counts as valuable training data. The standard practice in reasoning-data curation is to filter out errors — train only on correct solutions, verified trajectories, human-approved reasoning chains. The SoS result challenges this orthodoxy directly: the SoS model, trained on data that is 43% incorrect (unsuccessful search trajectories containing dead ends and arithmetic mistakes), nearly doubles the accuracy of the model trained on 100% correct optimal paths. This is not a marginal improvement from a clever data-augmentation trick; it is a diagnostic finding about the nature of the reasoning capability being trained. Learning to reason from only clean solutions teaches the model that reasoning is a forward-only process where every step is correct by construction. Learning to reason from search trajectories teaches the model that reasoning involves exploration, recognition of dead ends, and recovery — and it is precisely this recovery capability that the OP model lacks when it encounters its own errors at test time. The paper provides a concrete mechanism for this: exposure bias in autoregressive generation means the model must handle contexts containing its own mistakes, and only the SoS model has seen such contexts during training.

This finding reconciles conflicting intuitions in the literature. On one side, Huang et al. (2023) found that "large language models cannot self-correct reasoning yet" — but those experiments used models trained on clean data and prompted to self-correct at inference time, precisely the setting where no recovery behavior has been learned. On the other side, process-supervision work (Lightman et al., 2023) showed that training verifiers on intermediate-step correctness improves reasoning — but required expensive human annotation. SoS suggests a middle path: the model can learn recovery and self-correction from automatically generated search traces, without per-step human labels, provided the training data includes the full process of search including failures. The inability to self-correct is not an inherent limitation of autoregressive LMs; it is a consequence of never being shown what correction looks like.

The paper also subtly shifts the conversation around world models in language models. The criticism that next-token predictors lack internal models of the world (LeCun, 2023) is met not with an architectural change but with a data change. By forcing the model to simulate state transitions itself — writing out arithmetic operations and their results as part of the search trace — the SoS training setup requires the model to learn the domain's causal structure to generate coherent trajectories. The low arithmetic error rate (~2 per trajectory, Appendix Table 2) and low exploration anomaly rate (0.8%) provide behavioral evidence that the model has acquired a functional world model for Countdown, learned entirely from sequence prediction. This does not prove that the model's internal representations are causally structured, but it demonstrates that the right training data can produce world-model-like behavior from a standard autoregressive architecture without explicit environment models or simulators.

Research directions that become more attractive after this work:

• Training-data design for reasoning capabilities. The paper makes a strong case that the structure of training data — whether it includes process or only outcomes, whether it shows recovery from failure or only clean success — can be as important as data quantity. This opens a design space for reasoning datasets that deliberately include annotated search traces, mistake-correction pairs, and multi-path explorations alongside final answers.
• Intrinsic search as a pretraining objective. If search can be learned in a toy domain like Countdown, it may be learnable as a general capability during pretraining from appropriately structured data, analogous to how next-token prediction on code improves reasoning. The paper's language-for-search vocabulary (Section 3) provides a template for what such data might look like.
• Self-improvement loops for reasoning models. The STaR and APA results (Section 6), though modest in absolute gain (~5–6%, Figure 4a), demonstrate that a search-trained LM can generate improved training data for itself by exploring solution spaces and filtering for correctness. This is a proof-of-concept for the kind of self-play loop that produced superhuman game-playing agents (Silver et al., 2018), applied to an LM in a reasoning domain.

Research directions that become less urgent after this work:

• Sole reliance on increasingly complex extrinsic search architectures. The paper does not invalidate extrinsic methods — they remain more practical for many current applications — but it demonstrates that intrinsic search is possible and learnable. Research that exclusively explores new extrinsic scaffolding without considering whether the underlying LM could be trained to internalize that scaffolding is missing a dimension of the design space.
• Human annotation of per-step correctness for verifier training. The SoS approach achieves its results without any human labels — the training data comes entirely from symbolic solvers. For domains where symbolic search strategies can be implemented (even suboptimal ones), this suggests that the expensive human annotation pipeline of process supervision (Lightman et al., 2023) may not be the only path to teaching models about intermediate reasoning quality.

---

Follow-Up Research This Work Enables

Scaling SoS to harder and more open-ended search domains. Countdown is a closed, deterministic domain with a small state space and objectively verifiable solutions. The paper leaves entirely open how SoS transfers to domains where the search space is not enumerable, intermediate states are ambiguous, or correctness is graded rather than binary. A natural follow-up would apply SoS to a domain with qualitatively different search characteristics: mathematical proof (where the branching factor is high but evaluation requires deep domain knowledge), code generation with test-case feedback (where correctness is binary but the state space is the set of all programs), or open-ended planning (where goal states are underspecified). For code generation, the experiment would train a model on search trajectories generated by a symbolic solver that explores program variants guided by test-case pass/fail signals, then measure whether the SoS-trained model generates programs that pass test cases more reliably than an optimal-trajectory baseline. The key measurement would be not just accuracy but the type of errors: does the SoS model exhibit structured debugging behavior (identifying a failing test case, modifying the relevant code section, retesting) versus the OP model's tendency to hallucinate or restart from scratch? The paper's low exploration error rate (0.8%, Appendix Table 2) on Countdown suggests the model can learn valid state transitions; the question is whether this transfers when state transitions are complex operations like code edits rather than arithmetic.

Characterizing the learned heuristic and world model through mechanistic interpretability. The paper's most significant claim — that the model "discovers new search strategies" — is supported only by aggregate behavioral evidence (Figure 5c: ~4% of difficult problems solved) without any characterization of what the model learned. A high-priority follow-up would apply interpretability tools to the trained SoS model to extract its implicit heuristic and world model. The experiment would train linear probes on the model's hidden states at points where it must make exploration decisions (choosing which frontier state to expand next) to predict either the ground-truth distance to the solution or the symbolic heuristics used in training (sum and multiply). If the probes recover the symbolic heuristics with high accuracy, the model has internalized them. If the probes achieve above-chance prediction but below-perfect recovery, the model has learned something partially aligned with but not identical to the training heuristics. If a newly trained probe predicting solution distance outperforms probes trained to recover the symbolic heuristics, that would be direct evidence for a novel learned heuristic — the model's internal state representations encode information about solution proximity that is not captured by either the sum or multiply heuristic. A parallel experiment would examine attention patterns during backtracking operations: does the model attend to the specific state it is returning to, or does it rely on positional encoding to "count back" through the trajectory? The answer would reveal whether the model learns an explicit state-pointer mechanism or an approximate positional heuristic for backtracking.

Difficulty-adaptive generation budgets for SoS models. The current implementation generates the same autoregressive trajectory for every problem until "Goal Reached" or context limit, with no mechanism to allocate more tokens to hard problems or stop early on easy ones. A direct extension would incorporate a learned difficulty estimator into the SoS framework. The experiment would add an explicit "difficulty assessment" operation to the SoS language (e.g., the model outputs a confidence score after initial exploration) and train the model to condition its subsequent search breadth on this assessment. The training data would pair problems with the symbolic solvers' search statistics (nodes explored before solution, whether the heuristic-guided search succeeded or failed) as difficulty labels. The evaluation would measure: (1) whether the model stops earlier on easy problems (reducing average tokens per correct solution compared to the uniform-budget baseline), (2) whether the model reallocates saved tokens to hard problems (improving solve rate on the hardest quintile), and (3) whether a total-token-matched comparison — where the adaptive model and the uniform model get the same total token budget across a batch of problems — shows a net accuracy gain from adaptive allocation. This would connect SoS to the test-time compute scaling literature (Snell et al., 2024), testing whether the principle of difficulty-conditioned allocation applies to intrinsic search as it does to extrinsic verifier-guided sampling.

Combining intrinsic and extrinsic search in a hybrid system. The paper positions SoS against extrinsic search methods but never combines them. A natural integration would use a SoS-trained model as the proposal and evaluation module within a symbolic search framework, replacing the generic LM calls in Tree of Thoughts with a model that has been trained to generate search-coherent continuations. The experiment would compare three conditions on a harder version of Countdown (e.g., 5–6 input numbers, where the SoS model's 4096-token context window is insufficient for full search): (1) pure extrinsic search (Tree of Thoughts with a standard pretrained LM), (2) pure SoS (the model generates until context limit), and (3) hybrid (Tree of Thoughts where each node expansion calls the SoS model, which generates a partial search subtree rather than a single proposed next step). The prediction is that the hybrid would outperform both pure conditions: the SoS model's learned heuristic and backtracking capability would make its proposed subtrees more efficient than the standard LM's single-step proposals, while the symbolic framework would handle context management by stitching subtrees together across model calls, overcoming the context-length ceiling. The measurement would include total LM calls (efficiency), accuracy on 5-input problems (scale), and qualitative analysis of whether the hybrid system's search behavior exhibits patterns not present in either pure system.

Stress-testing the limits of learned search: adversarial problems and distribution shift. The paper tests generalization to held-out targets (10–100) and held-out input combinations, but these are within-distribution shifts. A critical stress-test would evaluate whether the SoS model's learned search strategies are robust to distributional shifts that break the training heuristics. Two specific experiments: (1) Out-of-range targets: test on problems with targets outside the 10–100 training range (e.g., targets of 150–200). The sum and multiply heuristics are defined identically regardless of target range, so if the model has internalized these heuristics in their general form, performance should degrade gracefully. If the model has learned heuristics specific to the 10–100 range (e.g., memorizing common factor patterns), performance should collapse. (2) Heuristic-adversarial problems: construct problems where the sum heuristic and multiply heuristic actively mislead — e.g., where the correct solution requires temporarily producing a number far from any factor of the target, which both heuristics would penalize. The symbolic strategies guided by these heuristics would systematically fail on such problems; the question is whether the SoS model's learned heuristic — which combines and potentially transcends the training heuristics — succeeds. A failure on these adversarial problems would bound the model's strategy discovery: it would suggest the learned heuristic is a sophisticated interpolation of the training heuristics rather than a genuinely novel evaluation function. A success would strengthen the "strategy discovery" claim considerably.

Multi-turn self-improvement with exploration rewards. The STaR experiments filter for binary correctness: trajectories either contain the solution (kept) or don't (discarded). The APA experiments add a length penalty to the reward but still optimize for solution-finding. Neither approach explicitly rewards informative failures — trajectories that explore promising but ultimately unsuccessful branches, or that efficiently exhaust a subspace and correctly conclude it contains no solution. A follow-up could design a reward function that credits such behavior: trajectories that explore diverse states receive a diversity bonus; trajectories that correctly identify a subtree as dead-end (by proving no solution exists within it) receive partial credit even if they don't find the global solution; trajectories that find novel solutions to previously unsolved problems receive a discovery bonus. Training with such rewards — possibly through a learned curiosity module or an intrinsic motivation objective — would directly incentivize the kind of exploration that the paper speculates leads to strategy discovery. The experiment would measure: does adding exploration rewards to the STaR or APA objective increase the model's solve rate on the "difficult" problem set (Figure 5c) beyond the ~4% achieved with correctness-only optimization? Does it increase the diversity of state-visitation patterns (measurable through the alignment metrics in Figures 3c and 5a)? A positive result would connect SoS to the exploration literature in reinforcement learning, suggesting that the right training incentives can push learned search beyond imitation of the training heuristics into genuine discovery.

---

Practical Applications and Downstream Use Cases

Automated data generation for reasoning benchmarks. The most immediate practical application is using SoS-trained models to generate high-quality search traces for new problems, creating training data for larger models or for domains where symbolic solvers are unavailable. The paper's finding that the SoS model matches or exceeds the average accuracy of the symbolic strategies that generated its training data (51.27% vs. ~57–60%, Figure 3a) — despite having no access to a ground-truth environment model — means a SoS model could serve as a data-generation engine for problems where writing a symbolic solver is difficult. For a new reasoning domain, one could: (1) implement a few simple, suboptimal heuristic search strategies to generate an initial SoS dataset, (2) train a SoS model on this data, (3) use the trained model to generate search traces on a larger set of problems, and (4) filter these traces for correctness to produce a clean solution dataset for downstream fine-tuning. The efficiency gain over running exhaustive symbolic search comes from the SoS model's learned heuristic: it explores fewer states than exhaustive search and finds solutions faster than random exploration. The paper's efficiency data (Figure 6, right: fine-tuned models explore fewer states per correct solution than the base model) suggests this pipeline would improve with iteration, as the SoS model's self-improvement loop generates increasingly efficient traces.

Edge deployment of reasoning systems with limited model capacity. The paper's core result — that a relatively small model (250M parameters) can learn to solve search problems that naive optimal-path training fails on — has direct implications for deploying reasoning systems on resource-constrained devices. A 250M-parameter model is small enough to run on a single GPU or even a high-end mobile device. If SoS training transfers to other reasoning domains (an open question, per the limitations), it would enable on-device planning and problem-solving agents that do not require cloud-based extrinsic search systems. The benefit is both latency (no network calls to a larger model or search orchestrator) and privacy (the entire search process stays on-device). The specific numbers that ground this application: the SoS model achieves 51.27% accuracy on held-out Countdown problems with a 250M-parameter model trained from scratch (Figure 3a), compared to 25.73% for the optimal-paths baseline. Extrapolating to a hypothetical deployed system, an on-device SoS-trained model could handle the majority of problems without cloud fallback, while the much lower accuracy of the OP-trained model (~26%) would make it unsuitable for standalone deployment.

Curriculum design for teaching reasoning to LMs. The paper's finding that exposure to mistakes and recovery is more valuable than exposure to clean solutions alone — supported by the 25-percentage-point gap between SoS and OP (Figure 3a) — has direct implications for how reasoning datasets should be constructed. Practitioners building training data for mathematical reasoning, code generation, or planning tasks should consider including not just correct solutions but annotated search traces that show productive exploration and error recovery. The specific recipe from the paper: generate diverse solution attempts using (even suboptimal) heuristic-guided search, serialize them with explicit backtracking and goal-checking operations, and include both successful and unsuccessful traces in the training set. The paper shows that 43% incorrect trajectories in the training data (285,501 correct out of 500,000 total, Section 4) does not degrade performance — it substantially improves it. This is a concrete, counterintuitive guideline for data curation that practitioners can apply immediately, with the caveat that it has only been validated on Countdown.

Self-improving data flywheels for reasoning systems in production. The STaR and APA results (Section 6), though modest in absolute gain (~5–6%), demonstrate a working self-improvement loop: the model generates trajectories, the system filters for correctness, and the model retrains on the improved data. In a production setting where a reasoning system handles a stream of user queries, this loop could run continuously: user problems that the system solves correctly get added to the training pool (with their generated search traces as process supervision), and the model periodically fine-tunes on these accumulated successes. The benefit is that the system improves on the specific distribution of problems users actually pose, without requiring human annotation of intermediate reasoning steps. The paper's finding that fine-tuning on model-generated correct trajectories improves accuracy (Figure 4a, Figure 4c) provides proof that this loop is viable, and the finding that policy improvement increases search efficiency (Figure 6, right) means the loop also reduces per-problem cost over time. The ~36% solve rate on previously unsolved training-set problems (Figure 5b) suggests the improvement would be most significant early in deployment, as the model accumulates successes on problems it initially fails.

---

When to Prefer This Method

The paper does not provide a systematic head-to-head comparison against named alternatives (extrinsic search methods, process supervision) with quantified trade-offs, so a formal decision matrix would impose structure the paper itself does not establish. However, the paper's positioning and results support several qualitative guidelines for when SoS is likely to be the right approach:

• When a symbolic solver (even a suboptimal one) can be written for the domain. The SoS training pipeline requires search trajectories to train on — these must come from somewhere. If you can implement heuristic-guided BFS/DFS/other search for your domain (as the paper did for Countdown), SoS is viable. If no search algorithm exists and you cannot build even a weak one, SoS cannot be bootstrapped. The symbolic solver does not need to be optimal — the paper's training solvers succeeded on only 57% of problems — but it does need to exist.

• When inference-time efficiency matters more than absolute accuracy on the hardest problems. The SoS model executes search in a single autoregressive pass without calling an external controller. This is faster per search step than extrinsic methods (which make multiple LM calls per step) but is bounded by the model's context window — it cannot search deeper than what fits in 4096 tokens. For problems within this depth budget, SoS provides more efficient inference; for problems requiring deeper search, extrinsic methods with symbolic state tracking are necessary. The paper does not provide latency or token-count comparisons against extrinsic methods, so this trade-off is qualitative, not quantitative.

• When you want the search capability to improve through self-play or continued training. Extrinsic search systems wrap a fixed search algorithm around a frozen LM — the LM does not get better at search through use. SoS models can improve through STaR/APA-style self-improvement loops (Section 6) because the search policy lives in the model's weights. If the deployment involves a stream of problems where accumulating successful search traces and periodically retraining is feasible, SoS is preferable to static extrinsic scaffolding.

• When interpretability of the search process matters. The SoS model outputs the entire search trajectory as human-readable text (Figures 7–8 in the appendix show examples). This includes explicit backtracking statements, goal checks, and the sequence of states explored. Debugging a failed search means reading the generated trajectory and identifying where the model went wrong — the same process as debugging a symbolic search trace. Extrinsic methods typically expose the search tree structure through a separate interface, but the LM's internal reasoning at each node may be opaque. If downstream users or auditors need to inspect the full reasoning process, SoS's serialized format is an advantage.