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

Anhao Zhao; Haoran Xin; Yingqi Fan; Junlong Tong; Wenjie Li; Xiaoyu Shen

Source-linked AI summary

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

Anhao Zhao, Haoran Xin, Yingqi Fan, Junlong Tong, Wenjie Li, Xiaoyu Shen

arXiv:2605.16826v1May 16, 2026cs.LGcs.AIcs.CL

TL;DR

The paper addresses limited understanding of how prefix source and KL direction interact in LLM distillation. It decouples these axes into four objectives, derives their learning-regime interpretations, and evaluates them with RL follow-up; the resulting methods improve accuracy while controlling entropy, length, and stability.

Problem
Distillation design and its interaction with RL remain poorly understood despite the prevalence of off-policy and on-policy post-training pipelines.
Method
The paper decouples prefix source from token-level KL direction, derives four objectives and their gradient-level learning-regime interpretations, then evaluates them in standalone and distillation-then-RL settings.
Results
The entropy-gated length curriculum improves Avg@k by 3.6 points, raises Pass@k by up to 5.8 points, and cuts average response length by roughly 3× versus fixed long-horizon training.
Takeaways & Limitations
Reasoning distillation should balance accuracy, diversity, compute, and continued trainability rather than optimize standalone accuracy alone.
Takeaways & Limitations
The study focuses on mathematical reasoning and uses GRPO with an accuracy-only outcome reward, leaving other domains and RL algorithms for future investigation.

Abstract

from arXiv · show

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.

1 Introduction

The paper decouples prefix source from token-level KL direction, exposing four distillation objectives and three tradeoffs relevant to standalone reasoning distillation and subsequent RL.

1 Introduction: Decoupling prefix source from KL direction yields four objectives corresponding to classical training regimes.The conventional off-policy and on-policy pairings are therefore not intrinsic at the token level.
1 Introduction: Forward KL corresponds to soft-target SFT, while reverse KL corresponds to an RL-style policy-gradient objective with dense teacher–student log-ratio rewards.These gradient identities connect the four objectives to off-policy SFT, DAgger-style on-policy SFT, offline-RL-style distillation, and OPD.
1 Introduction: The study identifies accuracy–entropy, quality–compute, and accuracy–stability tradeoffs across distillation and downstream RL.Reverse KL can improve standalone accuracy while reducing entropy; student prefixes can improve matched-step performance, while teacher prefixes can be more compute-effective under matched FLOPs.
1 Introduction: KL mixing prevents entropy collapse and length inflation in long-sequence distillation with little or no accuracy loss when forward-KL weight is high.High reverse-KL mixtures remain unstable.
1 Introduction: The entropy-gated length curriculum improves Avg@k by 3.6 points, raises Pass@k by up to 5.8 points, and reduces average response length by roughly 3× versus fixed 4096-token training.It increases the training horizon only while held-out predictive entropy remains above a stability threshold.

2 Disentangling KL Direction and Prefix Source

Autoregressive KL decomposition explains conventional pairings between prefix source and KL direction, while token-level analysis shows the axes can be independently combined into four learning regimes.

2 Disentangling KL Direction and Prefix Source: Forward sequence KL pairs teacher-induced prefixes with token-level forward KL, whereas reverse sequence KL pairs student-induced prefixes with token-level reverse KL.The distinction follows from applying autoregressive factorization to sequence-level objectives.
2 Disentangling KL Direction and Prefix Source: Independently combining the two axes yields four token-level objectives, including two off-diagonal pairings absent from prevailing practice.The off-diagonal objectives use teacher prefixes with reverse KL or student prefixes with forward KL.
2 Disentangling KL Direction and Prefix Source: Prefix source determines the policy regime, while KL direction determines whether learning uses soft-target cross-entropy or an RL-style log-ratio reward.Teacher prefixes are off-policy and student prefixes are on-policy for the student.
2 Disentangling KL Direction and Prefix Source: The four objectives correspond to off-policy SFT, DAgger-style on-policy SFT, offline-RL-style distillation, and OPD.Teacher-prefix forward KL is off-policy SFT; student-prefix forward KL is DAgger-style SFT; teacher-prefix reverse KL is offline-RL-style; student-prefix reverse KL is OPD.
2 Disentangling KL Direction and Prefix Source: The evaluation tests matched-step standalone distillation and whether distilled checkpoints provide effective initializations for subsequent RL.The study compares resulting students in both regimes.

