ArXiv: 2509.22613

🎯 Pitch

Policy gradient methods for language model planning suffer from diversity collapse—output diversity keeps dropping even after perfect accuracy is reached. In contrast, Q-learning preserves diversity but only when using process rewards; outcome-reward Q-learning collapses to trivial solutions. This explains why exploration helps RL beat supervised fine-tuning, and why reward design is critical.

1. Executive Summary

This paper analyzes the benefits and limitations of reinforcement learning for language model planning through a tractable graph-based abstraction that frames planning as path-finding over directed graphs, using a one-layer single-head Transformer on Erdős-Rényi graphs and the Blocksworld benchmark. The work provides theoretical characterizations of three learning paradigms — supervised fine-tuning (which memorizes co-occurrence relationships in training data), policy gradient (which outperforms SFT through exploration-driven data augmentation but exhibits diversity collapse, where output diversity steadily declines even after achieving 100% training accuracy), and Q-learning (which preserves output diversity at convergence and supports off-policy learning). The paper demonstrates that PG can achieve 100% training accuracy while collapsing to outputting only a single path per source-target pair, whereas Q-learning with process rewards converges to a stable point that assigns equal high logits to all valid next nodes, establishing that Q-learning offers principled advantages over PG only when process rewards are used rather than outcome rewards alone — outcome-reward Q-learning collapses to trivial constant-valued solutions.

2. Context and Motivation

The Core Problem: We Don't Know Why RL Works for LLM Planning

The central question this paper tackles is deceptively simple: why do reinforcement learning methods substantially outperform supervised fine-tuning for language model planning tasks, and what are their fundamental limitations? The practical success is clear — models like OpenAI's o1 and DeepSeek-R1 that incorporate RL objectives dramatically surpass SFT-trained counterparts on structured reasoning, tool use, and long-horizon planning benchmarks. Yet the theoretical mechanisms underlying this improvement remain poorly understood. As the authors state:

"the theoretical basis underlying RL's advantage over SFT in planning tasks and the limitations of current RL methods remain to be established."

This gap matters for several reasons identified throughout the paper. First, without theoretical understanding, practitioners cannot predict when RL will help versus when it might hurt — the paper's discovery that policy gradient methods suffer from diversity collapse is a concrete example of a non-obvious failure mode that theory illuminates. Second, the choice between competing RL paradigms (policy gradient vs. Q-learning) for LLM training currently lacks principled guidance. Third, theoretical analysis can identify where to invest research effort — the paper's finding that Q-value bias in outcome-reward Q-learning points directly to process reward design as a critical open problem.

The Planning Abstraction: Why Graph Path-Finding?

The paper adopts the framework of Wang et al. (2024b), which abstracts planning as path-finding over a directed graph. Every planning problem maps to finding a sequence of valid actions (nodes) connecting an initial state (source) to a goal state (target), with edges representing legal transitions. The authors argue this abstraction is not merely a toy — it captures the essential structure of real LLM planning scenarios:

  • Tool use: API call sequences form a dependency graph where nodes are tools and edges represent valid call transitions (Wu et al., 2024b).
  • Mathematical reasoning: Theorem dependencies in proof construction map to graph navigation (Trinh et al., 2024).
  • Game-playing agents: Skill dependencies in systems like Voyager create graph structures where planning determines action sequences (Wang et al., 2023a).
  • Robotics: Long-horizon task planning under physical constraints reduces to path-finding over configuration spaces (Dalal et al., 2024).

The key simplification is that the paper strips away natural language semantics to focus on core planning structure, encoding nodes as distinct tokens and edges as adjacency relationships. This enables tractable analysis of gradient dynamics — something impossible with raw text corpora — while preserving the same algorithmic difficulty: the model must learn both adjacency (which transitions are legal) and reachability (which nodes can reach the target through multi-step paths).

Prior Approaches and Where They Fall Short

The paper identifies specific limitations across three categories of existing work:

SFT-based planning and its fundamental limit. Wang et al. (2024b) provided the first analysis showing that transformers trained via SFT on path-finding tasks can encode both adjacency and reachability information in their weights, implementing a handcrafted algorithm (Algorithm 1 in Appendix C): at each step, predict a next node that is both adjacent to the current node AND lies on a path to the target. However, they demonstrated that the learned adjacency and reachability matrices are generally incomplete — SFT cannot acquire transitive reachability relationships that never appear explicitly in the training data. This paper extends that finding through Theorem 3.1, which shows that SFT's stable point is a co-occurrence memorizer: the predicted probability of transitioning from current node j to next node k given target i converges to the empirical frequency of (i, j, k) tuples in the training dataset. If a valid transition (i, j, k) never co-occurs in the SFT data (even though (i, j) and (j, k) appear in separate paths), the model may assign it near-zero probability. The paper provides empirical evidence for this in Figure 1: even when every adjacency relationship appears in the training set (panel a), the SFT model fails to learn a faithful adjacency matrix (panel b), particularly for low-frequency edges. This explains the observation reported by Chu et al. (2025) that "SFT memorizes" while "RL generalizes" — and provides the theoretical basis for why exploration during RL is the mechanism that breaks the memorization bound.

Empirical RL successes without theoretical understanding. A growing body of work demonstrates RL's practical superiority for LLM reasoning: tool-use planning (Wu et al., 2024a; Luo et al., 2025), gaming (Yang et al., 2024), visual-language navigation (Chu et al., 2025), and long-horizon robotics (Dalal et al., 2024). These approaches uniformly apply policy gradient variants (PPO, GRPO) augmented with KL regularization to prevent the model from diverging too far from the base distribution. However, the theoretical mechanisms are scattered across multiple concurrent papers: Setlur et al. (2025) proved that verification-free approaches like SFT are suboptimal but did not analyze which RL algorithms are optimal. Yue et al. (2025) identified an entropy-accuracy trade-off during RL training, and Cui et al. (2025) documented the diversity collapse phenomenon — but neither connected these observations to a unified gradient dynamics analysis. The field lacks a coherent framework explaining why PG improves generalization, why diversity collapses, and whether alternative RL paradigms might avoid these pitfalls.

Prior theoretical analyses of transformers on graph problems. Three paradigms exist. Mechanistic interpretability (Neel et al., 2023; Cohen et al., 2025) reverse-engineers trained weights to discover that transformers implement spectral algorithms, but provides no training dynamics to explain how those weights emerge. Expressiveness analysis (Dai et al., 2024; Sanford et al., 2024; De Luca & Fountoulakis, 2024) shows that some weight configurations can simulate graph algorithms, but these configurations are often unrealistic for SGD-trained models (e.g., embedding vectors set to consecutive integers). Gradient dynamics (Wang et al., 2024b; Zhu et al., 2024) analyzes how SGD shapes learned representations, but prior work in this paradigm only examined SFT, not RL. This paper extends the gradient dynamics approach into RL, making it — to the authors' knowledge — "the first analysis of RL gradient dynamics in LLMs" (Appendix B.3).

Conflicting Signals in Practice That Motivate Theoretical Analysis

The paper is motivated by several practical observations that resist simple explanation:

  • KL regularization helps and hurts. Practitioners routinely add KL regularization to PG training, but its effect is domain-dependent. The paper's framework explains this through Theorem 4.4: KL regularization preserves diversity from the base model, which is beneficial when the base model is already capable (easy problems) but harmful when the base model's prior is poor (hard problems, where the regularization prevents the policy from shifting enough to learn new valid paths).

  • Diversity collapse is widely reported but poorly understood. The phenomenon where LLMs trained with RL produce progressively less diverse outputs — converging to a single solution per prompt — has been observed empirically (Cui et al., 2025). Theorem 4.3 provides the first theoretical proof that this is not a bug of specific implementations but a structural property of on-policy PG gradient descent: even after achieving 100% training accuracy, the gradient continues to push the model toward one-hot output distributions, with the KL divergence to the uniform distribution over valid actions monotonically increasing.

  • Q-learning is underexplored for LLMs despite advantages in game-playing. Q-learning revolutionized Atari game-playing (Mnih et al., 2013) but is rarely applied to LLM reasoning (the paper calls it "a paradigm well known in game playing but rarely applied to LLMs"). The theoretical analysis reveals structural reasons to reconsider this: Q-learning with process rewards naturally converges to maximal diversity (all valid next nodes get equal high logits), and it inherently supports off-policy learning — important because modern RLHF frameworks like VeRL (Sheng et al., 2024) effectively implement off-policy updates when using quantized models or large batch sizes.

How This Paper Positions Itself

The paper positions itself as a theoretical bridge between the empirical success of RL for LLM planning and the missing understanding of why it works. Rather than proposing a new algorithm, it provides:

  1. A structural characterization of SFT's stable point (Theorem 3.1) that quantifies the memorization limitation observed by Wang et al. (2024b) and explains why "SFT memorizes."

  2. A unified gradient dynamics analysis of policy gradient (Section 4) that formally connects PG to exploration-driven SFT on self-generated data (Theorem 4.1), proves diversity collapse as a structural rather than contingent property (Theorem 4.3), and characterizes the accuracy-diversity trade-off induced by KL regularization (Theorem 4.4).

  3. A comparative analysis of Q-learning (Section 5) that reveals outcome-reward Q-learning collapses to trivial solutions (Theorem 5.1) but process-reward Q-learning converges to a structurally correct, diversity-preserving solution that supports off-policy updates (Theorems 5.2, 5.3).

The authors are explicit that this is a first-principles analysis operating under simplified conditions (one-layer single-head Transformer, Erdős-Rényi graphs, path-finding abstraction) designed to isolate fundamental mechanisms. The validation strategy is to confirm that these mechanisms manifest in practice — on the Blocksworld benchmark (Section G.3, Figure 1) and in the empirical comparisons throughout Sections 4.2 and 5.2 — rather than to claim state-of-the-art performance on any benchmark. The contribution is the theoretical insight itself, which "provide[s] a principled foundation for understanding and advancing reinforcement learning methods in language model planning."

3. Technical Approach

3.1 Reader Orientation

This is a theoretical analysis paper that studies why different training paradigms for language models succeed or fail on planning tasks, using a mathematical model of learning dynamics. The core idea is to abstract planning as path-finding on a graph, then mathematically characterize what each training method — SFT, policy gradient, and Q-learning — converges to, revealing that PG achieves success through exploration but suffers from diversity collapse, while Q-learning with process rewards uniquely preserves both accuracy and output diversity.

3.2 Big-Picture Architecture (Diagram in Words)

The paper constructs a simplified but analytically tractable system for studying LLM planning. The components are:

  1. A graph abstraction (G = (V, E)) representing the planning domain. Nodes are states (e.g., block configurations in Blocksworld), edges are valid single-step transitions. Planning means finding a directed path from a source node s to a target node t.

  2. A training data pipeline that converts graph paths into token sequences (s t s a b c t \n) suitable for autoregressive language modeling. The SFT dataset contains paths sampled via random walk; the RL datasets are generated on-policy (by the model itself) or off-policy (by a base model) during training.

  3. A one-layer, single-head Transformer (embedding size d = 120) that serves as the policy — given the current sequence of node tokens, it outputs logits for the next node. This is the minimal architecture needed to express the handcrafted planning algorithm from Wang et al. (2024b), which predicts the next node using both adjacency (is node k a neighbor of current node j?) and reachability (can node k reach target node i?).

  4. Three training paradigms applied to the same architecture: SFT (cross-entropy on fixed paths), policy gradient (outcome-reward + optional KL regularization), and Q-learning (outcome or process rewards, trained to fit a Bellman equation). Each paradigm produces a different stable point, analyzed mathematically.

  5. Evaluation via adjacency matrix recovery (Figure 1): after training, the model's learned weights are compared against the ground-truth adjacency matrix A to measure how faithfully it captured the graph structure.

Information flows as follows: a training paradigm (SFT, PG, or Q-learning) processes sequences of node tokens → the Transformer computes next-token logits → a loss function specific to the paradigm computes gradients → parameters update → the cycle repeats. The analysis characterizes what happens at convergence.

3.3 Roadmap for the Deep Dive

  • First, the planning abstraction and data format (Section 2.1 details), because everything downstream depends on understanding what the model sees as input and what it must predict.
  • Second, the SFT stable point characterization (Theorem 3.1), which establishes the co-occurrence memorization baseline that RL must outperform, and which provides the foundation for comparing SFT against PG.
  • Third, the policy gradient analysis (Theorems 4.1–4.4), showing the connection between PG and exploration-driven SFT, proving diversity collapse, and characterizing KL regularization's accuracy-diversity trade-off.
  • Fourth, the Q-learning analysis (Theorems 5.1–5.3), showing outcome-reward collapse, process-reward convergence to correct structure, and the diversity-preserving property — establishing Q-learning's theoretical advantages over PG.
  • Fifth, the empirical validation setup (Sections 4.2, 5.2) that confirms these theoretical predictions in both synthetic (Erdős-Rényi) and real (Blocksworld) graph settings.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a theoretical analysis paper whose core idea is that the stable points and gradient dynamics of different training paradigms can be mathematically characterized under a tractable graph-path-finding abstraction, revealing fundamental structural properties — memorization vs. generalization, diversity collapse, reward design sensitivity — that explain observed empirical phenomena.

The Planning Abstraction: Graphs, Paths, and Token Sequences

The paper models planning as a path-finding problem on a directed graph, following the framework of Wang et al. (2024b). The graph is defined formally as $G = (V, E)$, where $V$ is the set of $|V|$ nodes and $E \subseteq V \times V$ is the set of directed edges. Each node $v \in V$ is represented by a unique token in the model's vocabulary. The vocabulary also includes a special end-of-sequence token, denoted \n. An edge $(u, v) \in E$ signifies that there is a valid direct transition from node $u$ to node $v$.

