PaperlayerView exact revision on arXiv

Source-linked AI summary

TRACE: Distilling Where It Matters via Token-Routed Self On-Policy Alignment

Jiaxuan Wang, Xuan Ouyang, Zhiyu Chen, Yulan Hu, Zheng Pan, Xin Li, Lan-Zhe Guo

arXiv:2605.10194v1cs.AIcs.LG

TL;DR

All-token self-OPD can waste privileged teacher gradients on redundant positions and destabilize long-horizon reasoning. TRACE routes FKL, optional RKL, or no-KL to annotator-marked spans while decaying the KL channel back to GRPO. It improves the five-benchmark Qwen3-8B average by 2.76 percentage points, preserves GPQA-Diamond base performance, and shifts from FKL to RKL as the base model weakens.

  • Problem

    All-token self-OPD spends KL on mostly redundant positions and is associated with entropy rise, shortened reasoning, and out-of-distribution degradation.

  • Method

    TRACE routes FKL to key spans, optional RKL to error spans, and GRPO elsewhere, with span masking and KL decay limiting privileged-gradient exposure.

  • Results

    TRACE improves the Qwen3-8B five-benchmark average over GRPO by 2.76 percentage points and preserves the base score on GPQA-Diamond.

  • Takeaways & Limitations

    The best routed action depends on the base model’s residual error mode: FKL on key spans for Qwen3-8B and RKL on error spans for Qwen3-1.7B.

  • Takeaways & Limitations

    The reported theoretical exposure bound assumes bounded masked privileged variance and does not directly characterize empirical baselines that combine clipping, decay, or routing.

Abstract

from arXiv · show

On-policy self-distillation (self-OPD) densifies reinforcement learning with verifiable rewards (RLVR) by letting a policy teach itself under privileged context. We find that when this guidance spans the full response, all-token KL spends gradients on mostly redundant positions and amplifies privileged-information leakage, causing entropy rise, shortened reasoning, and out-of-distribution degradation in long-horizon math training. We propose Token-Routed Alignment for Critical rEasoning (TRACE), which distills only on annotator-marked critical spans: forward KL on key spans of correct rollouts, optional reverse KL on localized error spans, and GRPO on all remaining tokens, with the KL channel annealed away after a short warm-up. Our analysis explains TRACE through two effects: forward KL provides non-vanishing lift to teacher-supported tokens that the student under-allocates, while span masking and decay keep cumulative privileged-gradient exposure finite. On four held-out math benchmarks plus GPQA-Diamond, TRACE improves over GRPO by 2.76 percentage points on average and preserves the Qwen3-8B base OOD score on GPQA-Diamond, where GRPO and all-token self-OPD baselines degrade. Gains persist under online self-annotation (+1.90 percentage points, about 69% of the strong-API gain), reducing the concern that TRACE merely imports external annotator capability. Across scales, the best routed action is base-dependent: on Qwen3-8B it is forward KL on key spans, while on Qwen3-1.7B it shifts to reverse KL on error spans.

1 Introduction

TRACE addresses long-horizon instability in all-token self-OPD by routing privileged teacher guidance to critical spans while leaving most tokens GRPO-driven. It combines span-specific KL actions with finite coverage and decay to limit leakage and preserve reasoning stability.

  • Motivation: All-token self-OPD can destabilize long-horizon math training, with entropy rising above 4× GRPO and validation accuracy dropping 18–19 percentage points.The observed collapse also includes shortened responses, and SRPO’s EMA teacher appears to amplify actor entropy.
  • Motivation: ∼97% of probability mass lies on shared student-teacher tokens, motivating distillation focused on spans where the student’s next-token distribution needs correction.The paper frames all-token KL as spending gradients on mostly redundant positions and risking privileged-information leakage.
  • Method: TRACE routes FKL to key spans, optional RKL to error spans, and GRPO to non-span tokens.The span mask caps KL support at α = 0.25, while the KL channel decays to zero within 40 steps.
  • Mechanism: TRACE’s default FKL-on-key-spans choice provides non-vanishing lift to teacher-supported tokens while masking and decay keep cumulative privileged-gradient exposure finite.RKL can vanish when the student assigns little mass, whereas FKL remains effective in that under-allocation regime.
  • Evidence: TRACE improves the Qwen3-8B five-benchmark average by 2.76 percentage points, preserves GPQA-Diamond out-of-distribution performance, and retains a 1.90-point gain with online self-annotation.The reported results also show a shift from FKL on key spans to RKL on error spans for the weaker Qwen3-1.7B base.