3 Experimental Setup

The experiments distill Qwen3 teachers into a Qwen3-0.6B student, compare short and long horizons, and evaluate standalone and RL-follow-up behavior on mathematical reasoning benchmarks.

3 Experimental Setup: Qwen3-4B and Qwen3-8B teachers are distilled into Qwen3-0.6B-Base within the same model family.Using one model family avoids cross-tokenizer artifacts in token-level KL computation.
3 Experimental Setup: Standalone distillation compares 128-token and 4096-token horizons using bf16 training, learning rate 5 × 10^-7, batch size 32, and 1000 steps.All objectives use matched optimizer and training settings in the standalone comparison.
3 Experimental Setup: The RL follow-up uses GRPO with an accuracy-based outcome reward, group size 8, batch size 32, temperature 1.0, top-p 0.95, and 1000 steps.Maximum decoding length is 4096 tokens.
3 Experimental Setup: Evaluation uses DeepScaleR training data and AIME24, AMC23, MATH500, and GSM8K, reporting Avg@N, Pass@N, and average response length.N is 5 for AIME24 and AMC23, and 3 for MATH500 and GSM8K.
3 Experimental Setup: Training dynamics track task accuracy, mean per-token predictive entropy, and average response length as diagnostics of performance, exploration capacity, and generation behavior.A fused full-vocabulary KL kernel makes exact full-vocabulary KL feasible without materializing vocabulary-sized intermediates.

4 Empirical Analysis of the Four Decoupled Objectives

The four objectives expose distinct tradeoffs: reverse KL improves standalone accuracy but can collapse entropy and inflate lengths, student prefixes improve quality but cost more compute, and long training helps accuracy while reducing stability.

4.1 KL Direction: Forward versus Reverse KL: +2.45 Avg@k points: reverse KL improves standalone accuracy on average, but 4096-token runs often approach entropy collapse and the generation limit.The average gain is larger at 128 tokens (+3.68) than at 4096 tokens (+1.21).
4.1 KL Direction: Forward versus Reverse KL: Forward-KL warm starts are more reliable for subsequent RL, while reverse-KL warm starts can begin more accurately but degrade or plateau as low entropy constrains exploration.After 4096-token distillation, one reverse-KL run fell from roughly 45% to 36% MATH500 accuracy, while forward KL improved from around 40% to about 45%.
4.2 Prefix Source: Teacher versus Student Prefixes: Student prefixes improve quality under matched steps, whereas teacher prefixes can be more compute-efficient under matched FLOPs by reusing cached logits.In a 128-token reverse-KL comparison, teacher prefixes reached roughly 38–40% MATH500 accuracy within 15–20k cumulative TFLOPs, while student prefixes required substantially more compute.
4.2 Prefix Source: Teacher versus Student Prefixes: After 4096-token distillation, student prefixes outperform teacher prefixes by 3.55 Avg@k and 2.95 Pass@k points under matched teacher scale and KL direction.Across matched comparisons, student prefixes improve Avg@k by +1.80 points and Pass@k by +2.11 points on average.

5 Methods: Balancing Accuracy and Entropy

The paper frames distillation as a balance among KL direction, training length, accuracy, entropy, and response-length stability. KL mixing and an entropy-gated curriculum target the instability of long-sequence reverse-KL training.