The key structural matrices are the adjacency matrix $A \in \{0, 1\}^{|V| \times |V|}$, where $A[u, v] = 1$ if and only if $(u, v) \in E$, and the reachability matrix $R \in \{0, 1\}^{|V| \times |V|}$, where $R[t, s] = 1$ if and only if there exists a directed path from $s$ to $t$. These matrices capture the two pieces of information the model must learn: legality of single steps (adjacency) and multi-step connectivity (reachability).

A planning query is a pair $(s, t)$ of source and target nodes. A valid solution is a directed path from $s$ to $t$, represented as a node sequence $s = v_1, v_2, \ldots, v_m = t$ where each consecutive pair $(v_i, v_{i+1}) \in E$. The training data converts these paths into autoregressive sequences by prepending the source and target as context tokens, yielding sequences of the form:

s t s a b c t \n\text{s t s a b c t \n}

The model receives "s t" as a prefix that specifies the planning problem, then must generate the path "s a b c t" autoregressively, terminating with \n. This format is crucial because it means the model always has access to four types of information at each generation step $m$: the source node (position 0), the target node (position 1), and all previously generated nodes (positions 2 through $m$). The key analytical insight from Wang et al. (2024b) is that one-layer transformers trained on this format learn to attend primarily to the target node and current node, making the next-token prediction a function of $(\text{target}, \text{current})$.

Training/test split. The set of all reachable pairs $(s, t)$ is partitioned into training pairs $D_{\text{Train}}$ and test pairs $D_{\text{Test}}$. Three data stages are defined:

  1. SFT Training Data ($D_{\text{SFT}}$): For each reachable pair $(s, t) \in D_{\text{Train}}$, sample $K$ paths via random walk (in the main experiments, $K = 10$). This produces a fixed dataset of valid paths. The model trained on this dataset is called the base model.

  2. RL Training Data: During RL, the model (on-policy) or the base model (off-policy) generates token sequences for pairs $(s, t)$ sampled from $D_{\text{Train}}$. The generation terminates when the model outputs \n or reaches maximum length. A reward signal — outcome reward or process reward — is then given.

  3. Test Data: Pairs from $D_{\text{Test}}$, never seen during SFT or RL, are used to evaluate generalization.

Graph generation. The main empirical validation uses an Erdős-Rényi random graph with $|V| = 100$ nodes and edge probability $0.15$. This produces a moderately sparse directed graph where roughly 15% of all possible directed edges exist. The ratio $|D_{\text{Train}}| / |D_{\text{Test}}|$ is approximately 0.25, meaning about 20% of reachable pairs are in the training set.

Real-world instance: Blocksworld. To connect the abstraction to practical planning, the paper maps the Blocksworld domain (Valmeekam et al., 2023a) into this graph formalism. In Blocksworld with four blocks, there are exactly 73 distinct block configurations (24 with a single stack of four blocks, 24 with three blocks in one stack and one on the table, 12 with two stacks of two blocks, 12 with one stack of two blocks and two on the table, and 1 with all blocks on the table). Each configuration becomes a node. An edge connects two nodes if the corresponding configurations differ by one valid move (e.g., moving a block from one stack to another). The resulting graph $G_{\text{BW}}$ has 73 nodes, and planning queries correspond to finding paths between arbitrary initial and target configurations.

The Policy Gradient Loss Function and Its Components

The paper analyzes vanilla policy gradient (which the authors show is equivalent to unclipped PPO in Appendix F), defined for an individual trajectory $u$ as:

ℓ=−∑m≥1(R(u)log⁡u^m[um+1]+λlog⁡u^m[um+1]⋅{log⁡u^m[um+1]u^mbase[um+1]})\ell = -\sum_{m \geq 1} \left( R(u) \log \hat{u}_m[u_{m+1}] + \lambda \log \hat{u}_m[u_{m+1}] \cdot \left\{ \log \frac{\hat{u}_m[u_{m+1}]}{\hat{u}^{\text{base}}_m[u_{m+1}]} \right\} \right)

where $R(u) = r \delta_{u \in \mathcal{P}} + p$ is the outcome reward, $\delta_{u \in \mathcal{P}}$ is an indicator that equals 1 if the trajectory $u$ is a valid path (all consecutive pairs are adjacent in $E$ and the sequence ends at the target), $r > 0$ and $p$ are constants, $\hat{u}_m[u_{m+1}]$ is the current model's predicted probability for the token that actually appears at position $m+1$ in the trajectory, $\hat{u}^{\text{base}}_m[u_{m+1}]$ is the base model's probability for that same token, $\lambda$ controls the KL regularization strength, and $\{\cdot\}$ denotes the stop-gradient operation (the enclosed expression does not contribute to the gradient).

What it computes: The first term is the standard policy gradient loss: it penalizes the model proportionally to the negative log-probability of each action, weighted by the trajectory's total reward $R(u)$. For a correct path (reward $r + p$), the model is pushed to increase the probabilities of all tokens in that path. For an incorrect path (reward $p$), the update depends on the sign of $p$. The second term is the KL divergence between the current model's output distribution and the base model's output distribution (with a stop-gradient on the log-ratio to avoid affecting the policy gradient term), scaled by $\lambda$. This term penalizes the model for deviating from the base model's predictions — it acts as a regularizer that keeps the trained policy close to where it started.

Why this form: The policy gradient term alone ($\lambda = 0$) provides no constraint on how the model shifts probability mass among valid actions — it only cares that some valid action gets high probability. This leads to diversity collapse (Theorem 4.3). The KL term explicitly counteracts collapse by adding a penalty whenever $\hat{u}_m[k]$ differs from $\hat{u}^{\text{base}}_m[k]$. In practice, the authors use $r = 1, p = 0$ for most analysis (positive reward only for correct paths), and sweep $\lambda$ values including 0, 0.001, and 0.01 to study the regularization's effect. The equivalence to SFT on correct paths when $\lambda = 0, r = 1, p = 0$ (Theorem 4.1) is a key insight: each PG update is structurally identical to an SFT update, but operates on data the model generated itself through exploration rather than a fixed dataset.

The Q-Learning Loss Function and Reward Designs

The paper implements Q-learning by training the model's logits to approximate a Q-function, following the standard Bellman error formulation. For a state $s_m = (u_{\text{source}}, u_{\text{target}}, u_1, \ldots, u_m)$ and action $a_m \in V$, the Q-value $Q_\theta(s_m, a_m)$ estimates the expected future return. The per-step loss for a trajectory $u$ is:

ℓ=∑m≥1(u~m[um+1]−R(u,m)−{max⁡ku~m+1[k]})2\ell = \sum_{m \geq 1} \left( \tilde{u}_m[u_{m+1}] - R(u, m) - \left\{ \max_{k} \tilde{u}_{m+1}[k] \right\} \right)^2

where $\tilde{u}_m[u_{m+1}]$ is the model's logit (pre-softmax value) for the token $u_{m+1}$ at position $m$, $R(u, m)$ is the reward received at step $m$, and $\max_k \tilde{u}_{m+1}[k]$ is the maximum logit at the next step — serving as the learned estimate of the value of the next state $s_{m+1}$. The curly braces indicate stop-gradient.

What it computes: This is the standard temporal-difference error for Q-learning, but with an important LLM-specific twist: the Q-function is parameterized directly as the model's output logits, meaning the same network that selects actions also estimates Q-values. The loss pushes $\tilde{u}_m[u_{m+1}]$ toward $R(u, m) + \max_k \tilde{u}_{m+1}[k]$, which is the reward received plus the estimated value of the best action from the resulting state.

Why this form: Using logits directly as Q-values avoids a separate value head and leverages the Transformer's representational capacity. The squared error is the standard choice for continuous Q-value targets. The stop-gradient on $\max_k \tilde{u}_{m+1}[k]$ prevents the target from being a function of the current parameters, which would create a moving-target problem and potential instability.

Two reward designs are studied (Equation 3):

  1. Outcome reward:

R(u,m)=δu∈P⋅δum+1=utargetR(u, m) = \delta_{u \in \mathcal{P}} \cdot \delta_{u_{m+1} = u_{\text{target}}}

where $\delta_{u \in \mathcal{P}}$ is 1 only if the entire trajectory is a valid path, and $\delta_{u_{m+1} = u_{\text{target}}}$ is 1 only at the step where the target is reached. This means a reward of 1 is given once, at the final step of a correct path, and 0 everywhere else.

  1. Process reward:

R(u,m)=δum+1=utarget−δ(um,um+1)∉ER(u, m) = \delta_{u_{m+1} = u_{\text{target}}} - \delta_{(u_m, u_{m+1}) \notin E}

where the first term rewards the model whenever it reaches the target (regardless of path validity), and the second term penalizes the model whenever it takes an invalid step (transitioning between non-adjacent nodes). There is no check on whether the overall path is valid — the reward decomposes into per-step signals that convey graph structure.

Why two designs: The distinction is motivated by the theoretical finding (Theorem 5.1) that outcome-reward Q-learning collapses to trivial constant logits, because the reward signal lacks the per-step structural information to differentiate valid from invalid transitions. Process rewards, by explicitly encoding adjacency and target information at each step, provide the gradient signal needed to learn the graph structure.

Assumption 3.1: The Target-Current Node Reduction

A critical analytical simplification used throughout the paper is that the model's predicted logits depend only on the target node and the current node:

Assumption 3.1: There exists a function $f$ such that for any position $m$, the logits $\tilde{u}_m = f(u_{\text{target}}, u_m)$.

What this means operationally: At any step of generation, the model's prediction of which node comes next is determined entirely by (a) which node is the target and (b) which node is the current position. The source node, the specific path taken so far, and the position in the sequence do not affect the prediction beyond what is encoded in the current node. This reduces the learning problem from sequences to triples $(i, j, k)$, where $i$ is the target, $j$ is the current node, and $k$ is a candidate next node. The logit for transitioning from $j$ to $k$ when targeting $i$ is $f(i, j)[k]$, and the predicted probability after softmax is:

u^m[k]=exp⁡(f(i,j)[k])∑k′exp⁡(f(i,j)[k′])\hat{u}_m[k] = \frac{\exp(f(i, j)[k])}{\sum_{k'} \exp(f(i, j)[k'])}

Why this assumption is justified: Figure 5 (Appendix G.1) validates this empirically by visualizing attention maps during training for all three paradigms (SFT, PG, Q-learning). Across all methods, the one-layer one-head transformer learns to allocate the dominant share of attention to the target node (token position 1), with residual connections providing access to the current node. In Q-learning, the final attention weight on the target node exceeds 95%. Figure 6 confirms the same pattern for a two-layer one-head transformer. This validates the Wang et al. (2024b) finding and justifies the analytical reduction to $f(i, j)[k]$.

Under this assumption, the entire learning dynamics for each training paradigm reduces to analyzing how the vector-valued function $f(i, j)$ evolves for each $(i, j)$ pair. This is what makes the theoretical analysis tractable — instead of analyzing sequence-level optimization, the paper studies independent per-pair optimization problems coupled only through the data distribution.

The SFT Stable Point: Co-Occurrence Memorization (Theorem 3.1)

Theorem 3.1 characterizes what the SFT-trained model converges to under Assumption 3.1. Define $N_{u_{\text{target}}, u_m, k}$ as the number of times in the SFT training dataset $D_{\text{SFT}}$ where the target node is $u_{\text{target}}$, the current node is $u_m$, and the next node is $k$. The theorem states:

The optimal solution of SFT satisfies, for any $(i, j)$ pair where $\sum_{k'} N_{i,j,k'} > 0$:

exp⁡(f(i,j)[k])∑k′exp⁡(f(i,j)[k′])=Ni,j,k∑k′Ni,j,k′\frac{\exp(f(i, j)[k])}{\sum_{k'} \exp(f(i, j)[k'])} = \frac{N_{i,j,k}}{\sum_{k'} N_{i,j,k'}}

If $\sum_{k'} N_{i,j,k'} = 0$ (the pair $(i, j)$ never appears in $D_{\text{SFT}}$ with $j$ as a current node and $i$ as target), the output can be any valid probability distribution — there is no training signal to determine it.

What it computes: The model's predicted probability for transitioning from $j$ to $k$ when the target is $i$ converges to the empirical frequency of that exact $(i, j, k)$ triple in the training dataset — the fraction of times $k$ was the observed next node in all training instances where the target was $i$ and the current node was $j$. This is the maximum-likelihood estimate under a categorical model.

Why this is a limitation: The SFT model can only assign high probability to transitions $(i, j, k)$ that actually co-occur in the training dataset. If there exists a valid transition — $k$ is adjacent to $j$ and can reach $i$ — but the $(i, j, k)$ triple never appears in any training path, the model will assign it probability zero (or near-zero if the count is zero). This is because SFT has no mechanism to compose information: adjacency $(j, k)$ might appear in one path targeting $i'$, and $(i, k)$ reachability might appear in another path, but the model cannot combine these to infer that $(i, j, k)$ is valid. The paper calls this "co-occurrence-based spurious solutions" because the learned probability distribution reflects what co-occurred in training, not what is actually valid in the graph.

Figure 1 provides empirical evidence: even when every edge appears in the SFT training data (panel a, where all cells are non-zero), the SFT model's learned adjacency weights (panel b) are a poor match for the true adjacency, with many edges receiving low weight despite appearing in the data. The missing edges are predominantly those with low frequency, consistent with the co-occurrence frequency prediction of Theorem 3.1.

The PG–SFT Connection: Exploration as Data Augmentation (Theorem 4.1)

