PaperlayerView exact revision on arXiv

Source-linked AI summary

Balancing Reference Guidance and Free Generation in Trajectory Rollouts for Reasoning RL

Hanyu Wang, Nakul Agarwal, Hossein Nourkhiz Mahjoub, Ehsan Moradi Pari, Makoto Fukushima, Jinghui Chen, Vaishnav Tadiparthi

arXiv:2610.11128v1cs.AI

TL;DR

On difficult questions, a model may rarely generate correct responses, while direct reference training can target trajectories unlike its own. The paper analyzes reference-prefix continuation and learns a shared prefix selector from continuation outcomes. ARG achieves the highest aggregate pass@12 across five mathematical reasoning benchmarks on two Qwen3 scales, with competitive average sampled accuracy.

  • Problem

    On difficult questions, models may rarely generate correct responses, while reference solutions can differ substantially from trajectories the model would produce itself.

    Source: §1 · p. 1

  • Method

    ARG learns a prefix selector shared across questions from continuation outcomes, then generates from selected reference prefixes within a fixed budget and falls back to the reference if no continuation succeeds.

    Source: §1 · p. 2

  • Results

    ARG achieves the highest aggregate pass@12 across five mathematical reasoning benchmarks on Qwen3-4B and Qwen3-8B, with competitive average sampled accuracy.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    Continuation success alone does not determine preferred guidance because longer prefixes tend to raise success probability while lowering reference surprisal.

    Source: §1 · p. 2

    ARG is applied only to GRPO, and experiments cover mathematical reasoning with one model family.

    Source: §7 · p. 10

Abstract

from arXiv · show

A verified reference solution provides a correct trajectory for training a reasoning model. Alternatively, a prefix of the reference can guide the model in generating a trajectory of its own. How much reference guidance should we provide? We study this question through prefix continuation, where the model continues from a reference prefix and keeps the resulting trajectory if it passes verification, falling back to the reference otherwise. Since both procedures produce correct trajectories, we compare their distributions with the ideal distribution, the model's own distribution conditioned on successful verification. For one continuation, we derive the KL divergence in closed form, which, up to a bounded term, decreases with the product of the probability of generating a different correct trajectory and the reference surprisal, the negative log probability of the reference suffix given the prefix. Since a longer prefix tends to raise the former but lowers the latter, continuation success alone does not determine the preferred amount of guidance. From this analysis, we learn a prefix selector shared across training questions from continuation outcomes, without estimating success probabilities or additional generation. The resulting Adaptive Reference Guidance (ARG) constructs correct trajectories within a fixed generation budget, and we apply it to all-failure groups in Group Relative Policy Optimization (GRPO). Experiments on Qwen3-4B and Qwen3-8B across five mathematical reasoning benchmarks show that ARG achieves the highest aggregate pass@12 among the evaluated methods with competitive average sampled accuracy.

1 INTRODUCTION

On difficult questions, a model may rarely generate correct responses, while a reference solution may differ substantially from its own trajectories. Reference-prefix continuation offers adjustable guidance, but choosing its amount requires balancing success against how much of the trajectory the model generates itself.

1 INTRODUCTION

  • Longer prefixes make continuations more likely to be correct but leave more reference content the model would not generate itself.A shorter prefix better matches the model’s own generation but increases the chance of failure; the two extremes correspond to direct reference training and standard RLVR.Source: §1 · p. 2
  • Up to a small bounded term, KL divergence to the model’s own correct-sample distribution decreases with d_tS_t, combining successful replacement probability and reference surprisal.Longer prefixes tend to raise d_t while lowering S_t, so continuation success alone does not determine preferred guidance.Source: §1 · p. 2
  • ARG learns a prefix selector shared across questions from continuation outcomes, avoiding per-prefix success-probability estimates and additional generation.The selector chooses among guidance levels and is updated using the same continuations used for training.Source: §1 · p. 2
  • ARG samples guidance levels, generates continuations within a fixed budget, and falls back to the reference if none succeeds.The method is applied to all-failure groups in GRPO, which otherwise provide no learning signal.Source: §1 · p. 2
  • On Qwen3-4B and Qwen3-8B across five mathematical reasoning benchmarks, ARG achieves the highest aggregate pass@12 at 67.33% and 70.00%, respectively.Its average sampled accuracy remains competitive, with a small gap to the strongest baseline.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries define correctness through a binary verifier and describe how reference-prefix continuation produces accepted trajectories. They contrast the resulting distribution with both the reference-only distribution and the model’s own correct-sample distribution.

