Escaping the KL Agreement Trap in On-Policy Distillation

• URL: http://arxiv.org/abs/2606.09471v1
• Date: Mon, 08 Jun 2026 13:28:54 GMT
• Title: Escaping the KL Agreement Trap in On-Policy Distillation
• Abstract summary: We propose an online OPD termination rule that detects persistent low-KL agreement with a dynamic training-adaptive threshold.<n>KAT improves avg@k accuracy by 2.66% and pass@k by 3.43% across four mathematical benchmarks, while reducing average rollout length by 59.73%.
• Score: 23.908556717663057
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: On-policy distillation (OPD) provides dense token-level supervision by asking a teacher to score student-generated rollouts. However, when the student drifts into an unrecoverable prefix, the teacher may locally agree with the degraded state, producing low reverse KL but little corrective training signal. We identify this persistent regime as a low-KL agreement trap. Further analyses show that tokens during and after such traps produce less useful supervision signals. We propose KAT (KL Agreement Trap Termination), an online OPD termination rule that detects persistent low-KL agreement with a dynamic training-adaptive threshold. By filtering weak supervision from degenerate agreement, KAT improves avg@k accuracy by 2.66% and pass@k by 3.43% across four mathematical benchmarks, while reducing average rollout length by 59.73%.

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)
• Beyond Teacher Likelihood: Group-Calibrated On-Policy Distillation for Long-Context Reasoning [19.29382518431932]
  Group-Calibrated On-Policy Distillation (GC-OPD) trains a student on its own responses using dense token-level guidance from a stronger teacher.<n>Group-Calibrated On-Policy Distillation (GC-OPD) normalizes verifier rewards and trajectory-level OPD scores within each rollout group.<n>Relative-advantage-based credit assignment (RACA) distributes this trajectory-level residual across tokens according to their relative OPD advantages.
  arXiv  Detail & Related papers  (2026-08-19T17:54:58Z)
• CADENCE: Closing the Reasoning Gap via Coverage-Adaptive On-Policy Distillation [2.587194279527956]
  CADENCE is a unified framework with a targeted fix for each failure mode.<n>DRIFT schedules a per-token convex mixture of forward-KL and reverse-KL surrogate objectives on student-sampled trajectories.<n>Six components extend it: (A) COVA, a coverage-adaptive schedule accelerating the forward-to-reverse transition; (B) FTB, a forking-token boost concentrating gradient at high-entropy positions via a globally-normalized entropy reference; (D) LAP, brevity-preferential correct-rollout reinforcement; (E) EMR, an entropy-gnostic calibration
  arXiv  Detail & Related papers  (2026-07-18T20:16:26Z)
• Reward-Gated On-Policy Distillation [58.187458236302575]
  On-policy distillation is a powerful way to transfer reasoning ability from a strong teacher to a smaller student.<n>We introduce RG-OPD: Reward-Gated On-Policy Distillation that uses verifier feedback to decide when teacher logits should be trusted.<n>Across reasoning and coding benchmarks, RG-OPD produces stronger distilled students, outperforming both vanilla reverse-KL distillation and the recent TSD-KD baseline.
  arXiv  Detail & Related papers  (2026-07-04T21:51:22Z)
• Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance [10.04930078540686]
  On-Policy Distillation improves large language model reasoning by training a student model on trajectories sampled from its own policy under teacher supervision.<n>We show that this "trajectory-sampled but token-learned" mechanism cannot reliably bridge student trajectories toward teacher trajectories.<n>We propose Trajectory-aware OPD, which uses near-future trajectory information to identify real divergent states and distribute guidance across multiple future tokens.
  arXiv  Detail & Related papers  (2026-05-29T19:32:07Z)
• 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)
• TRACE: Distilling Where It Matters via Token-Routed Self On-Policy Alignment [20.277178104190536]
  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.
  arXiv  Detail & Related papers  (2026-05-11T08:45:03Z)
• Asymmetric On-Policy Distillation: Bridging Exploitation and Imitation at the Token Level [23.59808224440603]
  On-policy distillation (OPD) trains a student on its own trajectories with token-level teacher feedback and often outperforms off-policy distillation and standard reinforcement learning.<n>We find that its standard advantage weighted policy gradient suffers from three structural weaknesses, including high variance updates, vanishing gradients in zero-advantage regions, and exploration bottlenecks when corrective signals are insufficient.<n>We propose Asymmetric On-Policy Distillation (AOPD), which replaces ineffective negative reinforcement with localized divergence minimization in non-positive advantage regions while preserving positive reinforcement learning.
  arXiv  Detail & Related papers  (2026-05-07T15:02:49Z)
• TCOD: Exploring Temporal Curriculum in On-Policy Distillation for Multi-turn Autonomous Agents [55.27396165691312]
  We identify a key limitation of vanilla OPD in multi-turn agent settings, which we term Trajectory-Level KL Instability.<n>We propose TCOD, a framework that controls the trajectory depth exposed to the student and progressively expands it with a curriculum schedule.<n> Experimental results across four student-teacher pairs show that TCOD mitigates KL escalation and enhances KL stability throughout training, improving agent performance by up to 18 points over vanilla OPD.
  arXiv  Detail & Related papers  (2026-04-27T03:38:27Z)
• Distillation Traps and Guards: A Calibration Knob for LLM Distillability [54.90137955363471]
  We propose a post-hoc calibration method that enables control over a teacher's distillability via reinforcement fine-tuning.<n>Our objective combines task utility, KL anchor, and across-tokenizer calibration reward.<n> Experiments show that students distilled from distillable teachers outperform SFT and KD baselines.
  arXiv  Detail & Related papers  (2026-04-21T01:22:35Z)
• Reinforcement-aware Knowledge Distillation for LLM Reasoning [63.53679456364683]
  Reinforcement learning (RL) post-training has recently driven gains in long chain-of-thought reasoning large language models (LLMs)<n>Most existing knowledge distillation methods are designed for supervised fine-tuning (SFT), relying on fixed teacher traces or teacher-student Kullback-Leibler (KL) divergence-based regularization.<n>We propose RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update.
  arXiv  Detail & Related papers  (2026-02-26T00:20:39Z)

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.