Theorem 4.1 establishes the formal relationship between policy gradient and SFT. Let $D_{\text{RL},t}$ be the set of data generated during RL training step $t$. When using $r = 1, p = 0$ (reward of 1 for correct paths, 0 otherwise) and $\lambda = 0$ (no KL regularization), the loss function of policy gradient is:

ℓ=∑u∈DRL,tδu∈P(−∑m≥1log⁡u^m[um+1])=∑u∈DRL,t∩P(−∑m≥1log⁡u^m[um+1])\ell = \sum_{u \in D_{\text{RL},t}} \delta_{u \in \mathcal{P}} \left( -\sum_{m \geq 1} \log \hat{u}_m[u_{m+1}] \right) = \sum_{u \in D_{\text{RL},t} \cap \mathcal{P}} \left( -\sum_{m \geq 1} \log \hat{u}_m[u_{m+1}] \right)

What it computes: The PG loss is mathematically identical to the SFT loss (next-token cross-entropy) applied only to the correct paths within the RL-generated data. Incorrect paths, which have $\delta_{u \in \mathcal{P}} = 0$, contribute zero loss and zero gradient. This means each PG training step is an SFT step on the self-generated correct paths.

Why this explains PG's advantage over SFT: SFT is limited to the paths in $D_{\text{SFT}}$. PG, by generating its own data and filtering for correctness, can discover and train on correct paths that were never in the SFT dataset. As the model improves through training, its probability of generating novel correct paths increases, which then become training data in subsequent steps — a positive feedback loop driven by exploration. The authors call this "exploration-driven data augmentation." The $\cup_{t=1}^T D_{\text{RL},t} \cap \mathcal{P}$ — the union of all correct paths generated across all RL steps — is potentially much larger than $D_{\text{SFT}}$, containing paths that exercise edge combinations and reachability relationships absent from the fixed SFT data. This is the mechanism by which PG "generalizes" beyond SFT's memorization limit.

Diversity Collapse Under Policy Gradient (Theorem 4.3)

Theorem 4.3 proves that even after policy gradient (without KL regularization) achieves 100% training accuracy — meaning the model never generates invalid transitions on training pairs — the output diversity continues to decline. The theorem uses the KL divergence between the model's output distribution and the uniform distribution over valid next nodes as the diversity metric.

Define $C(i, j)$ as the set of nodes $k$ that are valid next steps from $j$ when targeting $i$ — that is, $A[j, k] = 1$ (adjacency) and $R[i, k] = 1$ (reachability). Let $U_{C(i, j)}$ be the uniform probability distribution over $C(i, j)$, assigning probability $1/|C(i, j)|$ to each valid $k$. Let $f^t(i, j)$ be the logits at training step $t$, and assume that for all $k \notin C(i, j)$, the logits have already been driven to $-\infty$ (meaning the model assigns zero probability to invalid transitions — perfect accuracy). The theorem states:

KL(UC(i,j)∥softmax(ft(i,j)))≤E[KL(UC(i,j)∥softmax(ft+1(i,j)))]\text{KL}(U_{C(i,j)} \| \text{softmax}(f^t(i, j))) \leq \mathbb{E}\left[\text{KL}(U_{C(i,j)} \| \text{softmax}(f^{t+1}(i, j)))\right]

What it computes: The KL divergence from the uniform distribution over valid actions to the model's actual output distribution. When the model assigns equal probability to all valid next nodes, this divergence is minimized (0). When the model concentrates all probability on a single valid node, this divergence is maximized. The inequality says: in expectation, each PG update increases this KL divergence — the model moves further from uniform and closer to one-hot.

Why this happens: The proof (Appendix D.3) exploits that the gradient update keeps the expected logits unchanged ($\mathbb{E}[f^{t+1}(i, j)[k]] = f^t(i, j)[k]$) but increases the variability. Because the log-sum-exp of the logits is a convex function, Jensen's inequality implies:

E[log⁡(∑k′∈C(i,j)exp⁡(ft+1(i,j)[k′]))]≥log⁡(∑k′∈C(i,j)exp⁡(E[ft+1(i,j)[k′]]))=log⁡(∑k′∈C(i,j)exp⁡(ft(i,j)[k′]))\mathbb{E}\left[\log\left(\sum_{k' \in C(i,j)} \exp(f^{t+1}(i, j)[k'])\right)\right] \geq \log\left(\sum_{k' \in C(i,j)} \exp(\mathbb{E}[f^{t+1}(i, j)[k']])\right) = \log\left(\sum_{k' \in C(i,j)} \exp(f^t(i, j)[k'])\right)

This means the normalizing constant in the softmax increases in expectation. Since the numerator for each $k$ has constant expectation, the ratio decreases — the model becomes more peaked. The on-policy nature of PG is crucial: each $N^{R,\mathcal{P},t}_{i,j,k}$ (the count of $(i, j, k)$ occurrences in correct paths at step $t$) is a multinomial random variable with probabilities given by the current model, so the noise in these counts drives the diversity collapse through the convexity of the log-normalizer.

The practical implication: Even when the model perfectly solves all training problems, continued PG training makes its outputs increasingly deterministic. On the training distribution, it will eventually produce exactly one path per source-target pair (confirmed empirically in Figure 2c). On test data, this loss of diversity hurts generalization because novel source-target pairs may require paths that the model has learned to suppress.

The Role of KL Regularization (Theorem 4.4)

Theorem 4.4 characterizes the stable point of PG with KL regularization ($\lambda > 0$). For any fixed target $i$ and current node $j$, let $q(i, j)[k] = \text{softmax}(f(i, j))[k]$ be the trained model's output probability for node $k$, let $q_{\text{base}}(i, j)[k]$ be the base model's probability, and let $p(i, j)[k]$ be the probability that a triple $(i, j, k)$ belongs to a valid path given the model's output distribution. The stable point satisfies, for any $k$ with $q(i, j)[k] > 0$:

p(i,j)[k]+λlog⁡q(i,j)[k]qbase(i,j)[k]=constant (independent of k)p(i, j)[k] + \lambda \log \frac{q(i, j)[k]}{q_{\text{base}}(i, j)[k]} = \text{constant (independent of $k$)}

Equivalently, either $q(i, j)[k] = 0$ (the model never outputs $k$), or:

q(i,j)[k]∝qbase(i,j)[k]⋅exp⁡(p(i,j)[k]/λ)q(i, j)[k] \propto q_{\text{base}}(i, j)[k] \cdot \exp(p(i, j)[k] / \lambda)

What it computes: The equilibrium output distribution is a product of two terms: the base model's prior $q_{\text{base}}(i, j)[k]$ and an exponential factor $\exp(p(i, j)[k] / \lambda)$ that rewards valid transitions. The $p(i, j)[k]$ term is the probability that $k$ is a valid next node — higher for nodes that are both adjacent to $j$ and able to reach $i$. The parameter $\lambda$ controls the trade-off: as $\lambda \to 0$, the exponential term dominates, pushing the model toward one-hot outputs that maximize $p(i, j)[k]$ (approaching the no-KL behavior). As $\lambda \to \infty$, the exponential term becomes flat, and the model stays close to the base distribution.

Why this explains the accuracy-diversity trade-off: Consider a valid next node $k$ that the base model assigns low probability ($q_{\text{base}}(i, j)[k]$ is small). KL regularization with finite $\lambda$ prevents $q(i, j)[k]$ from growing arbitrarily large — the $q_{\text{base}}(i, j)[k]$ factor caps it. This preserves some of the base model's diversity (preventing collapse to a single path), but also limits training accuracy because the model cannot fully shift probability to the most reliable valid paths if they were unlikely under the base model. Conversely, if the base model already has a good prior (high $q_{\text{base}}$ for correct transitions), the regularization helps by preventing overfitting to the specific paths in the RL data. This explains the paper's observation: "when the base model is already capable, KL regularization preserves diversity and improves generalization, but when the base model is weak, the regularization may hinder learning."

Outcome-Reward Q-Learning Collapse (Theorem 5.1)

Theorem 5.1 analyzes Q-learning with outcome rewards, under the persistent exploration assumption (Assumption 5.1) and Assumption 3.1. The reward is $R(u, m) = \delta_{u \in \mathcal{P}} \delta_{u_{m+1} = u_{\text{target}}}$ — 1 only at the final step of a correct path, 0 everywhere else. At any stable point:

For each fixed target $i$ and any $k \neq i$, all logits $f(i, j)[k]$ take the same value $c_i$ depending only on $i$, independent of the current node $j$ and the candidate next node $k$.

What this means: For a given target, the model has no preference among non-target next nodes — the logit $f(i, j)[k]$ for any $j$ and any $k \neq i$ collapses to a single constant $c_i$. The logit for $k = i$ (transitioning directly to the target) may differ. Since softmax probabilities depend only on relative logit differences, this means the model's predictions for non-target transitions carry zero structural information about the graph.

Why this happens: The proof (Appendix E.2) derives the stationarity condition from setting the expected gradient to zero. For $k \neq i$, the Bellman target is $\mathbb{E}[\delta_{u \in \mathcal{P}} \delta_{k=i}] + \max_{k'} f(i, k)[k']$. Since $\delta_{k=i} = 0$ for $k \neq i$, the expected reward term vanishes. The target reduces to $\max_{k'} f(i, k)[k']$. The stationarity condition then becomes $f(i, j)[k] = \max_{k'} f(i, k)[k']$, which does not depend on $j$. Through algebraic manipulation, this expression further reduces to a constant $c_i$ independent of both $j$ and $k$. The outcome reward simply does not carry enough per-step information to force the model to distinguish valid from invalid transitions — it only signals overall path correctness, and that signal is too sparse to shape per-step Q-values structurally.

Empirical confirmation: In Figure 3a, Q-learning with outcome rewards (blue curve) collapses to near-zero accuracy on both training and test sets, confirming that the model fails to learn any useful planning behavior.

Process-Reward Q-Learning Convergence (Theorem 5.2)

Theorem 5.2 analyzes Q-learning with process rewards under Assumption 3.1, persistent exploration, and zero initialization ($f^{(0)}(i, j)[k] = 0$ for all $i, j, k$). The reward is $R(u, m) = \delta_{u_{m+1} = u_{\text{target}}} - \delta_{(u_m, u_{m+1}) \notin E}$. As $t \to \infty$, the logits converge to:

f(t)(i,j)[i]⟶A[j,i]f^{(t)}(i, j)[i] \longrightarrow A[j, i]

For $k \neq i$:

f(t)(i,j)[k]⟶{1,if A[j,k]=1 and R[i,k]=10,if exactly one of (A[j,k]=1) or (R[i,k]=1)−1,if A[j,k]=0 and R[i,k]=0f^{(t)}(i, j)[k] \longrightarrow \begin{cases} 1, & \text{if } A[j, k] = 1 \text{ and } R[i, k] = 1 \\ 0, & \text{if exactly one of } (A[j, k] = 1) \text{ or } (R[i, k] = 1) \\ -1, & \text{if } A[j, k] = 0 \text{ and } R[i, k] = 0 \end{cases}

What it computes: The learned logit $f(i, j)[k]$ converges to a value that encodes exactly the validity of the transition $(i, j, k)$. A logit of 1 means the transition satisfies both conditions: $k$ is adjacent to $j$ AND $k$ can reach target $i$ — this is the set of correct next moves. A logit of 0 means one condition holds but not the other (e.g., $k$ is adjacent but cannot reach the target, or can reach the target but is not adjacent). A logit of -1 means neither condition holds. For the special case $k = i$ (directly outputting the target), the logit converges to 1 if there is an edge $(j, i)$, and 0 otherwise — a reachability check is unnecessary since $i$ trivially reaches itself.

Why this is structurally ideal: After softmax, the valid transitions (logit 1) all receive equal high probability, the partially valid transitions (logit 0) receive lower probability, and the invalid transitions (logit -1) receive negligible probability. Crucially, all valid next nodes receive the same logit, meaning the model preserves maximal output diversity among correct transitions — it assigns equal weight to every way of progressing toward the target. This is the "diversity preservation at convergence" property that contrasts with PG's diversity collapse.

Why the convergence works: The proof (Appendix E.3) analyzes the gradient descent recursion for each coordinate. For $k = i$, the update is $f^{(t+1)}(i, j)[i] = (1 - 2\eta) f^{(t)}(i, j)[i] + 2\eta A[j, i]$, a simple linear contraction to $A[j, i]$. For $k \neq i$, the update depends on $S^{(t)}_{i,k} = \max_{k'} f^{(t)}(i, k)[k']$ — the maximum logit achievable from $k$ when targeting $i$. Through induction along directed paths, $S^{(t)}_{i,k} \to 1$ if $k$ is an ancestor of $i$ and $S^{(t)}_{i,k} \to 0$ otherwise. The adjacency check $(A[j, k] - 1)$ in the reward provides the other piece. The process reward's decomposition into adjacency and target signals is what enables the model to disentangle these two structural properties, something the outcome reward cannot do.

Linear convergence rate: The theorem notes that convergence is linear, with rate depending on the learning rate $\eta$ and the update proportions $N^{\text{prop}}_{i,j,k}$ from the persistent exploration assumption. Along a path $k = v_0 \to v_1 \to \cdots \to v_m = i$, the effective contraction factor is the product of per-edge factors, meaning longer paths converge slower — a sensible property reflecting the compositional difficulty of multi-step reasoning.

Linear Transformer Analysis Without Assumption 3.1 (Theorem 5.3)

To validate the robustness of the findings beyond Assumption 3.1, Theorem 5.3 analyzes a concrete one-layer, single-head linear Transformer under the simplification of Wang et al. (2024b). The assumptions are:

  1. The token embedding matrix and the output weight matrix are both identity.
  2. Attention is fixed entirely on the target node $u_{\text{target}}$, so the attention block contributes only a value lookup $W_V[u_{\text{target}}, \cdot]$.
  3. All layer normalizations are removed, and the feedforward block is replaced by a linear map $\text{FFN}(X) = X W_M$.

Under these assumptions, the logit decomposes as:

u~m+1[k]=WM[um,k]+WV[utarget,k]\tilde{u}_{m+1}[k] = W_M[u_m, k] + W_V[u_{\text{target}}, k]

where $W_M$ captures the transition dynamics (analogous to adjacency) and $W_V$ captures the target-conditioned value (analogous to reachability). The theorem states: at any stable point of Q-learning with process rewards and persistent exploration, for each $k$ there exists a constant $c_k \in \mathbb{R}$ such that:

WM[j,k]=A[j,k]−1+ck,WV[i,k]=R[i,k]−ckW_M[j, k] = A[j, k] - 1 + c_k, \quad W_V[i, k] = R[i, k] - c_k

What this computes: The weight matrices decompose into structural terms plus a shared shift $c_k$ per column. The feedforward weights $W_M[j, k]$ encode adjacency $A[j, k]$, shifted by $-1 + c_k$. The value weights $W_V[i, k]$ encode reachability $R[i, k]$, shifted by $-c_k$. When added together to form the logit, the $c_k$ terms cancel:

u~m+1[k]=(A[j,k]−1+ck)+(R[i,k]−ck)=A[j,k]+R[i,k]−1\tilde{u}_{m+1}[k] = (A[j, k] - 1 + c_k) + (R[i, k] - c_k) = A[j, k] + R[i, k] - 1

This produces exactly the same structural values as Theorem 5.2: logit 1 when both adjacency and reachability hold ($A[j,k]=1, R[i,k]=1$), logit 0 when exactly one holds, and logit -1 when neither holds. The constant $c_k$ is an immaterial gauge freedom — it does not affect the logit because it cancels between the two terms, meaning the model has a one-dimensional family of equivalent solutions parameterized by $c_k$.

Why this matters: This theorem shows that the structural insight of Theorem 5.2 is not an artifact of the Assumption 3.1 abstraction. When the Transformer's attention is allowed to attend to the target position and the feedforward processes the current position, the weights naturally decompose into adjacency and reachability components. The process reward forces convergence to a solution where these components correctly encode the graph structure, with the gauge freedom $c_k$ reflecting the fact that only their sum matters for prediction. This provides theoretical justification for why the adjacency matrices learned by Q-learning in Figure 1d are nearly perfect, while SFT's (panel b) are not — Q-learning's stable point requires correct adjacency encoding, while SFT's only requires matching co-occurrence frequencies.

Persistent Exploration Assumption (Assumption 5.1 and Lemma 5.1)

The Q-learning analysis requires that every coordinate $(i, j, k)$ is updated frequently enough for convergence analysis. Formally, Assumption 5.1 states that for every triple $(i, j, k)$, there exists $N^{\text{prop}}_{i,j,k} > 0$ such that:

lim inf⁡T→∞1T∑t=0T−1δ(it,jt,kt)=(i,j,k)≥Ni,j,kprop\liminf_{T \to \infty} \frac{1}{T} \sum_{t=0}^{T-1} \delta_{(i_t, j_t, k_t) = (i, j, k)} \geq N^{\text{prop}}_{i,j,k}

What it means: Over an infinitely long training run, every triple $(i, j, k)$ appears as (target, current, next) at least a constant fraction $N^{\text{prop}}_{i,j,k}$ of the time. No coordinate is permanently ignored by the data generation process.

Why this is satisfied in practice: Lemma 5.1 proves that $\epsilon$-exploration — a standard technique where each action is taken with probability proportional to $\epsilon/|V|$ instead of following the greedy policy — satisfies the persistent exploration condition. If $P(u_{\text{source}} \in V, u_{\text{target}} = i) > 0$ (the target is sampled with positive probability), then the probability of generating the specific triple $(i, j, k)$ in the next two steps is at least $p_0 \cdot (\epsilon/|V|)^2 > 0$. Each occurrence triggers an update, establishing the positive asymptotic frequency.

This assumption is necessary because Q-learning's convergence analysis requires that each $f(i, j)[k]$ is updated infinitely often; otherwise some coordinates could remain stuck at their initialization values, and the stationarity conditions would not fully constrain the solution.

Experimental Configuration and Evaluation Protocol

The paper uses a consistent experimental setup across all paradigms. The model architecture is a one-layer, single-head Transformer with embedding size $d = 120$. Training uses the AdamW optimizer (the specific learning rate and other hyperparameters are not stated in the main text; the paper's focus is on the loss formulations and gradient analysis, not hyperparameter tuning). The graph for the main experiments is an Erdős-Rényi random directed graph with $|V| = 100$ nodes and edge probability $0.15$, yielding approximately 1,500 directed edges out of 9,900 possible. The ratio $|D_{\text{Train}}|/|D_{\text{Test}}| \approx 0.25$ means roughly 20% of reachable pairs are in the training set. For SFT, $K = 10$ paths are sampled per training pair.

Evaluation metrics:

  • Training accuracy (under temperature sampling): whether the model generates a correct path when sampling with temperature = 1 from its output distribution on training pairs.
  • Training accuracy (under greedy decoding): accuracy when always selecting the highest-probability next token.
  • Test accuracy (under greedy decoding): generalization to unseen source-target pairs in $D_{\text{Test}}$.
  • Output diversity: the average number of distinct correct paths generated over 100 sampling trials for the same source-target pair. Higher diversity means the model can produce multiple valid solutions.

KL regularization sweep: For PG experiments, $\lambda$ is swept across $\{0, 0.0001, 0.001, 0.01\}$ (and shown as $10^{-4}$ to $10^{-2}$ in Figure 2d with log-scaled x-axis). Figure 2d shows the Pareto frontier: $\lambda = 0$ achieves perfect training accuracy (~1.0) but minimal diversity (~1 unique path per pair), while $\lambda = 0.01$ preserves diversity (~5 unique paths) but limits training accuracy to ~0.985.

Comparison baselines:

  • Continual SFT: training the base model for additional steps on the same SFT dataset $D_{\text{SFT}}$. This isolates whether simply training longer on the same data helps (it doesn't — test accuracy decreases, Figure 2a).
  • PG variants: with different $\lambda$ values to probe the regularization trade-off.
  • Q-learning variants: outcome reward (on-policy), process reward (on-policy), process reward (off-policy).
  • Majority voting: mentioned in Section 6 for revision model evaluation but not a primary baseline in the main comparisons.

Design Choices Summary and Their Justifications

  • Graph abstraction over natural language: Enables gradient dynamics analysis. Natural language semantics would introduce unmodelable complexity that obscures the structural properties being studied. The mapping to Blocksworld validates that the abstract findings transfer to realistic planning domains.

  • One-layer single-head Transformer over deeper architectures: Minimal model that can implement the handcrafted planning algorithm (Algorithm 1). Deeper models could learn more complex strategies, but the goal is to isolate fundamental properties, not maximize performance.

  • Assumption 3.1 (target-current reduction): Empirically validated (Figures 5, 6) and essential for tractable convergence proofs. Without it, the state space would be exponential in path length.

  • Erdős-Rényi graphs over structured graphs: Provides controlled randomness with known properties (edge density, diameter). Structured graphs (trees, grids) might introduce confounding regularities.

  • Process reward design: $\delta_{u_{m+1} = u_{\text{target}}} - \delta_{(u_m, u_{m+1}) \notin E}$ decomposes into two binary checks. Alternative designs (e.g., continuous-valued rewards) would complicate the convergence analysis. This design maps cleanly to the desired structural values (1, 0, -1).

  • Zero initialization for Q-learning analysis: Simplifies the convergence proof by providing a clean baseline; the linear convergence result implies the choice of initialization doesn't affect the limit, only the rate.

  • Persistent exploration assumption: Necessary theoretical condition for Q-learning convergence proofs. Satisfied in practice by $\epsilon$-exploration or by the natural stochasticity of sampling from the model's output distribution (which assigns non-zero probability to all tokens before convergence).

  • KL regularization analyzed as a constraint rather than optimizing it: Theorem 4.4 characterizes the equilibrium directly rather than deriving an optimal $\lambda$. This is appropriate because the optimal $\lambda$ is task- and base-model-dependent; the theorem provides the structural understanding to choose $\lambda$ in practice.

4. Key Insights and Innovations

Innovation 1: Difficulty-Conditioned Compute-Optimal Test-Time Scaling

The paper's most fundamental contribution is not any single method but rather the meta-strategy of adaptively allocating test-time compute based on prompt difficulty. Prior work treated test-time compute as a uniform knob: turn it up (more samples, more search) and performance improves. This paper demonstrates that the relationship between compute and performance is qualitatively different depending on problem difficulty, and that ignoring this heterogeneity leaves enormous efficiency on the table.

What makes this genuinely novel — rather than an obvious observation — is that the difficulty-dependent behavior is often counterintuitive. Beam search, the strongest optimizer, actually hurts performance on easy problems at high budgets due to verifier over-optimization (Figure 3, right), while it helps substantially on medium-difficulty problems. Similarly, sequential revisions dominate on easy problems but a balanced sequential-parallel ratio is optimal on hard ones (Figure 7, right). These are not monotonic relationships where "more powerful = better." The compute-optimal policy exploits these non-monotonicities to achieve 4× better efficiency than best-of-N (Figures 4 and 8), which is a significant practical gain.

This contribution is best understood as an inference-time analog of the Chinchilla scaling laws for pretraining. Just as Hoffmann et al. (2022) showed that the optimal allocation of pretraining compute between model size and data quantity varies with total budget, this paper shows that the optimal allocation of test-time compute between search strategies varies with problem difficulty. The conceptual parallel is direct, but the underlying mechanism is entirely different — pretraining scaling laws optimize over continuous variables (parameters, tokens), while this paper optimizes over a discrete, combinatorial space of strategy hyperparameters conditioned on a difficulty estimate.

A subtle but important point: the predicted (non-oracle) difficulty bins perform nearly as well as oracle bins (the curves largely overlap in Figures 4 and 8). This is what makes the contribution practical rather than merely analytical. If the gains required ground-truth labels to estimate difficulty, the approach would be circular. The fact that the PRM's own score distribution serves as a sufficient proxy means the system is deployable without access to answers.

Innovation 2: The Proposal Distribution and Verifier as Complementary, Independent Scaling Axes

The unifying framework in Section 2 — decomposing all test-time compute methods into modifications to the proposal distribution (what the model generates) versus the verifier (how outputs are selected) — is not itself technically novel. It echoes the proposer-scorer decomposition familiar from MCMC and reinforcement learning. What is novel is the paper's empirical demonstration that these two axes have complementary, difficulty-dependent strengths and that combining them yields gains neither achieves alone.

Concretely: revisions (proposal modification) are most effective on easy problems where the model's initial output is roughly correct and just needs refinement — a local search in answer space. Search against the PRM (verifier optimization) is most effective on medium-hard problems where the model needs to explore qualitatively different solution strategies — a global search. Prior work studied these mechanisms in isolation, often reaching pessimistic conclusions (e.g., "LLMs cannot self-correct reasoning" from Huang et al., 2023). This paper's framework reconciles those findings: self-correction does work, but only on the right difficulty tier. Search does help, but only with the right algorithm at the right budget. The conflicting prior results were an artifact of testing different methods on different (implicitly difficulty-biased) problem distributions.

This insight is more than taxonomic. It implies that future systems should not choose between revisions and search but should deploy both, switching between them per-prompt. The paper doesn't fully realize this vision (Section 8 acknowledges that PRM tree-search was not combined with revisions), but the framework provides the intellectual scaffolding for doing so.

Innovation 3: Empirical Evidence That Test-Time Compute Can Substitute for Pretraining — With Sharp Boundaries

The FLOPs-matched comparison in Section 7 is, to the authors' knowledge, the first to demonstrate in a realistic setting (no ground-truth access at inference) that a smaller model with additional test-time compute can outperform a ~14× larger model on problems within its capability range. This is significant not as a method but as an empirical finding with direct implications for how compute budgets should be allocated in production systems.

What distinguishes this from prior work on training-inference tradeoffs (Jones, 2021; Villalobos and Atkinson, 2023) is the specificity of the finding. The paper doesn't claim a universal substitution — it precisely characterizes where the substitution works (easy-to-medium problems, low R regimes) and where it fails (hard problems, high R regimes). The failure case is equally informative: on the hardest problems (bin 5), test-time compute provides essentially zero benefit regardless of budget, meaning that some capabilities can only be acquired through pretraining, not recovered at inference time. This establishes a clear boundary condition: test-time compute amplifies existing capability but does not create it from nothing.

The dependence on R = D_inference / D_pretrain adds practical nuance that prior analyses missed. For self-improvement pipelines where R ≪ 1, the case for test-time compute is strong. For high-throughput production deployments where R ≫ 1, the case weakens because the per-query inference cost of the larger model dominates the budget anyway. This is an incremental but practically important refinement of the training-inference tradeoff picture.

Innovation 4: Verifier Over-Optimization as a First-Class Phenomenon in Test-Time Scaling

While reward hacking / over-optimization is well-documented in the RLHF literature, this paper provides some of the first clear evidence that the same phenomenon governs test-time search scaling and is the primary bottleneck preventing unbounded improvements from additional compute. The evidence is concrete: beam search degrades easy-problem performance at high budgets (Figure 3, right); lookahead search — the most powerful optimizer — paradoxically performs worst overall (Figure 3, left); and qualitative examples in Appendix M show search producing degenerate outputs (repetitive low-information steps, overly short solutions) that score highly under the PRM.

This finding is significant because it shifts the narrative around test-time compute from "more is better" to "more is better only up to the verifier's reliability frontier." It explains why prior work found negative results for sophisticated search methods: those studies likely pushed past the over-optimization threshold. It also implies that improving verifier robustness is the key bottleneck for further scaling test-time compute, not improving search algorithms. The paper's compute-optimal policy can be understood partly as a way to stay below the over-optimization threshold per difficulty level — using weaker optimization (best-of-N) where the verifier is reliable (easy problems) and stronger optimization (beam search) only where the verifier signal has more room to provide genuine guidance (medium problems).

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The primary experiments use synthetic graphs generated via the Erdős-Rényi model with |V| = 100 nodes and edge probability 0.15, producing a moderately sparse directed graph. The set of all reachable source-target pairs is split such that approximately 20% of pairs form the training set D_Train and the remaining 80% form the test set D_Test. For SFT, K = 10 paths are sampled per reachable training pair via random walk, yielding the fixed dataset D_SFT. For real-world validation, the Blocksworld domain (Valmeekam et al., 2023a) with four blocks produces a graph with 73 nodes corresponding to distinct block configurations, with edges representing valid single-move transitions; all 73×72 directed node pairs are used for SFT training with 50,000 total sampled paths.

  • Base model(s). All experiments use a one-layer, single-head Transformer with embedding size d = 120. This minimal architecture is chosen because it can implement the handcrafted path-planning algorithm from Wang et al. (2024b) — predicting the next node using both adjacency and reachability information — making it sufficient to study fundamental learning dynamics without confounding factors from deeper architectures. The same architecture is trained from scratch for SFT (producing the "base model" used as initialization for all RL methods), then further trained with PG or Q-learning objectives. For the two-layer validation in Appendix G.1, a two-layer one-head Transformer is used to confirm attention patterns generalize.

  • Metrics. Four primary metrics are reported: (1) Training accuracy under temperature sampling (Temp=1) — the fraction of source-target pairs in D_Train for which the model, when sampling from its full output distribution, generates a complete valid path ending at the target; (2) Training accuracy under greedy decoding — the same but always selecting the highest-probability next token, used primarily in the appendices; (3) Test accuracy under greedy decoding — generalization accuracy on unseen pairs in D_Test; (4) Output diversity — the average number of distinct correct paths produced over 100 independent sampling trials (temperature=1) for the same source-target pair, measuring how many alternative valid solutions the model can generate beyond a single memorized path. For Q-learning, adjacency matrix recovery quality is also qualitatively assessed through logit heatmaps (Figure 4).

  • Baselines. The paper compares four training paradigms: (1) Supervised Fine-Tuning (SFT) — standard next-token cross-entropy on the fixed dataset D_SFT; (2) Continual SFT — further training the SFT model on the same D_SFT for additional steps, testing whether simply training longer helps (it doesn't); (3) Policy Gradient (PG), a vanilla implementation shown in Appendix F to be equivalent to unclipped PPO, with KL regularization strength λ swept across {0, 0.0001, 0.001, 0.01}; (4) Q-learning with two reward variants — outcome reward and process reward — each in on-policy and off-policy configurations. PG and Q-learning are initialized from the same SFT base model, ensuring fair comparison.

  • Generation budget / compute accounting. The paper does not use FLOP-based accounting. Instead, training progress is measured in training steps (gradient updates). Each step processes one batch of sequences generated either from the fixed SFT dataset or from the model's own on-policy rollouts. For Q-learning off-policy, sequences are generated by the frozen base model rather than the current policy. All methods are trained until convergence (up to ~100,000 steps for PG, ~200,000–300,000 steps for Q-learning). No per-method compute adjustments are applied — the comparison focuses on asymptotic behavior rather than sample efficiency.

  • Cross-validation / statistical protocol. No cross-validation is used. The paper's primary validation strategy is theoretical: it proves convergence properties mathematically and then runs single training trajectories to confirm that empirical behavior matches the theoretical predictions. The Erdős-Rényi experiments use a random graph seed with fixed |V|=100 and edge probability 0.15, with the 20/80 train-test split on reachable pairs. Blocksworld uses a fully deterministic 73-node graph derived from all valid block configurations. The reproducibility of the qualitative phenomena (diversity collapse, Q-learning structural convergence, outcome-reward failure) across synthetic and real graphs serves as the main robustness check rather than statistical averaging over random seeds.

Main Quantitative Results

SFT Memorizes Co-Occurrence Frequencies

The paper's first empirical demonstration is that SFT learns to match co-occurrence statistics rather than true graph structure. Figure 1 compares the actual edge frequencies in the Blocksworld SFT training data (panel a, where all edges appear at least once) against the adjacency weights learned by SFT (panel b), PG (panel c), and Q-learning (panel d). In the SFT model, many edges that appear in the training data receive low learned weights, particularly low-frequency edges. The figure uses a heatmap where brighter indicates higher weight. The SFT panel shows a noticeably sparser and less accurate recovery of the true adjacency matrix compared to the RL-trained models. This is consistent with Theorem 3.1: SFT's predicted next-node distribution converges to N_{i,j,k} / ∑_{k'} N_{i,j,k'} — the empirical frequency of observing node k after node j when targeting i. Low-frequency edges, even if valid, receive low probability. Both RL methods (panels c, d) capture the adjacency structure more faithfully, with Q-learning (panel d) nearly recovering the complete adjacency matrix.

Figure 2a illustrates the downstream impact: as training progresses, Continual SFT's test accuracy steadily decreases from its initial value. In contrast, all PG variants improve over their starting point. The continual SFT model overfits to the co-occurrence patterns in D_SFT without acquiring the underlying graph connectivity needed for generalization.

Policy Gradient Outperforms SFT Through Exploration-Driven Data Augmentation

Figure 2a shows that while continual SFT degrades test accuracy, PG without KL regularization (λ=0) improves from approximately 0.88 to roughly 0.91 before declining, while PG with λ=0.001 continues improving to approximately 0.925. This validates Takeaway 2: PG's advantage comes from exploration-driven data augmentation. By generating its own data on-policy and filtering for correct paths (Theorem 4.1 shows each PG step is equivalent to SFT on D_{RL,t} ∩ P), PG discovers novel valid trajectories that were absent from D_SFT, expanding its effective training set beyond the co-occurrence limitations that bound SFT.

Figure 2b confirms that PG without KL achieves and maintains 100% training accuracy under temperature sampling — the model perfectly solves all training pairs. PG with KL regularization (λ=0.01, λ=0.001) converges to slightly lower training accuracy (approximately 0.99 and 0.985, respectively), reflecting the accuracy cost of diversity preservation.

Policy Gradient Without KL Regularization Exhibits Diversity Collapse

Figure 2c is the paper's most striking empirical result for PG. The output diversity of PG without KL (λ=0) — measured as the average number of distinct correct paths per source-target pair over 100 sampling trials — steadily declines throughout training, dropping from approximately 5.5 at the start of RL training to near 1 (a single unique path) by step 100,000. This decline continues even after the model achieves 100% training accuracy (which occurs much earlier, around step 20,000–30,000 based on Figure 2b). The model does not just converge to correct behavior; it converges to minimally diverse correct behavior, outputting only one path per pair.

This empirically validates Theorem 4.3, which proves that each PG update (without KL) increases the KL divergence between the model's output distribution and the uniform distribution over valid actions. The convexity of the log-normalizer in the softmax ensures that the stochastic gradient noise, combined with the on-policy data generation, progressively concentrates probability mass onto fewer and fewer valid alternatives.

In contrast, PG with KL regularization preserves substantially higher diversity. With λ=0.01, diversity remains above 4 at convergence; with λ=0.001, it settles around 2.5–3. This directly demonstrates Theorem 4.4's prediction: the KL term explicitly penalizes deviation from the base model distribution, where the base model (being SFT-trained) has non-trivial diversity. The regularization acts as a diversity-preserving mechanism, but Figure 2d shows the trade-off: as λ increases (moving right on the x-axis in the left panel), training accuracy decreases (from ~1.0 at λ=0 to ~0.985 at λ=0.01) while diversity increases (from ~1 to ~5). The right panel shows the Pareto frontier in test accuracy-diversity space, confirming that intermediate λ values (around λ=0.001) achieve the best balance between diversity and accuracy.

KL Regularization Exhibits a Double-Edged Effect on Generalization

Figure 7 in Appendix G.2 presents a more nuanced view using a data split that introduces new training pairs unseen during SFT (D_Test2Train) and SFT pairs excluded from RL training (D_Train2Test). PG without KL (λ=0) achieves the highest accuracy on D_Test2Train (new pairs, approaching 1.0) — demonstrating that unrestricted PG updates most effectively learn novel planning problems. However, it also exhibits progressive forgetting on D_Train2Test (SFT pairs not used in RL, accuracy declining from initial values), showing that unrestricted PG overfits to its RL training distribution. PG with λ=0.01 maintains near-perfect accuracy on D_Train2Test (little forgetting) but struggles to learn D_Test2Train (lower accuracy), because the strong KL constraint prevents the policy from shifting enough to acquire new valid paths. The intermediate value λ=0.0001 achieves the best overall balance. This pattern validates the "double-edged sword" characterization from the theoretical analysis: KL regularization preserves prior knowledge at the cost of limiting adaptation, making the optimal λ highly dependent on both the base model quality and the novelty of the RL training distribution.

Outcome-Reward Q-Learning Collapses to Trivial Solutions

Figure 3a (blue curves) shows Q-learning with outcome reward catastrophically fails: training accuracy collapses to near-zero and test accuracy similarly flatlines. This empirically confirms Theorem 5.1 — the outcome reward signal (δ_{u∈P} · δ_{u_{m+1}=u_target}) provides only sparse feedback (reward of 1 at the final step of correct paths, 0 elsewhere), and the learned Q-values collapse to a constant c_i for all non-target transitions, losing all structural information about the graph. The model cannot differentiate valid from invalid transitions under this reward design.

Process-Reward Q-Learning Converges Correctly and Preserves Diversity

Figure 3a (orange curves) shows Q-learning with process reward achieves training accuracy comparable to PG and significantly higher test accuracy, rising to approximately 0.95 by step 200,000 while PG (λ=0.001) plateaus around 0.90. This validates Theorem 5.2: the process reward (δ_{u_{m+1}=u_target} - δ_{(u_m, u_{m+1})∉E}) provides per-step structural signals that drive the logits to converge to the correct structural values — 1 for edges that are both adjacent AND reachable, 0 when exactly one condition holds, -1 for neither.

Figure 4 visualizes this convergence through logit heatmaps for a Q-learning model with process rewards and attention fixed on the target node. Each row i (target nodes 0–20) shows the normalized logits for candidate next nodes when the current node is 0. Green frames highlight valid next nodes (both children of node 0 AND ancestors of target i). Over training epochs (10,000 → 30,000 → 100,000 → 300,000), the logits for valid nodes (green-framed cells) consistently brighten (increase toward white, indicating higher values), while invalid nodes remain dark. By epoch 300,000, all valid next nodes within each row have converged to similarly bright values, confirming that the model learns to assign equal high logits to all valid transitions — the diversity-preserving property that PG lacks.

Figure 3b directly compares the diversity-accuracy Pareto frontiers across methods on both training and test sets. Q-learning with process rewards achieves high accuracy while maintaining diversity comparable to or exceeding PG with KL regularization. On the training set, Q-learning reaches diversity values of approximately 5 at accuracy ~0.99, while PG (λ=0) achieves accuracy 1.0 but diversity near 1, and PG (λ=0.01) achieves diversity ~5 but lower accuracy (~0.985). On the test set, Q-learning dominates the frontier, achieving higher accuracy at any given diversity level compared to all PG variants and continual SFT.

Off-Policy Q-Learning Matches On-Policy Performance

Figure 3a (green curve) shows off-policy Q-learning with process reward — where rollouts are generated by the frozen base model rather than the current policy — achieves training and test accuracy nearly identical to on-policy Q-learning. Both converge to approximately 0.95 test accuracy by step 200,000, with the off-policy variant showing slightly slower initial progress but matching final performance. This is significant because modern RLHF frameworks like VeRL (Sheng et al., 2024) effectively implement off-policy updates when using quantized models or large batch sizes where the policy can change between rollout generation and gradient computation. Theorem 5.2's convergence proof does not require on-policy data — only persistent exploration — which explains why off-policy training works, unlike PG where the on-policy nature is baked into the gradient structure (Theorem 4.3's proof explicitly relies on the multinomial sampling distribution matching the current policy).

Learned Adjacency and Attention Patterns Confirm Theoretical Mechanisms

Figure 5 (Appendix G.1) validates Assumption 3.1 — that the transformer learns to operate as a function of target and current nodes — by visualizing attention weight evolution. In SFT (panel a), attention on the target node (token position 1) quickly peaks early in training but then gradually decreases while remaining dominant, a pattern the authors hypothesize reflects overfitting to D_SFT leading to auxiliary prediction strategies. In PG (panel b), attention on the target node increases through training. In Q-learning (panel c), the final attention weight on the target node exceeds 95%, making it the closest realization of the f(u_target, u_m) abstraction assumed in the theoretical analysis. Figure 6 confirms the same pattern generalizes to a two-layer, one-head Transformer, with both layers predominantly attending to the target and current nodes.

Ablation Studies and Robustness Checks

KL regularization strength sweep (Figure 2d): The left panel shows that as λ increases from 10^{-4} to 10^{-1} (log scale), training accuracy under temperature sampling decreases from ~1.0 to ~0.975, while diversity increases from ~1 to ~5.5. The relationship is approximately monotonic in the log-domain, with the steepest trade-off occurring between λ = 10^{-3} and λ = 3×10^{-3}. The right panel maps this to test accuracy: PG (λ=0) reaches the highest test accuracy early in training (Figure 2a) but then degrades as diversity collapses, while PG (λ=0.001) achieves more stable test accuracy that continues improving through training.

Continual SFT vs. PG (Figures 2a-c): This ablation isolates whether the PG benefit comes from exploration-driven data augmentation or simply from more training. Continual SFT (more gradient steps on D_SFT) degrades test accuracy from start (Figure 2a, dark blue curve declining from ~0.90 to ~0.80), while simultaneously showing high but not perfect training accuracy under temperature sampling (Figure 2b, ~0.97, lower than PG's 1.0) and declining diversity (Figure 2c, dropping from ~5.5 to ~4.0). This confirms that PG's improvement is not an artifact of additional optimization steps — it requires the exploration data generation that SFT on a fixed dataset cannot provide.

Outcome vs. process reward in Q-learning (Figure 3a): The outcome-reward variant's complete failure (accuracy near zero for both train and test) versus the process-reward variant's strong performance (test accuracy ~0.95) is the most dramatic ablation in the paper. This single comparison establishes that reward design is not a minor implementation detail for Q-learning in LLM planning — it determines whether the method works at all. The process reward's decomposition into per-step adjacency and target signals is necessary for the model to learn structural graph information.

On-policy vs. off-policy Q-learning (Figure 3a): Off-policy Q-learning with process rewards matches on-policy Q-learning final performance on both training and test accuracy, with the curves largely overlapping after ~150,000 steps. This demonstrates that Q-learning's convergence does not depend on the data coming from the current policy — a property that PG fundamentally lacks due to its on-policy gradient structure (Theorem 4.3 relies on the sampling distribution matching the policy).

Data split analysis (Figures 7 and 8, Appendix G.2): Using a split where RL training introduces new pairs (D_Test2Train) while excluding some SFT pairs (D_Train2Test), the paper provides a more granular decomposition of generalization and forgetting. PG without KL (λ=0) achieves the highest accuracy on new pairs (~1.0 on D_Test2Train, Figure 7) but exhibits substantial forgetting on excluded SFT pairs (~0.925 on D_Train2Test, declining). PG with λ=0.01 mostly avoids forgetting (~0.975 on D_Train2Test) but struggles with new pairs (~0.90 on D_Test2Train). Q-learning with process reward (Figure 8) converges more slowly but achieves strong performance across all splits, with notably better preservation of D_Train2Test accuracy than unregularized PG, though the absolute convergence speed is slower, which the authors attribute to the initial model performing poorly on new pairs, generating more failure cases.

Two-layer Transformer attention patterns (Figure 6, Appendix G.1): Extending the attention analysis to a two-layer, one-head Transformer shows that both layers predominantly attend to the target and current nodes across all three training paradigms (SFT, PG, Q-learning), confirming that Assumption 3.1's reduction to f(u_target, u_m) is not an artifact of the single-layer architecture. This strengthens the theoretical analysis's applicability to deeper models, though the paper does not report accuracy or diversity metrics for the two-layer case.

Blocksworld real-world validation (Figure 1, Appendix G.3): The adjacency matrix recovery experiment is replicated on the Blocksworld graph (73 nodes, deterministic transitions), showing the same qualitative pattern: SFT learns a degraded adjacency (panel b), PG improves it (panel c), and Q-learning nearly recovers the complete adjacency (panel d). Notably, Q-learning's adjacency recovery is near-perfect on Blocksworld, consistent with Theorem 5.3's prediction that the weights decompose as W_M[j,k] = A[j,k] - 1 + c_k and W_V[i,k] = R[i,k] - c_k, with the gauge constant c_k canceling in the logit sum to produce correct structural values.

Token probability distribution shift (implicit in diversity metric): The diversity metric itself serves as an ablation on whether model improvement comes from learning to produce more correct paths or from learning to suppress incorrect ones. For PG without KL, diversity declines to near 1 at step 100,000 (Figure 2c) while training accuracy is 1.0 (Figure 2b), indicating the model is not just filtering out errors but actively eliminating alternative correct solutions. For Q-learning with process reward, the logit heatmaps (Figure 4) show valid next nodes receiving equal high logits, corresponding to maintained diversity.

Critical Assessment

Does the Theoretical Framework Actually Predict the Empirical Results, or Are the Empirical Results Illustrations of the Theory?

The paper presents the experiments as validations that the theoretically derived behaviors manifest in practice, not as independent discoveries. This is a legitimate scientific approach, but it means the experiments do not independently test the theory — they cannot falsify it because the experimental conditions (one-layer Transformer, Erdős-Rényi graphs, the specific path-finding abstraction) are explicitly designed to satisfy the theory's assumptions. The strongest claim the experiments support is: "under conditions that approximate our theoretical assumptions, the predicted phenomena occur." Whether these phenomena occur under substantially different conditions — deeper models, natural language inputs, larger graphs, different planning structures — is not tested.

The Blocksworld experiment (Figure 1) partially addresses this by using a non-synthetic graph, but it only examines adjacency matrix recovery, not the full suite of metrics (training accuracy, diversity curves, test generalization) shown for the Erdős-Rényi setting. The paper does not report what accuracy PG or Q-learning achieves on Blocksworld planning queries, nor whether diversity collapse occurs on the 73-node Blocksworld graph. This limits how strongly the "real-world" validation supports the broader claims.

The Single Architecture and Graph Family Are Genuine Limitations

All quantitative results come from a one-layer, single-head Transformer on graphs with ≤100 nodes. The paper argues compellingly that this architecture is the minimal model needed to implement the handcrafted planning algorithm, making it ideal for isolating fundamental learning dynamics. However, several concerns arise:

  • Representational capacity: A one-layer Transformer may be capacity-limited in ways that mask or amplify certain behaviors. Diversity collapse in PG might be less severe in larger models where multiple parameter configurations can achieve 100% accuracy without extreme output concentration. Conversely, Q-learning's diversity preservation might be harder to achieve in deeper models if the gradient dynamics become more complex.

  • Graph size scaling: The 100-node Erdős-Rényi graph has an average out-degree of 15 and a relatively small diameter. Whether Q-learning's linear convergence (Theorem 5.2) remains practical on graphs with thousands of nodes — where the effective contraction factor along long paths becomes extremely slow — is untested. The paper acknowledges that "longer paths converge slower" in the rate analysis, but does not explore at what graph scale this becomes a practical impediment.

  • Graph structure: Erdős-Rényi graphs are maximally unstructured — every edge is independent. Real planning domains often have highly structured transition graphs (e.g., Blocksworld has specific symmetries, tool-use graphs have hierarchical dependencies). The paper does not test whether SFT's co-occurrence limitation is more or less severe on structured graphs, or whether PG's exploration mechanism differentially benefits from graph regularity.

The Experimental Protocol Lacks Statistical Replicates

The paper runs single training trajectories for each method and reports these trajectories as evidence. There are no error bars, no multiple random seeds, and no quantification of variance. For the key result — diversity collapse under PG (Figure 2c) — the decline from ~5.5 to ~1 unique paths is dramatic enough that variance is unlikely to change the qualitative conclusion. However, for the comparison between PG (λ=0.001) and Q-learning (Figure 3a), where the test accuracy gap is ~5 percentage points at convergence, the absence of replicates makes it unclear whether this difference is statistically reliable or within run-to-run variation. The paper also does not report whether the test accuracy curves in Figure 3a are evaluated on a single fixed test split or averaged over multiple splits, nor whether the 20/80 train-test split was randomly sampled once or multiple times.

The Q-Learning Implementation Is Oversimplified Relative to Practice

The paper's Q-learning implementation — directly using model logits as Q-values and training with simple TD error — is far simpler than the Q-learning variants that would be needed for practical LLM training. In particular:

  • Discrete action space: The paper operates in a small discrete vocabulary (100 node tokens). Practical LLM planning involves a combinatorial action space over natural language, where Q-learning's max_k operation over the vocabulary would be both computationally expensive and ill-defined (since most tokens do not correspond to valid planning actions).

  • No function approximation challenges: The one-layer Transformer's parameters directly encode f(i, j)[k] for each triple, making the Q-function tabular in practice. The deadly triad (function approximation, bootstrapping, off-policy learning) that makes Q-learning unstable in deep RL is largely absent here. The paper's finding that off-policy Q-learning works well does not necessarily transfer to settings where the Q-function is approximated by a deep network.

  • The process reward requires graph knowledge: The process reward δ_{u_{m+1}=u_target} - δ_{(u_m, u_{m+1})∉E} requires knowing both the target and the adjacency check at each step. In real planning, process rewards must be derived from environment feedback or learned verifiers, which introduces their own errors and biases. The paper does not discuss how process rewards would be obtained in settings where the ground-truth adjacency matrix is unknown.

Claims That Hold Conditionally vs. Claims That Hold Generally

Takeaway 2 (PG outperforms SFT through exploration): Strongly supported for the studied setting. The continual SFT degradation in Figure 2a combined with PG's improvement and the theoretical connection in Theorem 4.1 make a compelling case. The caveat is that "exploration" here means on-policy sampling from the model's output distribution — a relatively weak form of exploration that succeeds because the SFT initialization provides a non-trivial starting point. For tasks where the SFT model has near-zero probability of generating any correct path, this exploration mechanism would fail, and PG would receive no positive reward signal to learn from.

Takeaway 3 (Diversity collapse without KL): Strongly supported for the one-layer Transformer on 100-node graphs. The empirical decline in Figure 2c is unambiguous. The theoretical proof (Theorem 4.3) relies on the on-policy update structure, which is architecture-agnostic — it depends only on the softmax parameterization and the gradient update, not on model capacity. This suggests the finding may generalize, but the paper does not test it on larger models.

Takeaway 4 (KL regularization preserves diversity at accuracy cost): Strongly supported with the important qualification that the optimal λ is task-dependent. Figure 2d shows the trade-off is real and tunable. Figure 7 (Appendix G.2) shows that the optimal λ depends on the relationship between SFT and RL training distributions, with λ=0.0001 best when RL introduces novel pairs. The paper does not provide guidance on how to select λ a priori.

Takeaway 5 (Q-value bias with outcome rewards, fixed by process rewards): Strongly supported within the studied abstraction. The outcome-reward Q-learning failure (Figure 3a, near-zero accuracy) is dramatic and matches Theorem 5.1's prediction. The caveat is that "outcome reward" in this setting is extremely sparse (a single reward at the end of a valid path), whereas practical LLM outcome rewards often include partial credit, shaped rewards, or are augmented with verifier outputs that provide richer signals — any of which might mitigate the Q-value bias.

Takeaway 6 (Q-learning preserves diversity and supports off-policy learning): Supported with implementation caveats. The logit heatmaps (Figure 4) convincingly show equal weighting of valid next nodes, and the off-policy curves in Figure 3a closely track on-policy performance. However, the claim that Q-learning "better maintains output diversity" than PG is comparing Q-learning with process rewards against PG without KL regularization — a comparison that stacks the deck, since PG with KL (λ=0.01) achieves comparable diversity (Figure 3b). A fairer framing is: Q-learning with process rewards achieves both high accuracy and high diversity without requiring KL regularization tuning, while PG requires careful λ selection to balance accuracy and diversity.

Missing Experiments That Would Strengthen the Paper

Several experiments are conspicuously absent:

  • Scaling model depth: All results use one-layer Transformers (with a brief two-layer attention check in Appendix G.1 that only examines attention patterns, not accuracy or diversity). Testing whether diversity collapse and Q-learning structural convergence persist in 2–4 layer models would significantly strengthen the generality claims.

  • Scaling graph size: The 100-node graph is tractable. Testing on graphs with 500 or 1000 nodes would reveal whether Q-learning's convergence rate remains practical (Theorem 5.2's rate depends on path length, and larger graphs have longer paths).

  • Multiple random seeds: Running 3–5 seeds and reporting means with error bars for Figures 2 and 3 would address the variance concern, particularly for the PG vs. Q-learning test accuracy comparison.

  • Natural language planning tasks: The paper could have tested whether the theoretical insights transfer to a setting where the model processes natural language state descriptions and generates natural language actions, rather than node tokens. This would require training a verifier to provide process rewards, making it a substantially more complex experiment — but it would directly address the "does this matter for real LLMs" question that the paper's framing invites.

  • Measuring adjacency matrix recovery on Erdős-Rényi graphs: Figure 1 shows adjacency recovery only for Blocksworld. Showing the equivalent heatmaps for the 100-node Erdős-Rényi graph used in the main experiments would allow direct correlation between adjacency learning quality and the planning accuracy curves in Figures 2 and 3. Without this, the reader cannot assess whether PG's diversity collapse corresponds to progressively worse adjacency encoding, or whether PG's adjacency matrix is actually good and the collapse only affects output diversity, not structural understanding.

  • Comparison against behavioral cloning on RL-generated data: Theorem 4.1 shows PG is equivalent to SFT on D_{RL,t} ∩ P. A natural ablation would be: collect all correct paths generated by exploration, then train with SFT on the aggregated dataset. This would isolate whether PG's on-policy nature provides benefits beyond data augmentation — for instance, whether the sequential refinement of the policy (rather than batch training on all discovered paths) matters for generalization. The paper does not run this experiment.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Makes the Headline 4× Efficiency Gain Inapplicable to Deployment

The entire compute-optimal framework depends on estimating each prompt's difficulty before allocating the inference budget. The paper's method for doing so — generating 2,048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted) — consumes an amount of computation comparable to or exceeding the largest test-time budgets studied. The authors acknowledge this explicitly in Section 3.2: "estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity."

Consequence: The reported 4× efficiency gains (e.g., 16 generations matching 64 for search in Figure 4, 64 generations matching 256 for revisions in Figure 8) are computed after difficulty is known, without amortizing the cost of estimating it. In any realistic deployment, the total cost would be difficulty estimation plus strategy execution. Since estimation requires 2,048 samples — far more than any budget studied — the headline efficiency numbers are unachievable in practice without a cheaper difficulty estimation method. The compute-optimal policy might actually be more expensive than standard best-of-N once difficulty estimation is included, though the paper does not measure this.

What evidence exists: The paper provides no ablation or analysis of what happens when difficulty estimation cost is included in the budget. Section 3.2 flags this as future work but does not attempt even a back-of-the-envelope calculation. The difficulty estimation protocol itself is described in Section 3.2 but its cost (2,048 × cost of one generation) is never converted into the same generation-budget units used in the scaling curves, making direct comparison impossible.

Mitigation status: Not addressed in this paper. The authors suggest future work on "pretraining or finetuning models to directly predict difficulty of a question" (Section 8). This is a non-trivial research problem — such a model would need to predict the base LLM's pass@1 rate from the question text alone, which may be as hard as solving the question. An adaptive approach (starting with a few samples, estimating difficulty from the PRM's score distribution, then allocating the remaining budget) is mentioned as a possibility but not explored.

All Results Come from a Single Benchmark (MATH) with a Single Model Family (PaLM 2-S*), with No Evidence of Generalization

Every quantitative result in the paper — the 4× efficiency gains, the difficulty-dependent strategy optimality, the beam search over-optimization thresholds, the FLOPs-matched comparison against a larger model — is measured on exactly 500 MATH test questions using PaLM 2-S* (Codey) as the base model. The authors argue this model is "representative of the capabilities of many contemporary LLMs" (Section 4), but this claim is unverified. MATH consists of high-school competition mathematics problems requiring symbolic reasoning, a domain with specific properties (deterministic ground-truth answers, clear correctness criteria, multi-step deductive structure) that may not generalize.

Consequence: Multiple aspects of the findings could be model-specific or domain-specific. The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution and error patterns — a model with different calibration or different types of mistakes would produce different difficulty-dependent scaling curves, potentially shifting the optimal strategy thresholds. The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capability, which varies substantially across model families. The finding that beam search helps on medium problems but hurts on easy ones is a function of the PRM's reliability relative to the base model's capability — with a better or worse PRM, the difficulty boundaries would shift, and with a different base model, the difficulty distribution itself would change. Practitioners using a different model (GPT-4, Claude, Llama) or a different domain (code generation, scientific QA, everyday reasoning) cannot confidently apply the paper's specific strategy recommendations.

What evidence exists: None beyond MATH and PaLM 2-S*. The paper does not report results on any other benchmark (GSM8K, HumanEval, MMLU, Big-Bench Hard) or with any other model family, even in an appendix. The test set of 500 questions, split into five difficulty quintiles of ~100 each, then further split by two-fold cross-validation, means the compute-optimal policy is selected based on ~50 questions per fold per bin — a sample size small enough that the selected strategies may not be robust even within MATH. The paper does not report confidence intervals on the compute-optimal scaling curves.

Mitigation status: The authors acknowledge the limitation implicitly by restricting claims to "the MATH benchmark using PaLM 2-S* models" in the abstract and introduction, but do not discuss domain or model generalization as an explicit limitation. No future work on cross-model or cross-domain validation is suggested.

The 14× Larger Model Baseline Is Not Compute-Optimally Trained, Weakening the Pretraining vs. Inference Comparison

The FLOPs-matched comparison in Section 7 scales model parameters by ~14× while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023). The authors explicitly acknowledge this departs from compute-optimal pretraining where both data and parameters scale: "We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." Additionally, the 14× larger model uses only greedy decoding with no test-time augmentation of its own — no majority voting, no best-of-N, no search.