2 PRELIMINARIES

  • A trajectory includes reasoning, the final answer, and a termination token; a binary verifier marks it correct when its reward is 1.The model’s success probability is the probability it assigns to the set of correct trajectories.Source: §2 · p. 2
  • Reference-prefix continuation splits a correct reference into a prefix and suffix, then samples a continuation conditioned on the question and prefix.At prefix length zero, generation uses only the original question.Source: §2 · pp. 2–3
  • The candidate trajectory is kept if correct and otherwise replaced by the reference, so the accepted-trajectory distribution assigns probability only to correct trajectories.Training on the reference alone instead gives the point-mass distribution δ_r.Source: §2 · p. 3
  • The ideal distribution is the model’s own correct-sample distribution, obtained by sampling from the model and keeping only correct trajectories.Likelihood learning on trajectories drawn from this distribution increases success probability in expectation at the current model.Source: §2 · p. 3
  • The paper relates this ideal distribution to likelihood learning (Phan et al., 2023, Singh et al., 2023) and measures distance from it using KL divergence.The likelihood-learning connection is stated for a finite trajectory space and a differentiable policy with positive probabilities.Source: §2 · p. 3

3 ANALYZING AND SELECTING REFERENCE PREFIXES

The section analyzes how reference-prefix length affects distance from the model’s own correct trajectories, then develops a shared selector and Adaptive Reference Guidance (ARG). ARG generates verified trajectories and applies them to all-failure groups in GRPO.

3.1 DISTANCE TO THE IDEAL DISTRIBUTION

  • The distance to the ideal distribution depends mainly on d_tS_t: longer prefixes can raise successful replacement probability but lower reference surprisal, so success alone does not choose the best prefix.Maximizing d_tS_t yields the smallest distance up to 1/e.Source: §3.1 · pp. 3–4

3.2 CHALLENGES IN FINDING THE PREFERRED PREFIX

  • Finding a question-specific preferred prefix is costly because d_t varies across questions and model updates, and estimating it across prefix lengths requires many continuations on hard questions.The reference surprisal is available from one model forward pass, but successful replacement probability is unknown.Source: §3.2 · p. 4

3.3 LEARNING A PREFIX SELECTOR ACROSS QUESTIONS

  • A shared selector samples guidance levels defined by fractions of reference surprisal, optimizing average d_tS_t across guided questions rather than searching separately for each question.Normalizing by total reference surprisal prevents long or unfamiliar references from dominating the average while preserving each question’s within-question optimum.Source: §3.3 · p. 4
  • Each continuation supplies selector feedback: successful replacements receive the remaining reference-surprisal fraction, while failures and exact reference copies receive zero, yielding an unbiased estimate of the objective.The feedback is pooled across questions and levels to learn the selector without estimating d_t or generating additional continuations.Source: §3.3 · p. 5

4 ARG: ADAPTIVE REFERENCE GUIDANCE

  • ARG builds on the distance analysis and selector update, while leaving the choice of guided questions and policy-learning procedure to the reinforcement-learning algorithm.Its generation and selector-update steps do not depend on a particular reinforcement-learning algorithm.Source: §4 · p. 5

4.1 GENERATING ACCEPTED TRAJECTORIES

  • Within a fixed continuation budget, ARG accepts correct candidates different from the reference and falls back to the reference if none succeeds.When limiting accepted trajectories, it prefers empty-prefix candidates first, then the shortest nonempty prefix.Source: §4.1 · pp. 5–6

4.2 UPDATING THE PREFIX SELECTOR

  • ARG updates the selector from continuations collected with the current model and selector fixed; the update follows the gradient of J in expectation.The same continuations provide feedback whether or not their candidate trajectories are accepted.Source: §4.2 · p. 6

4.3 APPLYING ARG TO GRPO

  • For GRPO all-failure groups, ARG uses K=8 continuations and inserts at most four accepted trajectories, retaining at least four original failures.Inserted trajectories are learned under the original question through their complete likelihood; for binary rewards, the expected update connects the distance from the ideal distribution to policy learning.Source: §4.3 · pp. 6–7

