On-Policy Context Distillation

• On-policy context distillation is a method where the student learns from its own trajectory data using teacher signals to directly mitigate exposure bias.
• It utilizes reverse-KL minimization and adaptive objective strategies to align student behavior with teacher policies across language, vision, and agentic domains.
• Empirical results show enhanced sample efficiency, greater stability, and the potential for student models to match or exceed teacher performance.

On-policy context distillation refers to a class of knowledge distillation and policy optimization algorithms in which a student model is trained to imitate, or extract structured knowledge from, a teacher model along the student's own generated trajectories. Unlike off-policy methods—which operate exclusively on static data or fixed teacher rollouts—on-policy context distillation injects dense, trajectory-aligned supervision into the student as it explores its own policy distribution, typically by minimizing a reverse Kullback-Leibler (KL) divergence or related objective. The approach mitigates exposure bias and has established itself as a foundational paradigm for reasoning transfer, efficient context internalization, and stability in model compression across language, vision-language, and agentic domains.

1. Foundations and Objectives

On-policy context distillation (OPCD) is fundamentally characterized by the choice to supervise the student model using teacher signals along rollouts sampled on-policy from the student itself, i.e., $y \sim \pi_\theta(\cdot \mid x)$, rather than conditioning solely on teacher-generated or offline data. The canonical objective is the expected per-token reverse-KL divergence: $$\mathcal{L}(\theta) = \mathbb{E}_{(x,c) \sim \mathcal{D}} \;\mathbb{E}_{y \sim \pi_\theta(\cdot \mid x)} \left[ \frac{1}{|y|} \sum_{t=1}^{|y|} D_\mathrm{KL} \left( \pi_\theta(\cdot|x,y_{<t}) \,\|\, \pi^*(\cdot|c,x,y_{<t}) \right) \right],$$ where $\pi_\theta$ is the student, $\pi^*$ the context-conditioned teacher, and $D$ the task distribution (Ye et al., 12 Feb 2026).

Contrasted with off-policy (forward-KL) distillation, on-policy methods avoid exposure bias—the error accumulation arising when students are never trained to recover from their own mistakes—and shift the optimization to high-probability regions of the student's own output, leveraging mode-seeking properties of the reverse KL to permit more aggressive imitation of the teacher's core competencies (Ye et al., 12 Feb 2026, Jang et al., 12 Jan 2026).

2. Algorithmic Methods and Variants

Numerous instantiations of on-policy context distillation have appeared across recent literature, differing in objectives, reward structure, context design, and stability refinements:

• Vanilla Reverse-KL-Based On-Policy Distillation: Minimizes reverse KL between the student and teacher distributions along the student-generated rollouts (Ye et al., 12 Feb 2026, Ko et al., 11 Mar 2026).
• Generalized On-Policy Distillation (G-OPD): Introduces a reward scaling factor $\lambda$ and an explicit reference model, interpolating between standard OPD, off-policy RL, and reward extrapolation (ExOPD). The objective generalizes to:

$$\mathbb{E}_{x,y\sim\pi_\theta} \left[ \lambda\sum_{t}\log\frac{\pi^*(a_t|s_t)}{\pi_\mathrm{ref}(a_t|s_t)} - D_\mathrm{KL}(\pi_\theta \| \pi_\mathrm{ref}) \right]$$

Reward extrapolation ($\lambda > 1$) can allow students to surpass teacher performance boundaries (Yang et al., 12 Feb 2026).

• Self-Distillation and Mixture Contexts: Student and teacher policies are realized as the same underlying model but under different contextual prompts (e.g., conditioning the teacher on privileged reasoning traces or on a "conciseness" instruction) (Zhao et al., 26 Jan 2026, Sang et al., 5 Mar 2026).
• Prefix-Only Distillation: The distillation signal is truncated to only the early prefixes of rollouts, motivated by the empirical observation that most of the on-policy loss mass appears in early tokens. This modification yields dramatic reductions in training cost with minimal degradation in downstream accuracy (Zhang et al., 16 Feb 2026).
• Entropy-Aware and Adaptive Objectives: Recent work proposes hybridizing reverse-KL with forward-KL at high-entropy tokens to preserve diversity where the teacher's prediction is uncertain—denoted Entropy-Aware On-Policy Distillation (EOPD):