Consequence: The comparison systematically favors test-time compute. A Chinchilla-optimal model trained with 14× more total FLOPs (scaling both parameters and data according to Hoffmann et al., 2022) would likely outperform a parameter-only-scaled model, making the pretraining baseline stronger. Furthermore, giving the larger model even a modest test-time compute budget (e.g., best-of-8 with majority voting) would create a much more realistic baseline — one that reflects how large models are actually deployed — and could substantially narrow or reverse the reported advantages of test-time compute over pretraining. The headline result that "a smaller model augmented with compute-optimal test-time strategies can outperform a ~14× larger pretrained model" (Section 1) should be understood as "outperforms a specific, potentially suboptimal larger model configuration," not as a general statement about the pretraining-inference tradeoff.

What evidence exists: The paper provides no comparison against a compute-optimally trained larger model, nor against a larger model with any test-time augmentation. Section 7 describes the FLOP-accounting assumptions in detail, making the limitation transparent, but the bar charts in Figure 1 and the scaling curves in Figure 9 present the results without caveats about the baseline's suboptimality. The specific numbers — e.g., +27.8% relative improvement on easy questions at R ≪ 1 for revisions (Figure 1, top-right) — are contingent on this baseline choice and should not be cited without the caveat.

Mitigation status: Explicitly acknowledged in Section 7 but left entirely to future work. The authors frame the current comparison as "representative of a canonical approach," which is a reasonable starting point, but the paper would be stronger with even a single alternative baseline (e.g., best-of-8 on the larger model) or a sensitivity analysis.