5 EXPERIMENTS

The experiments evaluate ARG on mathematical reasoning benchmarks and compare its results with supervised, reinforcement-learning, and reference-informed baselines. Ablations and training dynamics examine the roles of reference guidance, fallback, surprisal weighting, and the learned selector.

5.1 EXPERIMENTAL SETTING

  • Evaluation uses Qwen3-4B and Qwen3-8B on five mathematical benchmarks, drawing 12 responses per question and treating pass@12 as the primary metric.Baselines include the initial Qwen model, SFT, GRPO, OPSD, SAGE, BREAD, and LUFFY.Source: §5.1 · p. 7

5.2 RESULTS

  • ARG achieves 67.33% pass@12 on Qwen3-4B and 70.00% on Qwen3-8B, the highest aggregate results, improving over GRPO by 3.33 and 7.33 percentage points, respectively.It also exceeds reference-informed baselines in aggregate pass@12.Source: Table 1 · p. 7, §5.2 · p. 7

5.3 ABLATIONS

  • On Qwen3-4B, reference insertion reaches 56.67% pass@12 and 28.44% avg@12, below ARG’s 67.33% and 41.11%.This ablation replaces a random failure with the reference without extra generation.Source: §5.3 · p. 8, Table 2 · p. 8
  • On Qwen3-4B, ARG improves over resampling from the empty prefix from 60.00% to 67.33% pass@12 and from 34.00% to 41.11% avg@12.Candidate trajectories were accepted in 66.0% of all-failure groups with ARG versus 32.5% with resampling.Source: §5.3 · p. 8, Table 2 · p. 8
  • Removing surprisal weighting lowers pass@12 by 4.00 percentage points on Qwen3-4B and 6.00 points on Qwen3-8B.Both evaluation metrics fall on both models in this ablation.Source: §5.3 · p. 8, Table 2 · p. 8

5.4 TRAINING DYNAMICS

  • Across training windows, α = 1/4 had higher replacement success, but final feedback favored α = 1/2 (0.0991 versus 0.0802), whose selection probability ended highest at 49.0%.The final-window success rates were 43.84% for α = 1/4 and 23.65% for α = 1/2.Source: §5.4 · p. 8, Figure 1 · p. 9

6 RELATED WORK

  • Unlike methods using fixed or scheduled guidance, observed success, or prefix or hint gains, ARG derives guidance from accepted-trajectory distance and learns a selector across questions.The analysis indicates that continuation success alone does not determine the preferred prefix.Source: §6 · p. 9

7 CONCLUSION AND LIMITATIONS

The conclusion frames ARG as a selector learned from continuations, guided by the distance between accepted trajectories and the ideal distribution. It reports strong aggregate pass@12 results while noting limits in the algorithm, guidance levels, and evaluation scope.

7 CONCLUSION AND LIMITATIONS

  • Up to a small bounded term, distance to the ideal distribution decreases with d_tS_t, linking successful replacement with reference surprisal.This analysis motivates learning a prefix selector across questions from continuations themselves.Source: §7 · p. 9
  • Applied to all-failure GRPO groups, ARG achieves the highest aggregate pass@12 on Qwen3-4B and Qwen3-8B across five mathematical benchmarks, with competitive average sampled accuracy.Source: §7 · p. 10
  • The experiments apply ARG only to GRPO, use a predefined set of guidance levels, and cover mathematical reasoning with one model family.The authors identify other domains with verifiable rewards as a possible extension.Source: §7 · p. 10

A PROOFS FOR THE DISTANCE TO THE IDEAL DISTRIBUTION

A.1 SETUP

  • The setup uses a finite trajectory space and defines the ideal distribution ν as the current policy conditioned on correct trajectories; S_t is the negative log probability of the reference suffix given its prefix.For the empty prefix, S_0 is the reference surprisal.Source: §A.1 · p. 14

A.2 THE IDEAL DISTRIBUTION

  • The appendix states that learning likelihoods under the ideal distribution increases success probability in expectation; training on the reference has KL distance S_0 + log p, bounded between 0 and S_0.If ν assigns positive probability to another trajectory, the reverse KL from ν to the reference point mass is infinite.Source: §A.2 · p. 14