Tradeoffs: Reverse KL improves standalone Avg@k but sharpens the distribution, reducing diversity and weakening Pass@k, whereas forward KL better preserves entropy.Longer training improves accuracy but can drive entropy near zero and inflate response length under reverse KL.
KL Mixing: KL mixing interpolates between forward- and reverse-KL endpoints, with reverse-KL weight denoted by λ in the token-level mixture.The method evaluates the mixture in student-prefix, 4096-token training with Qwen3-4B as teacher.
KL Mixing: Forward-heavy KL mixing preserves most reverse-KL accuracy while increasing entropy and stabilizing response length.Pure reverse KL gives stronger accuracy but lowers entropy and increases length; pure forward KL is stable but weakest in accuracy.
Entropy-Gated Length Curriculum: The entropy-gated curriculum increases the training horizon only while held-out predictive entropy remains at least H_min.If entropy falls below the threshold, training stops increasing length and may terminate or continue at the last stable horizon.
Entropy-Gated Length Curriculum: 40.0 Avg@k and 54.1 Pass@k, versus 36.4 and 48.3, while average length falls from 8185 to 2688 tokens with Qwen3-4B.With Qwen3-8B, Avg@k rises from 36.6 to 40.2 and Pass@k from 49.2 to 51.8, while length falls from 8190 to 2793.

6 Conclusion

The paper separates prefix source from KL direction to unify four distillation objectives and studies their tradeoffs in accuracy, diversity, compute, and continued RL training. It proposes KL mixing and entropy-gated length control to improve stability while retaining distillation performance.

Conclusion: Separating prefix source from KL direction yields four objectives connected to off-policy SFT, DAgger-style on-policy SFT, offline-RL distillation, and OPD.The framework treats reasoning distillation as a controlled tradeoff among accuracy, diversity, compute, and continued trainability.
Conclusion: The objectives expose accuracy–entropy, quality–compute, and accuracy–stability tradeoffs across standalone distillation and downstream RL.These tradeoffs motivate methods that balance distillation performance with diversity and generation stability.
Conclusion: Off-policy and on-policy distillation remain prevalent post-training paradigms, but their design space and interaction with RL are poorly understood.The conclusion follows a motivation centered on the relationship between knowledge distillation and reinforcement learning.

B Proof of Proposition 1

The proof identifies the gradients induced by the two token-level KL directions. Forward KL produces soft-target cross-entropy matching, while reverse KL produces a policy-gradient estimator with a dense teacher–student log-ratio reward.

Forward KL: Forward-KL minimization has gradient equal to the gradient of cross-entropy H(p_T, q_θ) with teacher-distributed token targets.The teacher distribution is fixed, so differentiation acts through the student log-probability.
Reverse KL: Reverse-KL minimization yields the REINFORCE ascent estimator with dense reward r(s_t, y) = log p_T(y | s_t) − log q_θ(y | s_t).The reward is the teacher–student log ratio evaluated on sampled tokens.

C Derivation of Sequence-level KL Decomposition

Sequence-level KL decomposes into token-level KL terms evaluated along different trajectory distributions. The reverse-KL implementation uses a teacher–student log-ratio reward, corresponding to the k1 estimator rather than alternative KL estimators.

Sequence-level decomposition: The derivation applies autoregressive factorization and the tower property to sequence distributions, with end-of-sequence tokens accommodating variable-length responses.The proof initially presents a fixed-length sequence and notes how variable lengths are handled.
Sequence-level decomposition: Forward KL decomposes into token-level forward KL terms under the teacher-induced trajectory distribution, whereas reverse KL uses the student-induced distribution.The different state distributions couple token-level learning to different trajectory sources.
Reverse-KL reward: The reverse-KL log-ratio reward has expected value equal to the negative token-level reverse KL, so maximizing it minimizes that KL.This reward is evaluated under student samples.
Estimator comparison: The k1 estimator is unbiased in expectation but can be negative for individual samples, unlike k3, which is non-negative sample by sample.The non-negativity distinction concerns individual estimator values, not the expected KL.
Reverse-KL reward: The implementation uses the k1 log-ratio penalty estimator rather than the squared-log k2 approximation or non-negative unbiased k3 estimator.The paper explicitly identifies the reward as the negative of k1, not k2 or k3.

