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

1. Problem Statement

The paper addresses the limitations of On-Policy Self-Distillation (Self-OPD) when applied to long-horizon Reinforcement Learning with Verifiable Rewards (RLVR), particularly in mathematical reasoning tasks. While Self-OPD methods (e.g., SDPO, SRPO) densify the sparse reward signal by using a policy snapshot as a teacher to guide the student, the authors identify a critical failure mode in long-horizon training:

• Granularity Mismatch: Existing methods apply Knowledge Distillation (KL) loss across all tokens in a response. However, analysis shows that ~97% of probability mass in correct rollouts is already shared between student and teacher, meaning most KL gradients are redundant.
• The "Uniform Distillation Tax": Applying KL to all tokens forces the model to spend gradients on positions that do not need correction. This leads to:
  • Privileged Information Leakage: The student absorbs correlations present in the privileged context (e.g., verified traces) that are absent at inference, causing Out-of-Distribution (OOD) degradation.
  • Entropy Rise and Reasoning Collapse: The authors observe a "three-symptom collapse" in long-horizon training: response lengths shorten drastically, per-token actor entropy rises (often >4× the baseline), and validation accuracy drops significantly.
• The Core Question: The right unit of distillation is not the entire response, but the small subset of decisive reasoning tokens where the student's distribution actually needs correction.

2. Methodology: TRACE

The authors propose TRACE (Token-Routed Alignment for Critical rEasoning), a method that routes distillation signals only to specific token spans rather than applying them uniformly.

Core Components

1. Token-Routed Self-OPD: Instead of a global KL loss, TRACE uses a privileged annotator ($\pi_A$) to identify specific spans in a rollout:

   • Key Spans ($K_y$): Decisive reasoning steps in correct rollouts (Reward=1).
   • Error Spans ($E_y$): Root-cause reasoning errors in incorrect rollouts (Reward=0).
   • Non-Spans ($N_y$): All remaining tokens.
   • Constraint: The total span coverage is capped at $\alpha = 0.25$ (25% of tokens) to limit privileged information exposure.

2. Coarse Diagnostic Routing:

   • The annotator emits a coarse type label (e.g., case_split, missed_case) but not the span text or the full rollout content.
   • The Teacher ($\pi_T$) receives the problem and this coarse label as a private diagnostic prefix. It computes logits causally on the student's generated prefix but does not see the specific span locations or content.
   • The Student ($\pi_S$) is trained using a routed loss where the KL term is applied only to the masked spans.

3. Asymmetric Divergence Routing:

   • Forward KL (FKL) on Key Spans: Applied to correct rollouts. FKL ($KL(\pi_T \parallel \pi_S)$) provides a "lift" to tokens the student under-allocates, regardless of the student's current probability mass.
   • Reverse KL (RKL) on Error Spans (Optional): Applied to incorrect rollouts. RKL ($KL(\pi_S \parallel \pi_T)$) suppresses tokens the student is over-confident about but the teacher disfavors.
   • No KL on Non-Spans: These tokens rely solely on the standard GRPO (Group Relative Policy Optimization) loss.

4. Finite Exposure Schedule:

   • The KL weight ($\lambda_k$) is annealed to zero after a short warm-up (e.g., 40 steps).
   • This ensures that the cumulative privileged-gradient exposure remains finite over the training horizon, preventing the long-term accumulation of OOD-degrading correlations.

3. Theoretical Analysis

The paper provides theoretical justification for the routing strategy:

• Asymmetric Pointwise Pressure (Corollary 2):
  • FKL is effective for under-allocated tokens (where student probability $p \ll$ teacher probability $q$) because the gradient pressure scales with the gap $(q-p)$ and does not vanish as $p \to 0$.
  • RKL is effective for confident-wrong tokens (where $p \gg q$) because the gradient scales with the student's over-confidence ($p$).
• Cumulative Exposure Bound (Proposition 3): By combining span masking (limiting bandwidth to $\alpha$) and KL decay (limiting duration), TRACE bounds the cumulative privileged-gradient exposure to $O(\alpha \Lambda^2)$, preventing the unbounded growth seen in persistent all-token KL baselines.
• Signal vs. Leakage: The method maximizes the alignment signal on key spans while minimizing the risk of leaking privileged information that causes spurious correlations.

4. Key Contributions

1. Diagnosis: Identification of a "uniform distillation tax" in all-token Self-OPD, linking rising entropy, epistemic suppression, and privileged-gradient leakage to a granularity mismatch.
2. Method & Theory: Introduction of TRACE, which routes $\{FKL, RKL, \emptyset\}$ based on token class (Key, Error, Non-span) and decays the privileged channel. The theory explains why FKL is dominant for key spans in strong-base regimes and why the optimal corner shifts to RKL for error spans in weaker-base regimes.
3. Empirical Evidence:
   • Performance: On Qwen3-8B, TRACE improves the 5-benchmark average by 2.76 percentage points over GRPO.
   • Stability: Unlike SDPO/SRPO, TRACE preserves the base model's OOD performance on GPQA-Diamond, where baselines degrade.
   • Self-Annotation: The method works with an online self-annotator (the student policy itself), recovering ~69% of the gain from a strong external API annotator, proving the gains come from the routing mechanism, not just external capability.
   • Cross-Scale Adaptation: On a weaker base (Qwen3-1.7B), the dominant routing corner naturally inverts from FKL-on-Key to RKL-on-Error, matching theoretical predictions.

5. Results

• Qwen3-8B (Strong Base): TRACE (FKL on Key Spans) achieves 81.51% average accuracy across 5 benchmarks (MATH-500, AIME 24/25, AMC 23, GPQA-Diamond), outperforming GRPO (78.75%) and the strongest prior Self-OPD baseline RLSD (78.99%). Crucially, it maintains the base model's GPQA-Diamond score (58.33% vs 58.27% base), whereas GRPO and all-token baselines drop significantly.
• Qwen3-1.7B (Weaker Base): TRACE (RKL on Error Spans) achieves 60.16% average, the only trained method to exceed the base model (58.77%). All-token baselines (SDPO/SRPO) collapse, falling 11–19 points below GRPO.
• Ablation Studies:
  • Mask Localization: Replacing TRACE's specific span masks with random 25% masks or all-token KL results in performance drops, confirming that which tokens are selected matters more than the quantity.
  • Annotator Quality: Using the actively training student as its own annotator yields +1.90 points, demonstrating that the method does not rely on importing external "super-human" reasoning capabilities.

6. Significance and Claims

The paper claims that TRACE offers a loss-geometric solution to the instability of long-horizon Self-OPD, rather than a domain-specific fix. Its significance lies in:

• Decoupling Distillation from Instability: By restricting the "privileged" signal to critical spans and decaying it, TRACE avoids the entropy rise and reasoning collapse observed in SDPO and SRPO.
• Adaptive Routing: It demonstrates that the optimal distillation direction (FKL vs. RKL) is not a fixed hyperparameter but depends on the base model's error profile (under-allocation vs. over-confidence).
• Practical Viability: It proves that high-quality distillation can be achieved without external teachers or expensive per-query rubrics, relying instead on a lightweight, coarse diagnostic annotation that can be self-generated.

The authors emphasize that TRACE is a "corner-routed framework" rather than a fixed recipe, where the specific routing action (FKL vs. RKL) should be chosen based on the residual error mode of the base model.