A.3 PROOF OF THEOREM 1

The proof derives the accepted-trajectory distribution and uses the monotonicity of reference surprisal to establish the closed-form distance comparison. With fallback, every prefix length yields a distance no greater than training on the reference.

A.3 PROOF OF THEOREM 1

  • The closed-form derivation combines the accepted-trajectory distribution with the relation between prefix probability and reference surprisal.The proof uses two lemmas to express the likelihood ratios and then collect terms in the KL divergence.Source: §A.3 · p. 15
  • Non-reference trajectories accepted under a prefix are correct continuations of that prefix, and their total probability is the successful-replacement probability.The reference receives the remaining probability under the accepted-trajectory distribution.Source: §A.3 · p. 15
  • Reference surprisal is nonincreasing as the prefix grows, because each longer prefix has probability no greater than a shorter prefix.The proof obtains this from autoregressive factorization of the reference probability.Source: §A.3 · p. 15
  • Fallback keeps the accepted-trajectory distribution no farther from the ideal distribution than training on the reference, for every prefix length.The closed form combines the successful-replacement probability and reference surprisal, with a bounded correction term.Source: §A.3 · pp. 15–16

A.4 CONSEQUENCES OF THEOREM 1

The distance to the ideal distribution can be read through the reference content retained in accepted trajectories, and examples show why prefix choice depends on both successful replacement and reference surprisal. The appendix then develops selector feedback and its relation to normalized distance.

A.4 CONSEQUENCES OF THEOREM 1

  • Apart from a prefix-independent term and a bounded correction, the distance is the expected surprisal of reference content not replaced by the model.Both nearly full prefixes and empty prefixes on hard questions can retain nearly all of the reference surprisal, through the prefix or fallback respectively.Source: §A.4 · p. 16

A.5 A FINITE EXAMPLE

  • In both finite examples, the empty prefix has larger successful-replacement probability times surprisal and smaller distance, despite supplying H increasing continuation success.The second example also shows successful-replacement probability alone can favor a different prefix from the theorem's criterion.Source: §A.5 · pp. 16–17

B PROOFS FOR LEARNING THE PREFIX SELECTOR

  • The appendix organizes the selector analysis around feedback properties, its theorem guarantee, sharing and randomized guidance, GRPO updates, and accumulated sampling noise.It identifies these as the topics of Appendices B.1 through B.5.Source: B PROOFS FOR LEARNING THE PREFIX… · p. 17

B.1 SETUP AND SELECTOR FEEDBACK

  • Expected selector feedback is proportional to successful replacement probability times reference surprisal, normalized by the empty-prefix surprisal; larger feedback corresponds to smaller normalized distance up to a bounded term.Across prefixes of one reference, the normalization preserves the ranking of the product, though not necessarily the exact distances.Source: §B.1 · p. 17

B.2 PROOF OF THEOREM 2

The selector objective tracks average normalized distance to the ideal distribution up to a bounded correction, and its feedback supports updates across shared guidance levels and GRPO batches. The appendix also analyzes the costs of sharing and randomization and bounds accumulated sampling noise.

B.2 PROOF OF THEOREM 2

  • A larger selector objective corresponds to smaller average normalized distance up to a bounded correction, but the two can rank selectors differently.The correction is informative when the expected reciprocal empty-prefix surprisal is small and can be large when some surprisals approach zero.Source: §B.2 · p. 18

B.3 SHARING AND RANDOMIZING GUIDANCE LEVELS

  • Sharing one guidance-level distribution across questions can incur a cost, while randomizing levels gives a distance no greater than their average distance.The sharing cost is zero if one level maximizes feedback for almost every question; the randomization comparison follows from the mutual information between level and accepted trajectory.Source: §B.3 · pp. 19–20

B.4 SELECTOR UPDATES IN THE GRPO APPLICATION

  • For a fixed-step batch, the expected selector update scales with batch size, continuation count, and all-failure probability; changes in those probabilities also change question weights and update direction.The expected update follows the current objective in expectation as the model, question distribution, and prefix mapping change.Source: §B.4 · pp. 20–21

B.5 ACCUMULATED SAMPLING NOISE

  • The appendix bounds accumulated sampling noise around changing expected selector directions over a deterministic number of iterations.It first bounds individual update terms and then uses centered errors with vanishing cross terms.Source: §B.5 · pp. 21–22