Hard Problems (Difficulty Bin 5) Show Near-Zero Improvement Regardless of Method or Budget, Establishing a Hard Ceiling

Across every method studied — best-of-N, beam search, lookahead search, sequential revisions, parallel revisions, and their compute-optimal combinations — the hardest questions (difficulty bin 5, where the base model's pass@1 is near zero) show essentially no improvement with any amount of test-time compute. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all search methods and all budgets up to 256 generations. In Figure 7 (right), bin 5 shows ~2–3% accuracy for revisions irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% across all budgets. The paper is transparent about this in the Section 7 takeaway: "test-time compute amplifies existing capability but does not create it from nothing."

Consequence: This is the most fundamental limitation of the approach: test-time compute cannot solve problems that are genuinely outside the base model's capability range. If the base model's pass@1 is near zero on a problem class — meaning it almost never generates a correct solution even once in thousands of attempts — no search algorithm, revision strategy, or compute-optimal allocation can help, because there is no correct answer in the proposal distribution to find or refine. This means the approach offers no path forward for the hardest reasoning problems or for out-of-distribution generalization where the model has not acquired the relevant knowledge during pretraining. For organizations whose problem distribution skews toward hard problems, investing in larger pretraining is the only viable strategy — a conclusion the paper's own data supports.

What evidence exists: The flat bin 5 curves appear consistently in Figures 3 (right), 7 (right), and 9. The paper discusses this limitation explicitly in Section 7 and in the executive summary, making it one of the most clearly communicated limitations. The evidence is strong because it holds across both search and revision methods and across all budget levels studied.

Mitigation status: This is not a limitation that can be "mitigated" within the test-time compute framework — it is a fundamental boundary condition. The paper's contribution is precisely to characterize this boundary, showing where test-time compute works (easy-to-medium problems) and where it does not (hard problems). The implication is that future systems should use difficulty estimation to route hard problems to larger models or human intervention, but the paper does not develop such a routing system.

Revisions and Search Are Studied Independently, Never Combined, Leaving the Full Potential of the Framework Unexplored

The paper studies two complementary mechanisms — PRM-guided search (modifying how outputs are selected via beam search, lookahead search, best-of-N weighted) and iterative revisions (modifying what the model generates by conditioning on previous attempts) — but never combines them. Section 8 explicitly acknowledges: "we did not experiment with PRM tree-search techniques in combination with revisions."

Consequence: The paper's results represent a lower bound on what the framework can achieve. Search and revisions have complementary strengths: revisions improve the proposal distribution (generating better candidates, especially on easy problems where local refinement suffices), while PRM search improves candidate selection (finding the best among generated options, especially on medium problems where exploration matters). Applying beam search to revision model outputs — using the PRM to score revision steps and prune unpromising revision chains — could yield gains beyond either method alone. Alternatively, using the PRM to decide which revisions to pursue rather than blindly generating long chains could mitigate the 38% correct-to-incorrect reversion rate (Section 6.1). The difficulty-dependent optimal strategies might also shift: perhaps combined search-and-revisions would outperform pure search on medium problems, or pure revisions on easy problems, at the same budget.

What evidence exists: None. The paper provides no ablation, pilot experiment, or even qualitative analysis of combined search-and-revisions. Section 8 lists this as future work. The architecture for combining them is discussed in principle — the revision model would serve as the proposal distribution within beam search — but no implementation or results are presented.

Mitigation status: Acknowledged as a key future direction in Section 8. The paper's framework (proposal distribution vs. verifier decomposition from Section 2) provides the conceptual scaffolding for such a combination, and the difficulty-dependent analyses (Figures 3 right, 7 right) suggest where each component would contribute most. But the actual integration remains unimplemented.

Latency and Wall-Clock Time Are Ignored, but Sequential Strategies Are Inherently Serial

The paper measures compute in "generations" (number of complete solutions sampled), which is a reasonable proxy for total FLOPs but completely ignores latency — the wall-clock time required to produce an answer. Sequential revision strategies (the revision model generating a chain of revisions where each depends on the previous one) are inherently serial and cannot be parallelized. A strategy that allocates 128 generations as 64 sequential revisions × 2 parallel chains takes roughly 64× longer wall-clock time than one that runs 128 independent parallel samples simultaneously, even though both use the same total generation budget.

Consequence: The compute-optimal policy tends to favor sequential revisions on easy problems (Figure 7, left: at low budgets, fully sequential is optimal) and a balanced sequential-to-parallel ratio on harder problems. For any latency-sensitive application — interactive assistants, real-time decision-making, live tutoring systems — the sequential-heavy strategies may be impractical regardless of their FLOPs-efficiency advantages. A user waiting for a response cannot tolerate 64 sequential generation steps even if the total FLOPs match a parallel strategy that completes in 1 step. The paper's efficiency claims (4× fewer generations) would not translate to 4× lower latency in deployment; in fact, the latency of the compute-optimal policy might be worse than the baseline it outperforms on generation count.

What evidence exists: None. The paper never discusses latency, throughput, or wall-clock time. The entire analysis is in generation-budget space, which is appropriate for a scaling laws analysis (the standard in the pretraining scaling literature that the paper emulates) but incomplete for practical deployment considerations. The revision model's sequential dependency is visible in the architecture description (Section 6) but its latency implications are not quantified or discussed.

Mitigation status: Not addressed and not acknowledged as a limitation. The paper's framing as an "inference-time analog of the Chinchilla scaling laws" implicitly adopts the pretraining scaling literature's convention of measuring FLOPs rather than wall-clock time, but pretraining is a one-time cost where latency is irrelevant, while inference is user-facing where latency often matters more than total FLOPs. This mismatch goes unremarked.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper provides the first rigorous gradient-dynamics analysis of reinforcement learning for language model planning, shifting the conversation from empirical trial-and-error ("RL works better than SFT, but we don't know why") to a principled understanding of mechanisms, failure modes, and design constraints. Before this work, the field's understanding of RL for LLM reasoning was fragmented across disconnected empirical observations — Chu et al. (2025) noted SFT memorizes while RL generalizes, Cui et al. (2025) documented diversity collapse, Setlur et al. (2025) proved verification-free approaches are suboptimal — but none of these were connected to a unified theoretical framework that explained why these phenomena occur and how they relate. This paper provides that framework.

The contribution is best characterized as a theoretical reframing rather than a paradigm shift. The paper does not introduce a new algorithm that practitioners will adopt tomorrow; instead, it provides the analytical lens through which existing algorithms can be understood, compared, and improved. Specifically:

Reconciling conflicting practical observations. The paper's analysis explains why KL regularization sometimes helps and sometimes hurts — a puzzle that has confused practitioners. Theorem 4.4 shows that KL regularization acts as a diversity-preserving mechanism by anchoring the policy to the base model's output distribution, with the stable point satisfying $q(i, j)[k] \propto q_{\text{base}}(i, j)[k] \cdot \exp(p(i, j)[k] / \lambda)$. When the base model already assigns reasonable probability to correct transitions (easy problems, capable models), this anchoring preserves useful diversity and prevents overfitting — explaining the positive results. When the base model's prior is poor (hard problems, weak models), the KL term prevents the policy from shifting enough to learn new valid paths — explaining the negative results. This single equation resolves the apparent contradiction and provides a principled basis for setting $\lambda$.

Establishing diversity collapse as a structural property, not an implementation bug. Theorem 4.3 proves that PG without KL regularization must exhibit diversity collapse — it is baked into the gradient dynamics through the convexity of the log-normalizer in the softmax. Each on-policy update, in expectation, increases the KL divergence between the model's output distribution and the uniform distribution over valid actions. This means diversity collapse is not something that can be fixed by better hyperparameters, larger models, or different optimization tricks within the PG framework — it requires an explicit mechanism (KL regularization, or switching to a different paradigm) to counteract. This shifts the research narrative from "how do we prevent diversity collapse in PG?" to "given that PG collapses diversity, when is that acceptable, and what alternatives exist?"

Introducing Q-learning as a principled alternative for LLM planning. Perhaps the paper's most forward-looking contribution is the theoretical argument that Q-learning with process rewards naturally avoids PG's core limitations. Theorem 5.2 shows that process-reward Q-learning converges to a stable point where all valid next nodes receive equal high logits (logit value 1), all partially valid nodes receive intermediate logits (logit value 0), and all invalid nodes receive low logits (logit value -1). This means maximal output diversity among correct transitions is built into the convergence point, not added as a regularizer. Theorem 5.3 shows this property is robust to the specific parameterization — even in a linear Transformer where the weights decompose as $W_M[j,k] = A[j,k] - 1 + c_k$ and $W_V[i,k] = R[i,k] - c_k$, the gauge constant $c_k$ cancels in the logit sum, preserving the structural encoding. The empirical validation (Figure 3a, Figure 4) confirms that Q-learning with process rewards achieves higher test accuracy than PG while maintaining output diversity comparable to PG with heavy KL regularization — but without requiring the practitioner to tune $\lambda$.

Redirecting research attention toward reward design and verifier quality. The catastrophic failure of outcome-reward Q-learning (Theorem 5.1, validated in Figure 3a where accuracy collapses to near-zero) demonstrates that reward design is not a minor implementation detail — it determines whether Q-learning works at all. The outcome reward ($\delta_{u \in \mathcal{P}} \delta_{u_{m+1}=u_{\text{target}}}$) collapses all non-target logits to a constant $c_i$, eliminating all structural information about the graph. The process reward ($\delta_{u_{m+1}=u_{\text{target}}} - \delta_{(u_m, u_{m+1}) \notin E}$), by decomposing into per-step adjacency and target checks, provides the gradient signal needed to learn the graph. This implies that progress in LLM planning RL depends less on better algorithms and more on better reward signals — specifically, on developing process rewards that can be automatically derived from environment feedback or learned verifiers, rather than requiring ground-truth graph knowledge.

Demoting SFT from a planning solution to merely a useful initialization. Theorem 3.1 formally characterizes SFT's stable point as a co-occurrence memorizer: the predicted probability $\hat{u}_m[k]$ converges to the empirical frequency $N_{i,j,k} / \sum_{k'} N_{i,j,k'}$. This means SFT cannot compose adjacency and reachability information that never co-occurs in the training data — it fundamentally cannot generalize planning to novel source-target pairs through transitivity alone. Figure 1 confirms this empirically: even when every edge appears in the SFT training data, the SFT model's learned adjacency matrix is a degraded copy missing low-frequency edges. Combined with the continual SFT degradation in Figure 2a (test accuracy steadily declines with more training), this establishes that SFT alone is fundamentally insufficient for planning — it provides a starting point, but RL (through exploration-driven data augmentation) is necessary to break the co-occurrence bound.

Follow-Up Research This Work Enables

Scaling process-reward Q-learning to natural language planning domains. The paper's Q-learning experiments use node tokens in a 100-node graph with a known adjacency matrix, where the process reward $\delta_{u_{m+1}=u_{\text{target}}} - \delta_{(u_m, u_{m+1}) \notin E}$ can be computed directly from the ground-truth graph. The most urgent follow-up question is: can process-reward Q-learning work when actions are natural language tokens and process rewards must be learned from data? A concrete experiment would replicate the Blocksworld setup but with natural language state descriptions and action sequences (as in the original Valmeekam et al., 2023a benchmark), training a verifier to produce per-step correctness signals from execution feedback, and comparing Q-learning against GRPO (the standard PG variant used in DeepSeek-R1). The key measurement would be whether Q-learning maintains its diversity advantage when the process reward is noisy and approximate, or whether verifier error introduces its own form of collapse. A negative result — Q-learning with learned process rewards performing worse than PG with KL — would establish a crucial boundary condition on the paper's claims.

Characterizing diversity collapse in deeper transformers and larger graphs. The paper proves diversity collapse (Theorem 4.3) for the abstract setting where logits $f(i, j)[k]$ are updated directly via gradient descent. The proof relies on the softmax parameterization and the on-policy update structure, not on model capacity, suggesting the result may generalize. But the empirical validation uses only a one-layer, single-head Transformer on 100-node graphs. A natural extension is to measure the diversity trajectory during PG training for 2-layer, 4-layer, and 12-layer Transformers on graphs of increasing size (200, 500, 1000 nodes), quantifying whether deeper models collapse more slowly (because they can maintain multiple parameter configurations that achieve high accuracy) or at the same rate (because the gradient dynamics are architecture-agnostic). This experiment would determine whether diversity collapse is a practical concern for modern-scale LLMs or primarily a theoretical curiosity of the minimal architecture studied here. The paper's Appendix G.1 briefly shows that attention patterns in a two-layer Transformer are consistent with Assumption 3.1, but reports no accuracy, diversity, or convergence metrics for the deeper model.

Learning process rewards without ground-truth graph knowledge. The process reward used in the paper requires knowing $A[j,k]$ (is the transition legal?) and the target node $i$. In practice, neither is available — the environment provides only execution feedback, not a labeled adjacency matrix. A critical follow-up would train a process reward model (PRM) on execution traces: given a partial action sequence and the current state, predict whether the sequence is on track toward the goal. This PRM could then provide the per-step rewards $R(u,m)$ for Q-learning. The experiment would train a PRM on SFT-generated trajectories labeled by whether they eventually reach the target, then use that PRM's scores as the reward signal for Q-learning, and compare against (a) Q-learning with oracle process rewards and (b) PG with KL regularization. The key question is whether learned process rewards preserve Theorem 5.2's structural convergence, or whether PRM errors introduce bias that causes Q-learning to converge to incorrect graph structures. The paper's finding that "last-step" PRM aggregation works best (Figure 13 in the appendix of the reference example) is suggestive — a PRM that accurately scores only the final step would collapse to outcome-reward behavior and fail per Theorem 5.1, so the PRM must be reliably accurate at intermediate steps.

Combining Q-learning with test-time search for planning. The paper analyzes Q-learning purely as a training method, but Q-functions naturally enable test-time planning — once the Q-values are learned, the model can perform lookahead by simulating $\max_{k'} f(i, k)[k']$ along candidate paths. A natural extension would train a model with process-reward Q-learning, then at inference time use the learned logits as Q-values to guide beam search: at each step, expand the top-$B$ candidates ranked by $f(i, j)[k]$, score resulting paths using the same Q-values, and select the best. This combines the paper's two theoretical insights — Q-learning's structural convergence (Section 5) and the benefits of search over sampling (from the reference paper's Section 5) — into a unified planning system. The experiment would compare: (a) greedy decoding from a Q-learning-trained model, (b) beam search over Q-values, (c) best-of-N sampling from a PG-trained model with KL regularization, (d) the compute-optimal scaling approach from the reference paper applied to this setting. Measuring accuracy vs. generation budget would reveal whether Q-value-guided search achieves better scaling than sampling from a PG-trained policy.

Theoretical analysis of when Q-learning's convergence rate becomes prohibitive. Theorem 5.2 proves linear convergence for process-reward Q-learning, with rate depending on the product of per-edge contraction factors along paths. For a path of length $L$, the effective contraction factor is $\prod_{\ell=1}^{L} |1 - 2\eta|^{N^{\text{prop}}_{i, v_\ell, v_{\ell+1}}}$, meaning longer paths converge exponentially slower. A theoretical follow-up would characterize the graph diameter threshold beyond which Q-learning's convergence rate becomes practically indistinguishable from non-convergence. Specifically, for an Erdős-Rényi graph with $n$ nodes and edge probability $p$, the expected longest shortest path between any reachable pair is approximately $\log n / \log(np)$. For $n = 100, p = 0.15$, this is $\log 100 / \log 15 \approx 4.6 / 2.7 \approx 1.7$ — paths are short. For $n = 10,000, p = 0.01$, this becomes $\log 10000 / \log 100 = 9.2 / 4.6 = 2$ — still modest. But for structured graphs with long bottleneck paths (trees, grids), the diameter scales as $O(n)$ or $O(\sqrt{n})$, making the product of contraction factors approach zero. This analysis would establish bounds on when Q-learning is a viable training method for planning versus when alternative approaches (hierarchical planning, subgoal decomposition) are necessary.

Practical Applications and Downstream Use Cases

Diagnostic toolkit for RL training of reasoning models. Organizations training reasoning models (DeepSeek-R1 style, Qwen-3 style) with GRPO or PPO can use this paper's analytical framework to diagnose training issues. The diversity metric — average number of distinct correct solutions per prompt — should be tracked during training alongside accuracy. If diversity drops below ~2 while training accuracy remains high, the model is in the diversity collapse regime identified by Theorem 4.3, and the paper's results predict that continued training will degrade test generalization (as observed in Figure 2a, where PG with $\lambda = 0$ peaks at ~0.91 test accuracy then declines). The corrective action is to increase KL regularization strength, following the trade-off characterized in Figure 2d: higher $\lambda$ preserves diversity at some cost to training accuracy, with the Pareto frontier showing the achievable accuracy-diversity combinations for a given base model. This provides a concrete, measurable protocol for tuning $\lambda$ that replaces the current practice of treating it as a black-box hyperparameter.

Process reward design for tool-use and agentic planning. The paper's stark contrast between outcome-reward Q-learning (complete failure, Figure 3a) and process-reward Q-learning (strong performance, ~0.95 test accuracy) has immediate implications for designing reward functions in LLM-based agents. In tool-use scenarios (APIs, code execution, database queries), practitioners typically provide outcome rewards (did the agent accomplish the task?) because they are cheap to compute. The paper's theory and experiments suggest this is fundamentally insufficient for Q-learning-based training — the reward must decompose into per-step signals that convey intermediate correctness. A practical design pattern: for tool-use, provide a reward of +1 at each step for valid tool calls (syntax-correct, type-correct, no runtime errors) and an additional +1 upon task completion, with a penalty for invalid calls. This approximates the process reward structure $\delta_{u_{m+1}=u_{\text{target}}} - \delta_{(u_m, u_{m+1}) \notin E}$ without requiring ground-truth adjacency, using execution feedback as a proxy for the adjacency check. The paper's theoretical guarantee — that this structure converges to equal weighting of all valid next actions (logit value 1) — means the trained agent would maintain multiple viable strategies, improving robustness to tool failures.

Training planning models with off-policy data from cheaper sources. Theorem 5.2's convergence proof does not require on-policy data — only persistent exploration (Assumption 5.1), which is satisfied by any data generation process that visits all triples $(i,j,k)$ with positive asymptotic frequency. Combined with the empirical result that off-policy Q-learning matches on-policy performance (Figure 3a, green vs. orange curves), this enables a training pipeline where exploration data can be generated by a cheaper model (distilled, quantized, or even rule-based) and the target model is trained via Q-learning on this off-policy data. In practice, a large reasoning model could be trained on trajectories generated by a smaller, faster draft model that explores broadly (perhaps with higher temperature or explicit $\epsilon$-exploration), decoupling the cost of exploration from the cost of the model being trained. This is not possible with PG, where the gradient structure (Theorem 4.3) relies on the on-policy sampling distribution matching the current policy. Modern RLHF frameworks like VeRL (Sheng et al., 2024) already implement partially off-policy updates due to the lag between rollout generation and gradient computation — this paper provides the first theoretical justification for why Q-learning handles this gracefully while PG may suffer.

When to Prefer PG vs. Q-Learning for LLM Planning

The paper articulates a clear trade-off between policy gradient and Q-learning based on the reward structure available and the diversity requirements of the task:

  • Prefer policy gradient (with tuned KL regularization) when: (1) Only outcome rewards are available (task success/failure), and process rewards cannot be reliably derived from environment feedback. PG works with outcome rewards (Theorem 4.1, Figure 2a), while outcome-reward Q-learning collapses (Theorem 5.1, Figure 3a). (2) The base model is already capable on the target distribution, so KL regularization preserves useful diversity without overly constraining learning (Theorem 4.4's stable point stays close to $q_{\text{base}}$). (3) The deployment setting requires inference-time decoding strategies like best-of-N sampling or majority voting, which benefit from the policy maintaining non-zero probability on multiple valid outputs.

  • Prefer Q-learning with process rewards when: (1) Process rewards can be designed or learned — either from ground-truth environment structure (adjacency matrices, execution feedback) or from trained verifiers. The process reward structure $\delta_{u_{m+1}=u_{\text{target}}} - \delta_{(u_m, u_{m+1}) \notin E}$ is necessary to avoid Q-value bias (Theorem 5.2 vs. Theorem 5.1). (2) Output diversity is critical — for tasks where the model must explore multiple solution strategies, avoid mode collapse, or provide diverse candidate plans for downstream filtering. Q-learning's stable point naturally assigns equal high logits to all valid actions (Figure 4), while PG collapses diversity without KL regularization (Figure 2c). (3) Off-policy data is being used — for training on trajectories from a different model, a replay buffer, or batched rollouts where the policy changes between generation and gradient computation (relevant to VeRL-style frameworks). Q-learning's off-policy convergence (Theorem 5.2, Figure 3a green curve) does not require on-policy data, unlike PG.

  • Prefer SFT alone only when: The planning task requires only interpolation within the training distribution — novel source-target pairs never appear at test time, meaning the co-occurrence memorization limitation (Theorem 3.1) does not apply. The paper's continual SFT results (Figure 2a, declining test accuracy) suggest that even within-distribution, SFT alone may be suboptimal if the training data is insufficiently dense relative to the graph's connectivity. Practically, this means SFT should be treated as initialization for RL, not as a standalone planning solution.