PaperlayerView exact revision on arXiv

Source-linked AI summary

What Shared Prefixes Hide: Trajectory Dropout for On-Policy Distillation

Zizhuo Lin, Quanling Liu, Yi Yang, Yawei Luo

arXiv:2609.33455v2cs.AI

TL;DR

OPD supplies dense teacher feedback on student-generated trajectories, but the resulting learning signals can be concentrated and weakened by shared prefixes. The paper introduces Trajectory Dropout, which changes the student’s training view while preserving full-context rollouts and teacher targets, and reports gains across benchmarks, model scales, and OPD variants.

  • Problem

    Dense token-level feedback in OPD does not necessarily yield dense learning signals: approximately 70% of supervised tokens contribute less than 0.5% of measured own-logit signal mass.

    Source: §1 · p. 2

  • Method

    Trajectory Dropout masks selected parts of the saved reasoning history during student updates while keeping the sampled trajectory and teacher targets in the full-context view.

    Source: §1 · p. 2, §3.2 · p. 6

  • Results

    Trajectory Dropout improves average performance across six reasoning benchmarks and two out-of-domain benchmarks, across model scales, and when applied to existing OPD variants.

    Source: §1 · p. 3, §6 · p. 10

  • Takeaways & Limitations

    Across the reported settings, trajectory dropout strengthens corrective responses and selectively introduces recovery signals for prefix-dependent predictions.

    Source: §6 · pp. 9–10

    Strong masking can occasionally move the student probability below the teacher probability and reverse the correction direction; a sign-preserving fallback is discussed.

    Source: §3.2 · p. 6

Abstract

from arXiv · show

On-policy distillation (OPD) trains a student model on its own trajectories using dense token-level feedback from a stronger teacher model. Since each update is conditioned on the reasoning prefix already generated by the student, the prefix also shapes how effectively teacher feedback is converted into learning. We find that shared prefixes can lead to weak token-level updates, a phenomenon we call Prefix-Induced Supervision Attenuation (PISA). This attenuation arises in two common cases. (i) High student confidence can weaken corrective gradients even when the teacher disagrees. (ii) Tokens that rely on earlier reasoning can receive learning signals as weak as those for simple local continuations. To solve this problem, we propose Trajectory Dropout, a simple training-time intervention that exposes these weakened signals. The student first performs a standard full-context rollout to generate a complete trajectory. During training, we randomly drop a certain proportion of the student's reasoning trajectory, while the teacher continues to observe the complete trajectory for token-level supervision. This intervention strengthens corrections for overconfident predictions and introduces additional supervision at prefix-sensitive positions. Trajectory Dropout consistently improves average performance across teacher--student model pairs of different scales and six mathematical reasoning benchmarks, while also yielding gains on two out-of-domain benchmarks. It can also be flexibly integrated into existing OPD variants with negligible computational overhead, further improving their performance. These results demonstrate that Trajectory Dropout provides a simple mechanism for strengthening token-level supervision across model scales and OPD objectives.

1 INTRODUCTION

OPD provides dense token-level feedback, but its updates can remain concentrated and weak at useful positions. The paper identifies prefix-induced attenuation and proposes Trajectory Dropout to strengthen these signals, reporting gains across benchmarks and OPD variants.

1 INTRODUCTION

  • Approximately 70% of supervised tokens contribute less than 0.5% of measured own-logit signal mass, despite teacher feedback covering the trajectory.Weak updates can occur at positions where the teacher still disfavors the prediction or where full history makes qualitatively different predictions appear equally easy.Source: §1 · p. 2
  • Prefix-Induced Supervision Attenuation (PISA) arises when confidence weakens teacher corrections or full-history agreement obscures prefix-dependent reasoning.The latter can make reasoning-dependent predictions receive signals as weak as routine local continuations.Source: §1 · p. 2
  • Trajectory Dropout preserves standard full-context rollouts and full-prefix teacher supervision while changing the student’s access to saved reasoning during updates.It strengthens existing corrections and adds signal at prefix-sensitive positions while preserving nearby access to limit disruption to local predictive dependencies.Source: §1 · p. 2
  • Trajectory Dropout consistently improves average performance across six mathematical reasoning benchmarks and two out-of-domain benchmarks, including when added to TRD and Fast-OPD.The reported gains extend across teacher–student pairs at different scales.Source: §1 · p. 3
  • A short protection window after masked segments stabilizes training when masking is fragmented by preserving nearby predictive access.Ablations over dropout budget and protection window show that widely scattered masking can disrupt nearby dependencies.Source: §1 · p. 3