2 Related Work

Prior work uses teacher logits to densify sequence-level training, while self-OPD removes the need for a separate teacher but is associated with long-horizon collapse and privileged-information leakage. TRACE is positioned alongside routing and advantage-reweighting mitigations.

  • Prior approaches: Knowledge distillation and on-policy variants densify sequence-level training with teacher logits from external teachers or privileged-context self-teachers.Self-OPD removes the separate-teacher requirement by using privileged-context self-teachers.
  • Known limitations: Related work attributes self-OPD’s long-horizon collapse to epistemic suppression and privileged-information leakage, alongside sample-level routing and advantage reweighting as mitigations.

3 TRACE: Corner-Routed Span Distillation

TRACE separates student, privileged teacher, and span annotator roles, then selects discrete FKL, RKL, or no-KL actions by token class. The KL channel is short-lived so optimization ultimately returns to GRPO.

  • Pipeline: TRACE uses a student, a synchronized privileged-context teacher, and an annotator that emits sparse masks and coarse labels without logits or gradients.The teacher receives the type label as a private diagnostic prefix and computes logits causally on the rollout prefix.
  • Routing rationale: TRACE favors discrete routing because endpoint actions can dominate interior FKL/RKL mixtures under the paper’s endpoint-alignment and density-floor assumptions.
  • Span annotation: Annotated spans are selected from correct or incorrect rollouts according to verifier outcome, projected to token masks, and capped at 25% of rollout tokens.
  • Routing: The default routing applies FKL to key spans, optional RKL to error spans, and no KL to non-spans, with GRPO continuing on non-span tokens.The routed actions are chosen by token class rather than mixed continuously within a class.
  • Decay: The KL coefficient is held during warm-up, annealed linearly to zero, and followed by pure GRPO after the decay window.The implementation skips the teacher forward pass after decay, closing the privileged channel; the teacher syncs every 10 steps during the active phase.

4 Theoretical Analysis

TRACE’s analysis separates alignment benefit from privileged-gradient risk: FKL is effective for under-allocated teacher-supported tokens, RKL for confident teacher-disfavored tokens, while masking and finite decay bound exposure. A held-out lift proxy supports the predicted FKL behavior.

  • Pointwise asymmetry: FKL gives under-allocated teacher-supported tokens a mass-independent lift, whereas RKL’s lift vanishes as student probability approaches zero.This asymmetry motivates FKL for key spans in strong-base regimes.
  • Pointwise asymmetry: RKL gives confident-wrong student-supported tokens stronger down-pressure scaled by student over-confidence and teacher disagreement than FKL.This motivates RKL on error spans when local over-confidence is the dominant residual error mode.
  • Risk control: TRACE’s span coverage and finite-decay schedule keep cumulative privileged-gradient exposure finite as training horizon grows.Under the stated assumptions, exposure scales as O(αΛ2) rather than exhibiting the unbounded O(K) growth of persistent all-token KL.
  • Positive signal: The key-span signal is positive only under sufficient annotator precision, making annotation quality a condition for the default FKL benefit.The paper states that the positive lower bound requires annotator precision to exceed a threshold q∗K.
  • Empirical proxy: +0.145 nats, or +168% versus GRPO, is the reported TRACE-FKL key-token probability lift on teacher-supported spans.TRACE also reports an 8.5× credit concentration ratio versus 1.0–1.1× for all-token baselines.

5 Experiments

TRACE improves held-out math performance across scales while preserving the Qwen3-8B base score on GPQA-Diamond. Its best routing action depends on base capability, and gains persist with online self-annotation.

  • Main results: TRACE-FKL improves Qwen3-8B’s 5-benchmark AVG from GRPO’s 78.75 to 81.51 (+2.76 pp).
  • Main results: TRACE-FKL is best on MATH, both AIME splits, GPQA-Diamond, and the 5-benchmark AVG in the Qwen3-8B block.
  • Main results: TRACE matches the Qwen3-8B base score on GPQA-Diamond within evaluation resolution, while GRPO and all-token self-OPD baselines degrade.
  • Evaluation caveat: Per-benchmark gaps are treated as mean differences rather than standalone significance claims because the smallest held-out sets contain 30 AIME problems.
  • Cross-scale results: TRACE-RKL reaches 60.16 AVG on Qwen3-1.7B, exceeding the base and surpassing the FKL corner by 1.70 pp.
  • Ablations: All-token KL falls below GRPO, random sparsity trails TRACE, and inverted spans return to GRPO, showing that signal location matters more than sparsity alone.
  • Annotator ablation: Online self-annotation improves performance by +1.90 pp, approximately 69% of the strong-API gain, without an external supervisor.