$$\mathcal{L}_t^{\mathrm{EOPD}} = D_\mathrm{KL}(\pi_\theta \| \pi_\mathrm{te}) + \alpha_t D_\mathrm{KL}(\pi_\mathrm{te} \| \pi_\theta)$$

with $\alpha_t$ selected by an entropy threshold (Jin et al., 7 Mar 2026).

• Stabilization via Adaptive Target Reformulation: The Veto method blends teacher and student distributions geometrically in logit space, adjusting gradient flow via a parameter $\beta$ to suppress high-variance updates and control tradeoffs between decisiveness and diversity (Jang et al., 12 Jan 2026).
• Relaxed On-Policy Distillation (REOPOLD): Casts distillation as policy optimization, then adds reward clipping, entropy-based dynamic token sampling, and a unified exploration-refinement schedule to stabilize large-scale training (Ko et al., 11 Mar 2026).

These variants share the core principle of aligning students to teachers via on-policy sampling, but strategically differ in how approximation, context, and stability are controlled.

3. Integration with Reinforcement Learning and Policy Optimization

On-policy context distillation can be interpreted as KL-constrained reinforcement learning where the per-token log-likelihood ratio

$$r_{t} = \log \frac{\pi^*(a_t \mid s_t)}{\pi_\mathrm{ref}(a_t \mid s_t)}$$

functions as a dense reward and a policy-gradient update is performed directly via REINFORCE or PPO-style methods (Yang et al., 12 Feb 2026, Bousselham et al., 27 Oct 2025, Ko et al., 11 Mar 2026). For instance, in the context of vision-language reasoning, frameworks like VOLD combine Group Relative Policy Optimization (GRPO)—which uses normalized advantage estimation over sampled groups of trajectories—with on-policy reverse KL distillation, leveraging reward-guided masking to prevent overwriting correct student-discovered rollouts (Bousselham et al., 27 Oct 2025).

Furthermore, adversarial approaches (e.g., Generative Adversarial Distillation) integrate discriminators as on-policy reward models, evolving both generator (student) and discriminator in tandem to avoid reward hacking and staleness in feedback (Ye et al., 13 Nov 2025).

4. Applications Across Modalities and Tasks

On-policy context distillation frameworks have broad applicability:

• Mathematical and Logical Reasoning: Dense, trajectory-level teacher guidance yields substantial gains in pass@k accuracy and sample efficiency over conventional RL (Zhao et al., 26 Jan 2026, Jin et al., 7 Mar 2026, Zhang et al., 16 Feb 2026).
• Vision-Language and Multimodal Models: On-policy context distillation supports reasoning transfer from text-only LLMs into vision-language or multimodal students (e.g., VOLD), as well as causal distillation in real-time avatar video synthesis systems integrating textual, visual, and audio inputs (Bousselham et al., 27 Oct 2025, Chern et al., 29 Dec 2025).
• Agentic and Planning Models: In world-to-policy transfer, online on-policy distillation efficiently injects world-model rollouts and reward-driven planning expertise into compact students, achieving state-of-the-art safety and speed in driving tasks (Jiang et al., 25 Nov 2025).
• Knowledge, Prompt, and Behavioral Compression: On-policy context distillation enables models to internalize long system prompts, experiential knowledge, or context-specific behaviors, thus eliminating the test-time need for extended context or prompt injection (Ye et al., 12 Feb 2026).

5. Stability, Sample Efficiency, and Computational Considerations

Stability is a central challenge due to distribution mismatch and high-variance updates; to address this, several strategies have emerged:

• Initial Policy Alignment: Nearly all high-performing methods utilize a cold-start supervised fine-tuning phase (SFT) to reduce the initial gap between student and teacher, ensuring on-policy sampling remains in informative regions (Bousselham et al., 27 Oct 2025).
• Reward Masking and KL Selectivity: Masking distillation loss on successful student rollouts avoids stifling emergent correct strategies, maintaining a balance between imitation and innovation (Bousselham et al., 27 Oct 2025).
• Entropy-Guided Dynamic Sampling: Focusing optimization on high-entropy (i.e., uncertain) tokens maximizes signal utilization and reduces wasted compute (Ko et al., 11 Mar 2026, Jin et al., 7 Mar 2026).
• Prefix Distillation and Scheduling: Limiting distillation to early prefixes, or progressively increasing truncation length, achieves substantial training speedups (2×–47× reduction in FLOP) while retaining output accuracy, especially on lengthy reasoning tasks (Zhang et al., 16 Feb 2026).