2 PREFIX-INDUCED SUPERVISION ATTENUATION

In sampled-token OPD, token-level learning signals depend on both teacher–student disagreement and the student’s output sensitivity. Shared prefixes can weaken those signals through confidence saturation or by obscuring the role of earlier reasoning.

2.1 TOKEN-LEVEL LEARNING SIGNALS IN ON-POLICY DISTILLATION

  • Sampled-token OPD scores each generated token on its causal prefix, and the advantage is the teacher–student log-probability gap for that token.A positive advantage favors increasing the token probability; a negative advantage favors decreasing it.Source: §2.1 · p. 3
  • The output-logit signal combines a token’s advantage and output sensitivity, and signal mass sums these token-level norms; a full-context audit shows an uneven distribution.The token contribution is measured before optimizer transformations.Source: §2.1 · p. 3

2.2 CONFIDENCE AND AGREEMENT SATURATION

  • For fixed teacher probability, the own-logit correction approaches zero as student confidence approaches one, weakening correction despite teacher disagreement.This is the confidence-saturation route to a weak update.Source: §2.2 · p. 4, Figure 2 · p. 4
  • Near-zero teacher–student advantage can give weak updates to both routine local continuations and predictions that rely on earlier reasoning, hiding their different sources of predictability.The concern is not agreement itself, but that full-view agreement can bury useful prefix-dependent learning opportunities among routine continuations.Source: §2.2 · p. 4
  • 22.88%, 16.39%, and 17.70% were the signal shares of the same R_I ∪ R_II tokens across three checkpoints; 48.85%, 46.70%, and 45.58% of R_II tokens showed positive signal gain.The controlled prefix intervention revealed a substantial prefix-sensitive subset within the near-agreement region.Source: §2.2 · p. 5

3 TRAJECTORY DROPOUT FOR OPD

Trajectory Dropout preserves full-context student rollouts and teacher targets while changing the student’s training view. The resulting cross-view update can strengthen confidence-sensitive corrections and add signal at prefix-sensitive near-agreement positions.

3.1 FULL-CONTEXT ROLLOUTS WITH MASKED-VIEW TRAINING

  • TD improves the student’s training view while preserving standard full-context rollouts and teacher supervision on complete causal prefixes.The student generates and saves a complete trajectory before reprocessing it with restricted attention; evaluation uses full context.Source: §3.1 · p. 5, Figure 4 · p. 5
  • Training samples contiguous response spans totaling fraction ρ, masks later targets’ direct access to them, and can preserve the next w targets’ cues with a protection window.The prompt remains visible, and every response target contributes to the loss.Source: §3.1 · pp. 5–6

3.2 CROSS-VIEW LEARNING SIGNALS

  • The cross-view objective scores the saved trajectory with the same mask for current and pre-update student states, while the teacher scores complete causal prefixes.The student models share parameters and differ only in attention access; the student-side OPD scores are replaced with masked-view scores.Source: §3.2 · p. 6
  • Masking changes both the student’s sensitivity and the advantage shaping the token-level update.For highly confident tokens, lowering the masked-view probability can increase the own-logit derivative and strengthen a downward correction, though strong masking can reverse its direction.Source: §3.2 · p. 6
  • The sufficient condition for strengthening a correction is not universal: strong masking can move the masked-view probability below the teacher probability and reverse correction direction.The authors refer to a sign-preserving fallback for these cases.Source: §3.2 · p. 6
  • At near-agreement positions, prefix-sensitive predictions can acquire additional updates that distinguish them from positions remaining weakly updated.The induced signal differentiates positions that full-view agreement alone cannot resolve.Source: §3.2 · p. 6

4 EXPERIMENTS

The experiments evaluate Trajectory Dropout (TD) across mathematical reasoning and out-of-domain benchmarks, model scales, and OPD variants. Results include improved average performance, with ablations examining dropout ratio, local protection, teacher supervision, and the source of TD’s update advantage.