6 Discussion and Conclusion

TRACE routes KL by critical span and error mode, then returns optimization to GRPO, improving stability across model scales while changing its dominant action with the base model. Its design combines FKL/RKL asymmetry with finite privileged-gradient exposure, though the natural-gradient mass result does not describe standard Euclidean SGD.

  • Method: TRACE routes FKL, RKL, or no KL over annotator-marked spans, then decays the privileged channel to return non-span optimization to GRPO.The default uses FKL on key spans; RKL can target error spans.
  • Results: TRACE improves GRPO by +2.76 pp, preserves GPQA-Diamond OOD accuracy, and shifts to RKL on Qwen3-1.7B as predicted.The paper frames these results as support for choosing the routed corner by the base model’s dominant residual error mode.
  • Theory: The design limits privileged-gradient exposure through span masking and KL decay, keeping the cumulative exposure finite under the stated assumptions.The bound remains finite regardless of training horizon, unlike persistent all-token KL under matched assumptions.
  • Cross-scale findings: On Qwen3-8B, TRACE-FKL on key spans is dominant, while on Qwen3-1.7B the dominant corner shifts to RKL on error spans.Both routed corners exceed GRPO at both scales, while all-token SDPO/SRPO peak early and collapse.
  • Theory caveat: The natural-gradient analysis characterizes monotone movement toward teacher mass, but standard Euclidean SGD does not generally preserve that monotonicity.The paper treats the natural-gradient result as complementary; pointwise logit pressure holds universally under Euclidean SGD.

B.6 Proof of Proposition 3

The proof bounds cumulative privileged-gradient exposure by combining span coverage, bounded masked variance, and KL duration. This yields a finite exposure rate proportional to the span fraction and variance bound.

  • Per-token bound: The exposure bound uses a score-operator inequality to control cross-vocabulary covariance terms in the per-token privileged-gradient deviation.A bound based only on individual score norms would not control these covariance terms.
  • Span restriction: TRACE restricts privileged-gradient contributions to annotator-selected span positions, while non-span tokens receive no privileged-context teacher signal.The proof retains the joint expectation of coverage and span-positioned variance because selected spans may differ from rollout-average variance.
  • Exposure bound: With coverage at most α = 0.25 and uniformly bounded masked privileged variance, cumulative exposure satisfies E_K = O(αΛ^2) regardless of training horizon.The resulting rate remains finite even as the number of training steps grows.
  • Duration: Finite exposure additionally requires the KL weights to become zero after a finite decay window.The schedule condition is λ_k = 0 for k > t_start + T_decay.
  • Second-moment rationale: The proof bounds the cumulative second moment rather than the first moment because the privileged-gradient deviation has zero conditional expectation, leaving variance as the accumulating quantity.This connects the bound to SGD noise and gradient drift.

B.7 Proof of Proposition 4

Proposition 4 gives conditional per-token signal bounds for routed KL actions, then combines them into a three-class allocation under explicit alignment, density, and leakage assumptions.

  • Signal bounds: Under binary span weights, expected selected-span KL gradients are lower-bounded by precision-weighted alignment margins minus false-selection penalties.The bounds apply separately to error and key spans and require positive active-class brackets.
  • Interpretation: The oracle direction is an unobservable verifier-grounded gradient, so Proposition 4 is a conditional signal-lower-bound rather than a trained-network optimality theorem.Its role is to complement the exposure bound in describing the useful-guidance versus leakage trade-off.
  • Three-class allocation: The conditional corner solution assigns RKL to error spans, FKL to key spans, and no KL to non-spans when κN < κ < min{κE, κK}.The interval is non-empty only under the corresponding separation condition, and the figure presents this as idealized utility geometry.
  • Key-span allocation: Forward KL is favored on key spans because its teacher-supported-token pressure remains non-vanishing when student probability is low, unlike reverse KL’s student-mass scaling.This motivates the key-span endpoint in the under-allocation regime.
  • Dead-zone behavior: TRACE preserves a nonzero KL signal on all-correct rollout groups when GRPO advantages collapse to zero, provided student and teacher distributions differ on selected key spans.The gradient remains nonzero during the KL-active phase under generic distributional disagreement.

C.1 Training Hyperparameters