C ANALYSIS OF MULTIPLE CONTINUATIONS

The appendix extends the one-continuation analysis to repeated continuations from a fixed prefix and to general acceptance rules. Repeated attempts reduce distance toward the ideal distribution restricted to that prefix, while a surprisal condition gives a sufficient comparison with reference-only training for general acceptance.

C.1 REPEATED CONTINUATIONS FROM A FIXED PREFIX

  • With repeated continuations from a fixed prefix, the accepted distribution approaches the ideal distribution restricted to that prefix, and its distance decreases toward a prefix-dependent limit.For an empty prefix the limit is the full ideal distribution and the limiting distance is zero; a nonempty prefix has the same limit only if it contains all correct trajectories with positive probability.Source: §C.1 · pp. 22–23

C.2 A SUFFICIENT CONDITION FOR GENERAL ACCEPTANCE RULES

  • The general rule allows selecting among correct candidates, including excluding exact copies, prioritizing candidates, and accepting multiple candidates before selecting one.The analysis permits dependent continuations as long as each candidate has the specified conditional marginal law.Source: §C.2 · pp. 23–24
  • For general acceptance among correct candidates, if the sum of exp(−Stj) over selected prefixes is at most one, the accepted-trajectory distance is no greater than reference-only training.A sufficient condition is that every prefix surprisal is at least log K; this also applies when guidance levels are sampled before generation.Source: §C.2 · pp. 23–24

D ANALYSIS OF LEARNING FROM INSERTED TRAJECTORIES

The appendix analyzes how accepted trajectories contribute to GRPO updates, from their token-level gradient to the effects of inserting one or several correct trajectories. It also distinguishes learning a supplied prefix from learning its continuation.

D.1 THE GRADIENT OF INSERTED TRAJECTORIES

  • Under fixed advantages and active inserted tokens, Equation (13) gives each inserted trajectory a gradient that acts as positively weighted supervised fine-tuning when its advantage is positive.The advantage sets the learning weight, while candidate generation and acceptance determine the trajectory distribution.Source: §D.1 · p. 24

D.2 THE GROUP UPDATE WITH ONE INSERTED TRAJECTORY

  • With one inserted trajectory, the expected group update favors verification success and reflects which correct trajectories Q favors; reference training instead favors the particular reference.When Q is the ideal distribution, the expected update is proportional to the gradient of the original question’s success probability.Source: §D.2 · pp. 25–26

D.3 MULTIPLE INSERTED TRAJECTORIES

  • For multiple insertions, Q+ weights realized sets by their total positive advantage, accounting for both acceptance and the resulting group composition.With one insertion, the expression recovers the one-trajectory result.Source: §D.3 · p. 26

D.4 LEARNING THE PREFIX AND THE CONTINUATION

  • Learning a complete accepted trajectory covers producing its supplied prefix and completing it, whereas learning only the continuation concerns correctness given an external prefix.These events coincide for an empty prefix; fallback and successful continuations contribute through a mixture identity.Source: §D.4 · p. 27

E IMPLEMENTATION AND EVALUATION DETAILS

E.1 TRAINING PROCEDURE

The implementation uses screened DeepMath questions with matched offline references and specifies how ARG constructs prefixes, inserts accepted trajectories, updates the policy, and evaluates checkpoints.

E.1 TRAINING PROCEDURE

  • Training uses 11,068 screened DeepMath examples with matched offline DeepSeek-R1 references; questions solved in all eight initial Qwen3-4B attempts are removed.The procedure also filters by difficulty, verifier compatibility, and whether the reference ends in the gold answer.Source: §E.1 · p. 27

E.2 PREFIX BOUNDARIES

  • For α < 1, ARG selects a prefix near the chosen fraction of reference surprisal, adjusted to a valid sentence or formula boundary while leaving a nontrivial suffix.The empty prefix corresponds to α = 1; the prefix counts toward the maximum trajectory length.Source: §E.2 · pp. 27, 29

E.3 POLICY LOSS

  • GRPO recomputes group advantages from verifier and format rewards; inserted trajectories use a stopped-current-probability ratio denominator and include supplied prefixes and termination tokens.All-equal group rewards yield zero advantages, and ordinary responses use asymmetric clipping.Source: §E.3 · p. 29