4.1 MODELS AND DATASETS

  • The main setup uses Qwen3-4B-Instruct-2507 as teacher and Qwen3-1.7B as student, training on teacher-correct English DAPO-Math-17K data.Evaluation covers six mathematical reasoning benchmarks plus MBPP+ and GPQA Diamond, using full-prefix inference.Source: §4.1 · p. 7

4.2 RESULTS

  • TD raises Qwen3-1.7B OPD Avg from 28.27 to 32.78 (+4.5 percentage points), and also improves Fast-OPD and TRD across all six benchmarks.Fast-OPD rises from 29.48 to 33.76 and TRD from 18.76 to 22.68; with thinking enabled, OPD+TD reaches 47.29 Avg versus 45.40 for OPD.Source: §4.2 · p. 8

4.3 DROPOUT ANALYSIS

  • Across six mathematical benchmarks, performance rises as dropout ratio increases from 0 to 0.2, then gradually declines; the best result is at 0.2.The default ratio is ρ = 0.2.Source: §4.1 · p. 7, §4.3 · p. 8, Figure 6 · p. 8
  • Single-token masking at ρ = 0.2 sharply lowers accuracy and destabilizes training without protection; even a short protection window substantially recovers accuracy.Positions within the protection window retain direct causal access to each masked token.Source: §4.3 · pp. 8–9
  • Replacing the stronger teacher with the student itself yields no improvement, supporting complementary roles for teacher guidance and TD.The ablation keeps the teacher’s predictions full-context and applies TD only to the student during training.Source: Table 2 · p. 8, §4.3 · p. 9

4.4 ABLATION AND ANALYSIS

  • At strictly matched update norms, TD improves held-out teacher-alignment KL by +0.027 (95% CI [0.022, 0.034], p < 0.001), indicating a directional advantage.The analysis attributes the gain to a better learning signal rather than a larger step size.Source: §4.4 · p. 9
  • At matched norms, masking only the student’s forward computation (−0.004) or only the training signal (−0.002) underperforms full-context OPD; applying both yields +0.027.The reported benefit requires the student view and training-signal computation to be masked consistently.Source: §4.4 · p. 9

5 RELATED WORK

Prior OPD work changes efficiency, supervision, trajectories, objectives, or teacher context; this paper instead studies how the student’s reasoning prefix shapes updates. Related context-masking methods differ in when and why they mask.

On-policy distillation

  • Unlike OPD methods that modify objectives, weighting, trajectories, or teacher context, TD changes the student’s reasoning-history access during optimization while keeping rollout and full-prefix teacher targets fixed.The cited approaches include rollout scheduling, token selection or reweighting, trajectory refinement, and asymmetric teacher context.Source: §5 · p. 9

Context masking and asymmetric distillation

  • Unlike prior context masking for multimodal shortcuts or corrupted student inputs, TD masks only when replaying a standard full-context on-policy trajectory for optimization.This preserves the rollout distribution and full-prefix teacher supervision.Source: §5 · p. 9

6 CONCLUSION

The paper identifies uneven update strengths under shared reasoning prefixes as a source of attenuated OPD learning signals. It proposes TD and reports gains across reasoning benchmarks and existing OPD variants.

6 CONCLUSION

  • TD randomly masks reasoning history during student updates to strengthen corrective responses and selectively introduce recovery signals for prefix-dependent predictions.The paper names the two attenuation forms confidence saturation and agreement saturation.Source: §6 · pp. 9–10
  • Experiments across teacher–student pairs and six reasoning benchmarks show gains over standard OPD, with further improvements to TRD and Fast-OPD.The conclusion also reports that preserving local prefix continuity supports stable training.Source: §6 · p. 10

A TRAINING DYNAMICS AND RESULTS

Appendix results report training dynamics, Qwen3-0.6B performance, and variation across three training seeds. The TD-OPD traces have higher student entropy, while seed averages remain above the reported OPD result.

Training dynamics

Qwen3-0.6B results

Robustness Across Training Seeds

  • OPD+TD training runs average 32.72 ± 0.23 across three seeds, with each run above the OPD result of 28.27.The individual averages are 32.78, 32.91, and 32.46.Source: A TRAINING DYNAMICS AND RESULTS · p. 13

B EXPERIMENTAL CONFIGURATIONS

The main TD-OPD experiments use a 1.7B student and remove training examples that overlap evaluation benchmarks. Trajectory Dropout adds little measured training time and leaves rollout generation untouched.