E Connection between Student-Prefix Forward KL and DAgger

Student-prefix forward KL is equivalent to the supervised subproblem of DAgger when the teacher supplies soft token distributions. The correspondence preserves learner-induced state supervision, but not necessarily the original DAgger algorithm or guarantees.

Correspondence: The correspondence maps DAgger’s learner to the student, learner state distribution to student prefixes, and supervised loss to forward KL.
Soft-label DAgger: Student-prefix forward KL exactly matches DAgger’s supervised learning subproblem when the teacher provides soft-label token distributions.The student supplies the learner-induced prefix distribution, while forward KL becomes soft-label cross-entropy up to teacher-entropy terms independent of the student parameters.
Scope: Without aggregating previous student-prefix data, the method is DAgger-style rather than a literal implementation of original DAgger.
Hard-label limit: Hard-label DAgger is a special case of student-prefix forward KL when the teacher distribution is a point mass.
Scope: The connection is objective-level only and does not directly transfer DAgger’s original online no-regret or task-performance guarantees.

G.2 Training Data and Prefix Construction

The study holds the prompt distribution fixed while varying prefix source and KL direction, using offline teacher rollouts or on-policy student rollouts. It evaluates both short and long horizons under shared training and RL settings.

Training data: All standalone objectives use the same DeepScaleR prompt distribution, so prefix source and KL direction are the controlled variables.
Prefix construction: Teacher-prefix training reuses offline trajectories and cached teacher logits, whereas student-prefix training samples responses from the current student and queries the teacher on visited states.
Sequence lengths: Distillation compares short 128-token and long 4096-token maximum-generation regimes under common optimization settings.
Infrastructure: On-policy rollout generation uses colocated vLLM with tensor parallel size 1 for both distillation and intermediate evaluation.
RL follow-up: The RL follow-up uses GRPO with an accuracy-only outcome reward and 4096-token maximum responses.

H Implementation Details

The implementation accounts for the extra generation and teacher-evaluation costs of student-prefix objectives and uses fused kernels to make exact full-vocabulary KL feasible in long contexts.

Compute accounting: Student-prefix reverse KL adds online student generation and teacher evaluation to the student training cost.
Evaluation: General evaluation reports acc for MMLU and acc_norm for ARC-Challenge, HellaSwag, and PIQA, with Avg. as their unweighted mean.
Compute accounting: One online student-prefix reverse-KL update costs roughly 3.6 cached-logit teacher-prefix reverse-KL updates in the 128-token experiments.
Memory: Full-vocabulary KL creates B × L × |V| intermediate tensors, making memory the bottleneck in long-context distillation.
Kernel: A fused full-vocabulary KL kernel streams vocabulary tiles and accumulates exact KL without materializing full logits or probability tensors.

I GRPO Follow-up Evaluation

The GRPO follow-up evaluates distillation checkpoints across mathematical and general-language benchmarks. The supplied results indicate stronger general-language averages for on-policy models, while the study remains limited to math reasoning and one RL algorithm.

Math evaluation: GRPO follow-up evaluation compares checkpoints from 128-token and 4096-token distillation across Avg@k, Pass@k, and response length.
Metrics: The evaluation uses k = 3 for GSM8K and MATH500 and k = 5 for AMC23 and AIME24.
General evaluation: All four distilled models improve the unweighted general-language average over the base student, with on-policy models achieving stronger averages under both KL directions.
Limitations: The study covers mathematical reasoning and uses GRPO with an accuracy-only outcome reward, leaving other domains and RL algorithms for future work.