E.4 TRAINING CONFIGURATIONS AND QUESTION PROMPT

  • The compared GRPO-based runs share optimizer and student-rollout settings, while ARG generation uses the selected reference prefix and evaluation generation uses no reference prefix.Runs use 32 questions per update and one optimization epoch per collected batch.Source: §E.4 · pp. 29–30

E.5 CHECKPOINT AND EVALUATION RECORDS

  • All reported entries use the same 150 evaluation questions across five benchmarks, with 12 scored responses per question and a shared scoring protocol.Pass@12 records whether any response is correct; avg@12 is the fraction correct, averaged over questions.Source: §E.5 · pp. 30–31

F ADDITIONAL MAIN RESULTS

ARG improves aggregate avg@12 over GRPO on both model sizes, while benchmark-level ablation results show higher average sampled accuracy for ARG on every benchmark.

F ADDITIONAL MAIN RESULTS

  • 2.67 and 7.45 percentage points: ARG improves aggregate avg@12 over GRPO on Qwen3-4B and Qwen3-8B, respectively, and ranks second on both models.Table 7 reports avg@12 percentages averaged over the five benchmarks.Source: F ADDITIONAL MAIN RESULTS · p. 31, Table 7 · p. 31

G ADDITIONAL ABLATION RESULTS

  • The additional ablations examine benchmark-level results and how ablations affect accepted trajectories and prefix selection during training.Reference insertion and resampling use Qwen3-4B; surprisal-weighting comparisons cover both model sizes.Source: G ADDITIONAL ABLATION RESULTS · p. 31

G.1 EVALUATION RESULTS

  • For the ablations in Table 9, the reported comparison finds that ARG has higher average sampled accuracy than the ablations on every benchmark.Source: §G.1 · p. 31

G.2 ACCEPTED CANDIDATE TRAJECTORIES DURING TRAINING

The appendix reports accepted-candidate rates during training and compares the selector’s behavior with and without surprisal weighting. Acceptance frequency alone does not account for evaluation performance.

G.2 ACCEPTED CANDIDATE TRAJECTORIES DURING TRAINING

  • 318 of 482 Qwen3-4B all-failure groups accept candidates with ARG, versus 145 of 446 with resampling, using eight continuations per group.Across the first and last windows, ARG’s acceptance proportion rises from 59.4% to 71.8%, while resampling falls from 37.6% to 26.4%.Source: Figure 2 · p. 32, §G.2 · p. 32
  • Acceptance frequency alone does not explain evaluation results: on Qwen3-8B, the unweighted selector accepts 64.8% of all-failure groups versus 60.2% for ARG despite lower evaluation scores.On Qwen3-4B, the corresponding proportions are similar: 65.9% and 66.0%.Source: §G.2 · p. 32

G.3 SURPRISAL WEIGHTING AND PREFIX SELECTION

  • Figure 3 shows unsmoothed prefix-selector selection probabilities with and without surprisal weighting after each training step; steps without an update retain prior probabilities, and smaller α levels leave less reference surprisal to the continuation.Source: Figure 3 · p. 33

H TRAINING DYNAMICS

The appendix extends the training-dynamics analysis to Qwen3-4B and tracks accepted-trajectory sources for both models.

H.1 SUCCESSFUL REPLACEMENT, FEEDBACK, AND PREFIX SELECTION

  • Across the run, mean feedback is 0.0754 for α = 1/4 and 0.0588 for α = 1/2; α = 1/4 leads successful replacement in every window.The selection curves also show a persistent preference for α = 1/4.Source: §H.1 · p. 33, Figure 4 · p. 34

H.2 SOURCES OF ACCEPTED TRAJECTORIES

  • Candidate trajectories are accepted in 66.0% of Qwen3-4B and 60.2% of Qwen3-8B all-failure groups; most accepting groups use guided candidates, while fallback remains throughout training.The fallback share varies across windows as the all-failure questions change.Source: Figure 5 · p. 34, §H.2 · p. 34
  • Figure 5 counts each all-failure group once, classifying it by accepted guided candidates, unguided candidates, both, or reference fallback.A mixed group accepts both guided and unguided candidate trajectories.Source: Figure 5 · p. 34, §H.2 · p. 34
Loading 2610.11128v1…