Decoupling KL and Trajectories: A Unified Perspective for SFT, DAgger, Offline RL, and OPD in LLM Distillation

• URL: http://arxiv.org/abs/2605.16826v1
• Date: Sat, 16 May 2026 06:05:27 GMT
• Title: Decoupling KL and Trajectories: A Unified Perspective for SFT, DAgger, Offline RL, and OPD in LLM Distillation
• Abstract summary: 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.
• Score: 15.564152482735652
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Knowledge distillation is central to LLM post-training, yet its design space remains poorly understood, especially alongside reinforcement learning (RL). We show that the prevailing paradigms, off-policy distillation and on-policy distillation (OPD), implicitly couple two orthogonal choices: prefix source and token-level KL direction. This follows from decomposing sequence-level KL over autoregressive response distributions: forward KL pairs teacher prefixes with token-level forward KL, and reverse KL pairs student prefixes with token-level reverse KL. We argue this coupling is not intrinsic: decoupling the two axes yields four valid objectives. We establish gradient-level identities showing forward KL gives SFT-style cross-entropy matching with teacher soft targets, whereas reverse KL gives an RL-style policy-gradient objective with a dense teacher-student log-ratio reward, connecting them to off-policy SFT, DAgger-style on-policy SFT, offline-RL-style distillation, and OPD. We conduct an extensive controlled study on math reasoning, evaluating the four objectives both as standalone methods and as initializations for subsequent RL. The results reveal three tradeoffs: KL direction induces an accuracy-entropy tradeoff, prefix source a quality-compute tradeoff, and training length an accuracy-stability tradeoff. Motivated by these findings, we propose KL mixing and an entropy-gated length curriculum. KL mixing shows long-sequence distillation requires substantial forward-KL weight to prevent entropy collapse and length inflation without sacrificing accuracy. The entropy-gated length curriculum improves Avg@k and Pass@k by 3.6 and up to 5.8 points, and cuts average response length by roughly 3x versus fixed long-horizon training. Our results provide a framework and practical methods for designing reasoning distillation objectives that balance accuracy, diversity, compute, and RL behavior.

Related papers

• Unifying ICL, SFT, KL-Regularized RL Through a Bayesian Lens [0.0]
  This note develops a Bayesian perspective that puts these procedures on the same footing.<n>Part I formalizes few-shot ICL and SFT as amortized and-weights projections onto the Bayes posterior predictive.<n>Part II-IV show that KL-regularized RLHF/RLVR, reward-weighted SFT, reward-weighted ICL, and advantage-weighted SFT are all instances of forward-KL projection onto posteriors induced by rewards or advantages.
  arXiv  Detail & Related papers  (2026-09-04T13:06:29Z)
• AsyncOPD: How Stale Can On-Policy Distillation Be? [23.273855195545394]
  On-policy distillation (OPD) trains a student on its own rollouts guided by teacher feedback.<n>Asynchronous training pipelines can alleviate this bottleneck by decoupling rollout generation from learner updates.<n>We present the first systematic study of staleness in asynchronous OPD.
  arXiv  Detail & Related papers  (2026-06-23T04:50:49Z)
• DISA: Offline Importance Sampling for Distribution-Matching LLM-RL [56.9445657766829]
  We introduce DISA, which moves this calibration problem outside the RL loop.<n> DISA draws proposal trajectories offline, estimates the partition function via importance sampling, and freezes the resulting partition-function estimate.<n>On two open-weight backbones across six math and three code benchmarks, DISA matches or exceeds the online-coupled distribution-matching baseline Flow.
  arXiv  Detail & Related papers  (2026-05-17T07:14:44Z)
• 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)
• The Extrapolation Cliff in On-Policy Distillation of Near-Deterministic Structured Outputs [52.709361620508595]
  ListOPD brings a student to in-domain with an 8B-SFT baseline at one-fifth the parameters.<n>On Amazon Fashion, three pre-registered tests--a fine-grid cliff interval test, and a small-clip cross-prediction--fall within their locked prediction windows--with the small-clip value matching the closed-form prediction below grid resolution.
  arXiv  Detail & Related papers  (2026-05-09T06:48:00Z)
• Entropy-Aware On-Policy Distillation of Language Models [36.60992451188347]
  We introduce Entropy-Aware On-Policy Distillation.<n>Our key idea is augmenting the standard reverse KL objective with forward KL when teacher entropy is high.<n>It balances mode-seeking precision with mode-covering without sacrificing on-policy training efficiency.
  arXiv  Detail & Related papers  (2026-03-07T07:26:18Z)
• 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)
• 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)
• Rethinking Kullback-Leibler Divergence in Knowledge Distillation for Large Language Models [18.870276152694245]
  Kullback-Leiber divergence has been widely used in Knowledge Distillation (KD) to compress Large Language Models (LLMs)<n>Contrary to prior assertions, reverse Kullback-Leibler (RKL) divergence is mode-seeking and thus preferable over the mean-seeking forward Kullback-Leibler (FKL) divergence.<n>We propose a simple yet effective Adaptive Kullback-Leiber (AKL) divergence method, which adaptively allocates weights to combine FKL and RKL.
  arXiv  Detail & Related papers  (2024-04-03T11:40:17Z)
• 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)
• RL with KL penalties is better viewed as Bayesian inference [4.473139775790299]
  We analyze challenges associated with treating a language model as anReinforcement Learning policy. We show how avoiding those challenges requires moving beyond the RL paradigm.
  arXiv  Detail & Related papers  (2022-05-23T12:47:13Z)

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.