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

• URL: http://arxiv.org/abs/2605.10194v1
• Date: Mon, 11 May 2026 08:45:03 GMT
• Title: TRACE: Distilling Where It Matters via Token-Routed Self On-Policy Alignment
• Abstract summary: On-policy self-distillation (self-OPD) densifies reinforcement learning with verifiable rewards (RLVR) by letting a policy teach itself under privileged context.<n>We propose Token-Routed Alignment for Critical rEasoning (TRACE), which distills only on annotator-marked critical spans.<n>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.
• Score: 20.277178104190536
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation [59.43605629311855]
  We analyze the gradient of the per-token K2 estimator of reverse KL with respect to the student logits.<n>Our primary contribution is therefore a gradient-level characterization of reverse-KL OPD trained with the K2 estimator, with SuRe as one empirical instantiation.
  arXiv  Detail & Related papers  (2026-08-26T11:14:42Z)
• CRAFT: Counterfactual Credit Assignment from Free Sibling Rollouts for Self-Distilled Agentic Reinforcement Learning [4.663685189987781]
  Self-distilled agentic reinforcement learning augments trajectory-level reward with a token-level distillation loss.<n>The prevailing recipe gates this loss by a single scalar, the teacher-student log-probability gap.<n>We introduce CRAFT, a threepillar credit-assignment scheme that addresses both limitations.
  arXiv  Detail & Related papers  (2026-06-28T16:11:47Z)
• Learning from Own Solutions: Self-Conditioned Credit Assignment for Reinforcement Learning with Verifiable Rewards [49.1203423784326]
  We propose SC- GRPO (Self-Conditioned GRPO), which uses KL divergence mentioned before as a multiplicative weight on GRPO gradients.<n>Across five benchmarks spanning math, code, and agentic tasks, SC- GRPO consistently outperforms 8.1% over GRPO and 5.9% over DAPO with stronger OOD performance.
  arXiv  Detail & Related papers  (2026-06-17T08:26:02Z)
• Localizing Credit at the Divergence: Path-Conditioned Self-Distillation for LLM Reasoning [12.363923974156636]
  Reinforcement learning from verifiable rewards assigns a single scalar to each rollout.<n>On-policy self-distillation addresses this by letting the same model act as a teacher conditioned on privileged information.<n>We propose Hindsight Self-Distillation (HSD), which conditions the teacher on a successful peer rollout.
  arXiv  Detail & Related papers  (2026-06-14T03:37:27Z)
• Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation [81.10000755917712]
  On-policy distillation transfers reasoning capabilities by training a student model on its own generated trajectories using token-level feedback from a teacher.<n>As student-generated prefixes lengthen, the teacher's next-token distribution becomes less confident and less discriminative.<n>To mitigate SFD, we introduce textbfLookahead Group Reward (ours).
  arXiv  Detail & Related papers  (2026-05-29T04:39:20Z)
• Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation [28.572361799234784]
  We show that raw KL disagreement is a coarse proxy for learning value.<n>We formalize this local compatibility as token teachability.<n>We propose Teachability-Aware OPD, a lightweight token-position selection method.
  arXiv  Detail & Related papers  (2026-05-26T10:56:46Z)
• Decoupling KL and Trajectories: A Unified Perspective for SFT, DAgger, Offline RL, and OPD in LLM Distillation [15.564152482735652]
  We show that the prevailing paradigms, off-policy distillation and on-policy distillation, implicitly couple two choices: prefix source and token-level KL direction.<n>We propose KL mixing and an entropy-gated length curriculum.
  arXiv  Detail & Related papers  (2026-05-16T06:05:27Z)
• Self-Distilled Agentic Reinforcement Learning [65.24201057390938]
  Self-Distilled Agentic Reinforcement Learning maps detached token-level signals into a sigmoid gate.<n>SDAR consistently outperforms hybrid RL--OPSD baselines across model scales.
  arXiv  Detail & Related papers  (2026-05-14T17:51:26Z)
• Unleashing Implicit Rewards: Prefix-Value Learning for Distribution-Level Optimization [74.91418266859297]
  Implicit process reward models (PRMs) provide fine-grained reward signals along the reasoning process.<n>PRMs suffer from a train-inference mismatch: training only constrains a sequence-level aggregate, whereas inference requires token-level scores to reflect local step quality.<n>We propose a novel Implicit Prefix-Value Reward Model (IPVRM), which directly learns a prefix-conditioned value function estimating the probability of eventual correctness.<n>We also propose Distribution-Level RL (DistRL), which computes TD advantages for both sampled tokens and high-probability candidate tokens.
  arXiv  Detail & Related papers  (2026-04-14T18:19:54Z)
• Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes [31.95045602299568]
  On-policy distillation (OPD) is appealing for large language model (LLM) post-training because it evaluates teacher feedback on student-generated rollouts rather than fixed teacher traces.<n>We revisit OPD from the estimator and implementation sides.<n>We identify three failure modes of sampled-token OPD: an imbalanced one-token signal, unreliable teacher guidance on student-generated prefixes, and distortions caused by tokenizer or special-token mismatch.
  arXiv  Detail & Related papers  (2026-03-26T15:35:59Z)
• EMA Policy Gradient: Taming Reinforcement Learning for LLMs with EMA Anchor and Top-k KL [19.463033284177087]
  We propose two techniques to improve gradient policy algorithms for Large Language Models (LLMs)<n>First, we replace the fixed anchor policy during RL with an Exponential Moving Average (EMA)<n>Second, we introduce Top-k KL estimator, which allows for flexible between exact KL and sampled KL.
  arXiv  Detail & Related papers  (2026-02-04T10:50:17Z)
• On the Design of KL-Regularized Policy Gradient Algorithms for LLM Reasoning [59.11784194183928]
  Policy gradient algorithms have been successfully applied to enhance the reasoning capabilities of large language models (LLMs)<n>Regularized Policy Gradient (RPG) view shows that the widely-used $k_3$ penalty is exactly the unnormalized KL.<n>RPG-REINFORCE with RPG-Style Clip improves accuracy by up to $+6$ absolute percentage points over DAPO.
  arXiv  Detail & Related papers  (2025-05-23T06:01:21Z)
• Generalized Kullback-Leibler Divergence Loss [105.66549870868971]
  We prove that the Kullback-Leibler (KL) Divergence loss is equivalent to the Decoupled Kullback-Leibler (DKL) Divergence loss.<n>Thanks to the decoupled structure of DKL loss, we have identified two areas for improvement.
  arXiv  Detail & Related papers  (2025-03-11T04:43:33Z)
• Decoupled Kullback-Leibler Divergence Loss [90.54331083430597]
  We prove that the Kullback-Leibler (KL) Divergence loss is equivalent to the Decoupled Kullback-Leibler (DKL) Divergence loss. We introduce class-wise global information into KL/DKL to bias from individual samples. The proposed approach achieves new state-of-the-art adversarial robustness on the public leaderboard.
  arXiv  Detail & Related papers  (2023-05-23T11:17:45Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.