B EXPERIMENTAL CONFIGURATIONS

  • The 1.7B-student TD-OPD configuration removes overlapping training and evaluation examples during preprocessing to prevent potential data contamination.Other OPD experiments use the same configuration unless the main text states otherwise.Source: B EXPERIMENTAL CONFIGURATIONS · p. 13
  • TD adds no rollouts, teacher-scoring passes, or optimizer steps, and increases wall-clock time per step by 2.8% (273.3 s →280.9 s, +7.6 s).The increase is confined to student update and log-probability passes; rollout generation is unchanged.Source: B EXPERIMENTAL CONFIGURATIONS · p. 13

C EVALUATION DATASETS AND PROTOCOL

Evaluation covers six mathematical reasoning benchmarks and two out-of-domain benchmarks, with dataset versions and task-specific scoring procedures specified. Mathematical reasoning is summarized using average accuracy and pass-at-least-once measures.

C EVALUATION DATASETS AND PROTOCOL

  • The six mathematical benchmarks are AIME24, AIME25, AIME26, OlympiadBench, and the MathArena releases of AIME26 and HMMT.AIME versions each use 30 problems from that year's AIME I and II; the OlympiadBench snapshot contains 675 English, text-only competition-mathematics questions.Source: C EVALUATION DATASETS AND PROTOCOL · p. 13
  • The out-of-domain evaluations use MBPP+ with 378 EvalPlus tasks and GPQA Diamond with 198 graduate-level questions in biology, physics, and chemistry.The dataset table summarizes evaluation datasets and samples per problem, but the supplied passage gives no sample-count values.Source: C EVALUATION DATASETS AND PROTOCOL · p. 13, Table 5 · p. 14
  • Responses use zero-shot native-chat prompts and temperature 1.0, top-p 1.0 sampling; TD is disabled and full context is used at evaluation.The task-specific prompts request step-by-step reasoning for math and science, and Python solutions for MBPP+; generation is capped at 32,768 tokens within a 34,817-token context.Source: C EVALUATION DATASETS AND PROTOCOL · p. 14
  • Mathematical answers are graded by equivalence checks, GPQA Diamond by exact-match choice, and MBPP+ by execution-based test correctness.Unparseable answers and grading exceptions count as incorrect; length-truncated responses remain in the denominator, and intermediate reasoning is not separately scored.Source: C EVALUATION DATASETS AND PROTOCOL · pp. 14–15
  • Avg@n is mean accuracy over independently sampled responses, Pass@n is the fraction of problems solved at least once, and the six-benchmark Avg is their unweighted benchmark mean.MBPP+ and GPQA Diamond are reported separately and are excluded from that six-benchmark average.Source: C EVALUATION DATASETS AND PROTOCOL · p. 15

D SIGNAL ANALYSIS

Signal analysis finds that weakly updated tokens are prevalent across early, middle, and late OPD checkpoints. Span dropout substantially increases their share of training signal, with positive gains across both confidence- and agreement-saturated groups.

D SIGNAL ANALYSIS

  • Confidence-saturated tokens are defined by q_t ≥0.95 and |A_t^H| ≤0.05, while the disjoint groups are identified in the full view and held fixed across views.The analysis compares full-context scoring with 20% span dropout on the same next batch of 128 trajectories, using two independently seeded masks per trajectory.Source: D SIGNAL ANALYSIS · p. 15
  • At steps 10, 40, and 70, the two weak-signal groups comprise 74.85%, 73.32%, and 73.24% of response tokens but only 0.486%, 0.476%, and 0.479% of full-view signal mass.Under 20% span dropout, these same tokens' signal shares rise to 22.88%, 16.39%, and 17.70%, with absolute signal mass increasing 67.67, 45.36, and 49.78 times.Source: D SIGNAL ANALYSIS · p. 15
  • At the early, middle, and late checkpoints, 50.61%, 48.33%, and 47.29% of original weak tokens show positive signal gain under dropout.Positive gain occurs for 70.85%, 65.58%, and 66.41% of confidence-saturated tokens and 48.85%, 46.70%, and 45.58% of agreement-saturated tokens.Source: D SIGNAL ANALYSIS · p. 15

E NORM-CONTROLLED COMPARISON OF UPDATE DIRECTIONS