The table summarizes leading strategies for stability and efficiency:

Strategy	Key Mechanism	Source
SFT Cold-start	Init student with supervised traces	(Bousselham et al., 27 Oct 2025)
Entropy-aware loss	Extra FKL where teacher is uncertain	(Jin et al., 7 Mar 2026)
Prefix-only distillation	Truncate loss to early tokens; schedule K	(Zhang et al., 16 Feb 2026)
Reward-guided KL masking	Mask KL on correct student rollouts	(Bousselham et al., 27 Oct 2025)
Mixture-based reward clipping	Lower-bound rewards to prevent instability	(Ko et al., 11 Mar 2026)
Veto objective	Interpolate teacher/student in logit space with $\beta$	(Jang et al., 12 Jan 2026)
Reward extrapolation (ExOPD)	Scale dense reward $\lambda > 1$ to push student past teacher	(Yang et al., 12 Feb 2026)

6. Empirical Impact and Empirically Driven Variants

Empirical studies consistently demonstrate that on-policy context distillation methods outperform off-policy counterparts in accuracy, sample efficiency, and out-of-distribution robustness on reasoning, planning, and multitask benchmarks. Notable effects include:

• Accuracy Gains: EOPD yields Pass@8 increases of +1 to +5 points over baseline OPD (Jin et al., 7 Mar 2026); VOLD achieves +6–20% absolute improvement over GRPO-only on visual reasoning (Bousselham et al., 27 Oct 2025).
• Sample Efficiency: REOPOLD achieves 6.7–12× greater sample efficiency compared to prior RL-based methods (Ko et al., 11 Mar 2026).
• Knowledge Compression: OPSDC compresses reasoning chains by 57–59% while increasing mathematical accuracy by up to 16 points; no explicit difficulty estimators or ground truths are required—compression is adaptive to problem difficulty (Sang et al., 5 Mar 2026).
• Robustness Across Domains: OPCD allows small student models to internalize knowledge from large teachers without the need for extended or optimized contexts at inference, facilitating cross-size distillation (Ye et al., 12 Feb 2026).
• Generalization and Extrapolation: The G-OPD/ExOPD formulation demonstrates that by setting $\lambda>1$ and using reward correction from a teacher’s base model, the student can, in some cases, surpass the teacher’s original performance (Yang et al., 12 Feb 2026).

7. Limitations, Open Challenges, and Future Directions

Despite substantial empirical progress, current on-policy context distillation methods present several open questions:

• Computational Cost: Full on-policy rollouts and dense teacher query are expensive for long-form outputs. Prefix scheduling (Zhang et al., 16 Feb 2026) and sparse approximations (Zhao et al., 26 Jan 2026) partially alleviate cost, but further efficiency enhancements are needed.
• Stability at Scale: Pathological behaviors, such as gradient explosion (in forward-KL), diversity collapse (in reverse-KL), or overfitting to idiosyncrasies of the student’s exploration, require continued study (Jang et al., 12 Jan 2026, Jin et al., 7 Mar 2026).
• Adaptive and Hybrid Objectives: Tuning between mode-seeking and mode-covering, e.g., through entropy-aware mixtures, curriculum learning, or adaptive parameterization in objectives (Veto $\beta$, G-OPD $\lambda$), remains an active area (Jang et al., 12 Jan 2026, Yang et al., 12 Feb 2026).
• Generalization Beyond Reasoning: While most benchmarks are reasoning-centric (math, vision, agentic planning), broader applications and implications for safety, dialog, and multimodal contexts are ongoing research directions (Chern et al., 29 Dec 2025, Ye et al., 12 Feb 2026).
• Scalability of Self-Distillation: The extent to which self-rationalizing mechanisms (conditioning the same model on privileged or concise contexts) scale to very large models or to challenging domains is not yet fully understood (Zhao et al., 26 Jan 2026, Sang et al., 5 Mar 2026).

The consensus across the surveyed literature is that on-policy context distillation provides a flexible foundation for future research in knowledge transfer, model compression, and safe deployment of autoregressive models under dense, contextually aligned supervision.