The training setup uses standardized Qwen3-8B math RL settings and an annotator-driven TRACE loop that projects sparse spans to tokens, applies routed KL during warm-up, and returns to pure GRPO after decay.

  • Common settings: TRACE training uses Qwen3-8B, up to 30K OpenThoughts-114k pairs, 300 steps, learning rate 1×10^-5, and 10-step linear warm-up.Common settings include H100 GPUs, AdamW, batch size 32, eight rollouts per problem, and maximum response lengths of 32768 for Think training and 8192 for NoThink ablations.
  • KL schedule: The KL channel is annealed during a decay window, after which TRACE skips the teacher forward and reduces to pure GRPO.GRPO is smoothly reintroduced on span tokens during decay through ρk.
  • Training loop: The update loop samples rollouts, computes GRPO advantages, obtains annotator spans and type labels, forwards the privileged teacher, and combines routed KL with GRPO.Teacher context contains the original problem and a coarse diagnostic type, not rollout text, span locations, or span content.
  • Span annotation: The annotator marks error spans on incorrect rollouts and key spans on correct rollouts, with token masks capped at α = 0.25.Character-level annotations are projected to tokens by interval intersection and top-weight retention.
  • Evaluation protocol: The evaluation protocol uses thinking-mode decoding with T=0.6, p=0.95, k=20, and max_new_tokens = 38912 for main results and OOD evaluation.A nonthinking protocol with max_new_tokens = 8192 is reserved for the NoThink-eval ablation.

D.4 Evaluation Protocol Sanity Checks

Sanity checks show that the reported Qwen3-8B headline results are robust to grader choice and that uncertainty is quantified with problem-level bootstrap intervals.

  • Grader robustness: Across all five evaluation cells, strict-versus-ensemble grading differs by at most 0.03 pp, below the bootstrap confidence-interval width.The ensemble combines multiple math graders, supporting the use of the strict-grader headline numbers.
  • Uncertainty estimates: Table 7 reports bootstrap 95% confidence intervals over problem IDs, with AVG defined as the unweighted mean over five benchmark means.The intervals provide cell-wise uncertainty for the Table 2 results.

D.5 SDPO/SRPO Re-Implementation Sweep

The sweep re-implements SDPO and SRPO for long-horizon Qwen3-8B math RL, finding rapid collapse symptoms and interpreting these results as setting-specific rather than claims about the original systems.

  • Re-implementation setup: SDPO and SRPO are evaluated through re-implementations because their official training code is not publicly released.The configurations and selected values are reported in the sweep tables.
  • Transfer caveat: The reported SDPO and SRPO results are faithful re-implementations under long-horizon math RLVR, not claims about the original systems’ exact performance.The authors do not isolate a single cause for differences from original reports and note reproduction across Qwen3-8B and Qwen3-1.7B.
  • Observed failure pattern: SDPO and SRPO show length collapse, entropy rise, and validation collapse during long-horizon math training.The first two symptoms appear within 150–200 steps, while validation collapse is the headline observation of Section 5.2.
  • Observed failure pattern: SDPO collapses from 2027 to 1042 tokens and SRPO from 1873 to 759, reaching approximately one-quarter of their starting lengths.The minimum lengths are 528 for SDPO and 460 for SRPO.

D.8 NoThink-Eval Robustness Under Asymmetric Training

Under asymmetric training, the NoThink evaluation removes thinking-token padding and shows that TRACE remains the strongest method by average score, exceeding GRPO by 2.22 percentage points.

  • NoThink evaluation: +2.22 pp: TRACE with FKL on Ky exceeds GRPO on the NoThink evaluation average.The evaluation uses Thinking off, max_new_tokens = 8192, and avg@8.
  • NoThink evaluation: SDPO and SRPO collapse further when thinking-token padding is removed.This result rules out a padding-only explanation for the asymmetric variant’s behavior.

D.9 Negative Results

Two negative results show that TRACE’s span-localized routing matters: generic critique alternatives are costlier and weaker, while all-token FKL with decay still harms OOD performance.

  • Alternative supervision: Generic rubrics and free-form LLM critique underperform span-localized routing while incurring substantially higher annotator cost.Per-query rubric generation is reported as G× more expensive than per-rollout type-label annotation.
  • All-token FKL: By step 100, all-token FKL with the same decay schedule degrades GPQA-Diamond OOD performance relative to span-localized FKL.The matched comparison indicates that decay alone is insufficient; span localization is necessary in this setup.
Loading 2605.10194v1…