Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe

1. Problem Statement

On-Policy Distillation (OPD) has emerged as a core technique for post-training Large Language Models (LLMs), offering a dense, token-level reward signal that mitigates the exposure bias inherent in off-policy distillation. Despite its adoption by major industry players (e.g., Qwen3, MiMo, GLM-5), OPD remains poorly understood and fragile in practice.

The paper addresses three critical gaps:

1. Failure Modes: Stronger teachers do not always yield better distillation; in some cases, they cause the student to regress, while weaker teachers succeed.
2. Lack of Mechanism: It is unclear why OPD succeeds or fails at the token level and what specific signals drive the student's policy updates.
3. Scalability Limits: The reliability of dense token-level rewards over long reasoning trajectories is unverified.

2. Methodology

The authors employ a systematic empirical investigation combining controlled experiments, reverse distillation, and token-level dynamic analysis.

• Experimental Setup:
  • Models: Experiments utilize various model families (DeepSeek, Qwen, Skywork) ranging from 1.5B to 7B parameters.
  • Tasks: Primarily mathematical reasoning benchmarks (AIME 2024/2025, AMC 2023).
  • Distillation Variants: The study compares different supervision granularities: Sampled-Token, Full-Vocabulary, and Top-$k$ OPD.
• Key Techniques:
  • Reverse Distillation: Using a post-trained model (e.g., JustRL-1.5B) as a student to distill from its pre-trained checkpoint or a larger sibling. This isolates the effect of "thinking patterns" from raw capability.
  • Dynamic Metrics: The authors introduce and track specific metrics during training:
    • Overlap Ratio: The proportion of tokens shared between the student's and teacher's top-$k$ distributions.
    • Overlap-Token Advantage: The KL-divergence benefit specifically within the shared token set.
    • Entropy Gap: The difference in uncertainty between student and teacher distributions.
  • Ablation Studies: Decomposing the optimization target into "Overlap Tokens" vs. "Non-Overlap Tokens" to determine the source of gradient signals.
  • Remedial Strategies: Testing "Off-Policy Cold Start" (SFT before OPD) and "Teacher-Aligned Prompts" to rescue failing configurations.

3. Key Contributions & Findings

A. Phenomenology: Two Governing Conditions for Success

The paper identifies two necessary conditions for OPD to succeed, challenging the assumption that "stronger teacher = better student":

1. Thinking-Pattern Consistency: The student and teacher must share compatible reasoning patterns (e.g., similar token distributions).
   • Evidence: A 7B model with a different reasoning style (e.g., non-thinking vs. RL-trained) fails to improve a 1.5B student, even if the 7B model has higher benchmark scores. Conversely, a weaker model with the same reasoning pattern (e.g., pre-RL vs. post-RL of the same family) can successfully transfer knowledge.
2. New Knowledge (Beyond Scale): The teacher must possess capabilities the student has not yet seen.
   • Evidence: Distilling from a larger model within the same training pipeline (same data, same recipe) yields diminishing returns because the distributions converge. Significant gains only occur when the teacher has acquired new knowledge (e.g., via additional RL post-training) that the student lacks.

B. Mechanism: Progressive Alignment on High-Probability Tokens

The authors reveal the token-level mechanism driving successful OPD:

• Overlap is King: Successful OPD is characterized by a progressive increase in the overlap ratio of high-probability tokens (Top-$k$) between student and teacher.
• Concentration of Mass: In successful runs, the shared top-$k$ tokens concentrate 97%–99% of the total probability mass.
• Self-Reinforcing Cycle: Optimization focuses almost entirely on these overlap tokens. Non-overlap tokens contribute negligible gradients. As the student aligns with the teacher on these tokens, the overlap set grows, creating a virtuous cycle of alignment.
• Failure Mode: When conditions are unmet, the overlap ratio stagnates, and the entropy gap remains large, leading to no performance gain.

C. Practical Recipe: Rescuing Failing OPD

To bridge the thinking-pattern gap, the authors propose two strategies:

1. Off-Policy Cold Start: Perform Supervised Fine-Tuning (SFT) on teacher-generated rollouts before starting OPD. This raises the initial overlap ratio, allowing the student to immediately exploit the teacher's token-level signals.
2. Teacher-Aligned Prompt Selection: Using prompts from the teacher's specific post-training data improves alignment. However, this must be mixed with out-of-distribution prompts to prevent entropy collapse (over-confidence).

D. Limitations: The Cost of Dense Supervision

• Reward Degradation: The quality of the teacher's token-level reward degrades as the trajectory length increases.
• Instability: In long-horizon settings (e.g., >10k tokens), the teacher's reward becomes noisy for prefixes the teacher hasn't seen, causing the student's entropy to spike and training to collapse.
• Global vs. Local: Even in failing runs, the global reward signal correlates with correctness (AUROC ~0.75), but the local optimization geometry is anisotropic (gradients cancel out), preventing effective learning.

4. Results

• Performance Gains: Successful OPD (with aligned patterns and new knowledge) can recover >80% of the teacher-student performance gap.
• Reverse Distillation: Distilling a post-trained 1.5B model back to its pre-RL state causes it to regress to pre-RL performance, proving OPD actively overwrites the student's policy with the teacher's thinking pattern.
• Ablation: Optimizing only on the Overlap Top-$k$ tokens achieves performance nearly identical to full Top-$k$ optimization, confirming that non-overlap tokens are irrelevant for the gradient signal.
• Prompt Sensitivity: Changing the prompt template to match the teacher's training format significantly boosts overlap and final accuracy.

5. Significance

This paper fundamentally shifts the understanding of On-Policy Distillation:

1. Paradigm Shift: It moves the focus from "Teacher Strength" to "Thinking Pattern Compatibility" and "Knowledge Novelty." A stronger teacher is not a guarantee of success if the reasoning styles diverge.
2. Mechanistic Insight: It identifies high-probability token overlap as the primary locus of learning, explaining why OPD works and why it fails (stagnant overlap).
3. Practical Guidelines: It provides a concrete "Recipe" (Cold Start + Prompt Alignment) to stabilize OPD, making it more robust for industrial applications.
4. Scalability Warning: It highlights a critical limitation for long-horizon reasoning (e.g., agentic workflows), suggesting that pure dense token-level supervision may not scale indefinitely without hybrid approaches.

In summary, the paper argues that OPD is not a "free lunch" of dense rewards but a delicate process of pattern alignment that requires careful initialization and data selection to succeed.