With update norms matched, OPD+TD yields greater held-out KL improvement than OPD, and the advantage tracks update direction rather than step size. The comparison also finds that the gain requires jointly changing the student's attention and loss views.

E NORM-CONTROLLED COMPARISON OF UPDATE DIRECTIONS

  • The paired comparison uses identical starting weights and optimizer settings, rescales each update to the OPD norm, and evaluates full-vocabulary forward KL on held-out GSM8K solutions.It uses 16 training batches and 128 held-out reference solutions with no training-data overlap.Source: E NORM-CONTROLLED COMPARISON OF UPDATE… · pp. 15–16
  • At matched update norms, OPD+TD improves held-out KL more than OPD by +0.027 (95% CI [0.022, 0.034]); the excess grows with radius, supporting a directional rather than step-size effect.Across radii 0.25×, 0.5×, and 1.0×, excess gains are +0.008, +0.016, and +0.027; the advantage is positive in all 16 training batches at every radius.Source: Table 6 · p. 16, E NORM-CONTROLLED COMPARISON OF UPDATE… · p. 16, Figure 8 · p. 16
  • The directional advantage remains with an unpreconditioned fixed-step-size gradient update (+0.147, p < 0.001), and rescaling leaves the OPD+TD gain nearly unchanged at +0.027.The result therefore does not depend on Adam-style preconditioning or a larger natural update norm.Source: E NORM-CONTROLLED COMPARISON OF UPDATE… · p. 16
  • At matched norm, changing attention masking alone (−0.004) or loss-view masking alone (−0.002) underperforms full-context OPD, while changing both jointly yields +0.027.The figure's schematic illustrates alignment with the validation-objective direction; its angles are illustrative, not measured.Source: E NORM-CONTROLLED COMPARISON OF UPDATE… · p. 17, Figure 9 · p. 17

F.1 SCOPE OF THE SAMPLED OPD GRADIENT

These appendices scope the sampled OPD analysis to fixed rollout prefixes, specify how trajectory spans are masked, and analyze how masking changes token-level corrections and updates. They also characterize clipping, local regression behavior, and the frequency and handling of sign reversals.

F.1 SCOPE OF THE SAMPLED OPD GRADIENT

  • The sampled-token analysis characterizes the fixed-batch update on prefixes generated by the starting parameters, not derivatives through a parameter-dependent rollout distribution.The identities are evaluated on a fixed rollout batch before optimizer transformations.Source: §F.1 · pp. 17–18

F.2 ATTENTION MASK AND SPAN SAMPLING

  • The attention mask removes selected direct attention edges while retaining all source positions and target losses; span sampling limits fragmentation and preserves visible gaps.The selected-source fraction does not itself determine the removed-edge fraction, which also depends on span locations.Source: §F.2 · p. 18

F.3 CONFIDENCE SATURATION AND THE CORRECTION CURVE

  • For fixed teacher probability, the sampled-token correction vanishes as student probability approaches 1, and within the stated range it strengthens as that probability falls.The parameter-level update also vanishes in the high-confidence limit under the stated bound; the full parameter update additionally depends on remaining logits and the network Jacobian.Source: §F.3 · pp. 18–19

F.4 AGREEMENT AND UPDATE DECOMPOSITION

  • Near-agreement positions can acquire a non-negligible masked-view update through either an induced advantage or a change in the score vector.The update difference is bounded by the induced advantage times the masked score norm plus the full-view advantage magnitude times the score-vector difference.Source: §F.4 · p. 19

F.5 ADVANTAGE CLIPPING AND THE LOCAL REGRESSION INTERPRETATION

  • At the start of an update, raw TD has a local squared-regression interpretation, extended to Huber regression when the detached advantage is clipped.The interpretation is local to a fixed rollout batch and attention mask; its effect on a full-context validation objective depends on the update direction and optimizer.Source: §F.5 · pp. 19–20

F.6 SIGN CONSISTENCY OF CONFIDENCE-SATURATED CORRECTIONS

  • Masked-view advantage sign reversals occur for approximately 8% of confidence-saturated tokens and account for less than 5% of total signal mass; a safeguard can revert to full-context OPD when reversal occurs.The stated percentages describe the diagnostic analysis; the safeguard restores both the student computation view and corresponding training signal.Source: §F.6 · p. 21

Related papers

Loading 2609.33455v2…