Technical Summary: Decoupling KL and Trajectories in LLM Distillation

Problem Statement

Knowledge distillation has become a cornerstone of Large Language Model (LLM) post-training, particularly for reasoning tasks, often operating alongside Reinforcement Learning (RL). However, the design space of distillation remains poorly understood. Current pipelines typically follow one of two paradigms: off-policy distillation (training on teacher-generated traces, often via Supervised Fine-Tuning) or on-policy distillation (OPD) (training on student-generated rollouts with teacher feedback).

The paper argues that these prevailing paradigms implicitly couple two orthogonal design choices:

1. Prefix Source: Whether supervision is applied to prefixes generated by the teacher (off-policy) or the student (on-policy).
2. Token-level KL Direction: Whether the objective minimizes the forward KL ($KL(p_{teacher} \parallel q_{student})$) or the reverse KL ($KL(q_{student} \parallel p_{teacher})$).

Conventional practice couples teacher prefixes with forward KL and student prefixes with reverse KL. The authors contend this coupling is not intrinsic to the token level and that decoupling these axes reveals a broader design space with four valid objectives, each corresponding to distinct learning regimes.

Methodology

Theoretical Framework: Decoupling the Axes

The authors decompose sequence-level KL divergence using the autoregressive chain rule. They demonstrate that:

• Forward KL ($KL(p_T \parallel q_S)$) naturally pairs with teacher prefixes (off-policy) because the expectation is taken over the teacher's trajectory distribution.
• Reverse KL ($KL(q_S \parallel p_T)$) naturally pairs with student prefixes (on-policy) because the expectation is taken over the student's trajectory distribution.

By treating these as independent axes, the authors define four objectives:

1. Teacher Prefix + Forward KL: Corresponds to off-policy SFT (standard distillation).
2. Student Prefix + Forward KL: Corresponds to DAgger-style on-policy SFT.
3. Teacher Prefix + Reverse KL: Corresponds to an offline RL-style distillation objective.
4. Student Prefix + Reverse KL: Corresponds to On-Policy Distillation (OPD), viewed as dense-reward on-policy RL.

The paper establishes gradient-level identities showing that:

• Forward KL yields an SFT-style cross-entropy gradient with teacher soft targets.
• Reverse KL yields an RL-style policy gradient with a dense reward defined by the teacher-student log-ratio ($r = \log p_T - \log q_S$).

Experimental Setup

The authors conduct a controlled study on mathematical reasoning using Qwen3 models (0.6B student, 4B/8B teachers). They evaluate the four objectives in two regimes:

1. Standalone Distillation: Comparing the final performance of models trained solely with each objective.
2. Distillation-then-RL: Using the distilled checkpoints as initializations for Group Relative Policy Optimization (GRPO).

They track accuracy, predictive entropy, and response length across various sequence lengths (128 vs. 4096 tokens) and benchmarks (MATH500, GSM8K, AIME24, AMC23).

Key Contributions

1. Identification of Three Tradeoffs

The empirical analysis reveals three fundamental tradeoffs governing distillation design:

• Accuracy–Entropy Tradeoff (KL Direction): Reverse KL consistently achieves higher standalone accuracy (Avg@k) but drives predictive entropy toward collapse, reducing diversity (lower Pass@k) and making the model less stable for subsequent RL. Forward KL preserves entropy, supporting more stable RL improvement, though with lower initial standalone accuracy.
• Quality–Compute Tradeoff (Prefix Source): Student prefixes (on-policy) yield higher accuracy under matched training steps by supervising states the student actually visits. However, teacher prefixes (off-policy) are more compute-efficient under matched FLOPs because they allow reusing cached teacher logits and offline trajectories, avoiding the cost of online student generation.
• Accuracy–Stability Tradeoff (Training Length): Long-sequence distillation (4096 tokens) improves reasoning accuracy but, particularly under reverse KL, causes severe entropy collapse and "length inflation" (responses hitting the generation limit). Short-sequence training is more stable but less accurate.

2. Proposed Methods to Balance Tradeoffs

Motivated by these findings, the authors propose two practical methods:

• KL Mixing: A weighted combination of forward and reverse KL at each prefix. The authors find that in long-sequence distillation, a forward-heavy mixture is crucial. It retains the accuracy benefits of reverse KL while preventing entropy collapse and length inflation, stabilizing the training dynamics.
• Entropy-Gated Length Curriculum: A training schedule that starts with short horizons and increases sequence length only while the predictive entropy remains above a stability threshold. This approach avoids the entropy collapse associated with fixed long-horizon training.

Results

Standalone Distillation

• Reverse KL generally outperforms Forward KL in standalone accuracy (Avg@k), with gains of +2.45 points on average.
• Student Prefixes outperform Teacher Prefixes in accuracy, especially in long-sequence settings (+3.55 Avg@k gain).
• Entropy Collapse: Pure reverse KL in long-sequence settings drives entropy close to zero and causes response lengths to inflate to the maximum limit (e.g., 8192 tokens), degrading Pass@k.

Distillation-then-RL (GRPO)

• Initialization Stability: Models initialized with Forward KL (even if lower in standalone accuracy) show more reliable improvement during GRPO because they retain higher entropy.
• Reverse KL Pitfall: Models initialized with Reverse KL often start with higher accuracy but degrade or plateau during RL due to low entropy and reduced exploration capacity.
• Curriculum Success: The entropy-gated length curriculum significantly improves performance. Compared to fixed 4096-token training, it improved Avg@k by 3.6 points and Pass@k by up to 5.8 points, while reducing average response length by roughly 3× (e.g., from ~8185 to ~2688 tokens).

KL Mixing Results

In long-sequence settings, a mixture with high forward KL weight (e.g., 0.8 Forward / 0.2 Reverse) preserved most of the reverse KL's accuracy while stabilizing entropy and preventing length inflation.

Significance and Claims

The paper claims to provide a unified theoretical framework that connects disparate training recipes (SFT, DAgger, Offline RL, OPD) as specific instantiations of two orthogonal design choices. By decoupling prefix source and KL direction, the authors argue that practitioners can move beyond rigid paradigms to design reasoning distillation objectives that explicitly balance:

• Accuracy (via reverse KL and long sequences),
• Diversity/Exploration (via forward KL and entropy preservation),
• Compute Efficiency (via teacher prefixes/caching), and
• Downstream RL Trainability (via stable entropy).

The proposed methods (KL mixing and entropy-gated curriculum) are presented as practical tools to navigate the identified tradeoffs, enabling the training of smaller reasoning models that are both accurate and stable for further policy optimization. The authors emphasize that the "best" standalone objective is not necessarily the best initialization for an RL pipeline, highlighting the need for a holistic view of the post-training process.