On-Policy Context Distillation for Language Models

Paper: 2602.12275 Authors: Tianzhu Ye, Li Dong, Xun Wu, Shaohan Huang, Furu Wei Categories: cs.CL

Abstract

This paper introduces On-Policy Context Distillation (OPCD), a framework that enables language models to internalize contextual knowledge into their parameters through on-policy learning. Unlike traditional context distillation methods that rely on off-policy data, OPCD trains student models on their own generated trajectories while minimizing reverse KL divergence against a context-conditioned teacher model. The framework demonstrates effectiveness across two key applications: experiential knowledge distillation, where models learn from their historical solution traces, and system prompt distillation, where models internalize behaviors from optimized prompts. Experimental results across mathematical reasoning, text-based games, and domain-specific tasks show that OPCD consistently outperforms baseline methods while maintaining strong out-of-distribution generalization.

Key Contributions

• Introduction of OPCD, a principled framework that bridges on-policy distillation with context distillation using reverse KL divergence
• Demonstration of experiential knowledge distillation, enabling models to consolidate transferable knowledge from their own solution trajectories
• System prompt distillation capability that internalizes beneficial behaviors from optimized prompts into model parameters
• Comprehensive evaluation across diverse domains showing consistent improvements over baseline methods
• Evidence of effective cross-size distillation, where smaller models successfully internalize knowledge from larger teachers

Methodology

The OPCD framework operates by training a student model $\pi_\theta$ to match the behavior of a context-conditioned teacher model $\pi_{\text{teacher}}(\cdot|c)$, where $c$ represents the context (e.g., few-shot examples or system prompts). The key innovation lies in using on-policy sampling: the student generates its own trajectories, which are then used as training data.

The training objective minimizes the reverse KL divergence:

$\mathcal{L}(\theta) = \mathbb{E}_{x \sim \pi_\theta} [D_{\text{KL}}(\pi_{\text{teacher}}(\cdot|x,c) \| \pi_\theta(\cdot|x))]$

This formulation encourages the student to cover all modes of the teacher’s distribution while avoiding overconfident predictions on low-probability regions. The on-policy nature ensures that the student learns from realistic scenarios it will encounter during deployment, rather than from potentially off-distribution teacher-generated data.

For experiential knowledge distillation, the context $c$ consists of successful solution traces from the model’s own problem-solving history. For system prompt distillation, $c$ represents carefully crafted prompts that encode desired behaviors. The framework iteratively refines the student model by generating new trajectories and updating parameters based on the teacher’s guidance.

Experimental Results

The authors evaluate OPCD across three domains: mathematical reasoning (GSM8K, MATH), text-based games (ALFWorld, WebShop), and domain-specific tasks. In mathematical reasoning, OPCD achieves substantial improvements over baseline context distillation methods, with gains of 3-5% on GSM8K and 2-4% on MATH datasets. The experiential knowledge distillation variant shows particularly strong results, suggesting that models can effectively learn from their own successful problem-solving strategies.

In text-based game environments, OPCD demonstrates superior sample efficiency and final performance. On ALFWorld, the method achieves 15% higher success rates compared to standard fine-tuning approaches, while maintaining better generalization to unseen game scenarios. The system prompt distillation experiments reveal that OPCD can successfully internalize complex behavioral patterns, such as step-by-step reasoning or specific formatting preferences, without requiring explicit prompts at inference time.

Cross-size distillation experiments show promising results, with 7B parameter students successfully learning from 13B and 70B teachers. The performance gap between student and teacher narrows significantly compared to traditional distillation methods, suggesting that on-policy training helps smaller models better capture the reasoning patterns of larger models.

Implications and Future Directions

OPCD represents a significant advancement in making language models more efficient and capable. By internalizing contextual knowledge, models can achieve strong performance without the computational overhead of processing lengthy contexts at inference time. This has practical implications for deployment scenarios where latency and cost are critical factors.

The experiential learning aspect opens new possibilities for continual learning systems where models progressively improve by reflecting on their own experiences. Unlike traditional supervised learning that requires external labeled data, experiential knowledge distillation enables self-improvement through practice and reflection.

The framework also addresses a key challenge in prompt engineering: while carefully crafted prompts can dramatically improve model behavior, they add computational cost and complexity. System prompt distillation offers a path to “bake in” these improvements, making them available without runtime overhead.

Future research directions include exploring OPCD for multi-task learning scenarios, investigating optimal strategies for selecting which experiences to distill, and extending the framework to handle more complex forms of contextual knowledge such as retrieved documents or tool usage patterns.

Takeaways

1. On-policy training is crucial for effective context distillation, as it ensures the student learns from realistic scenarios rather than potentially off-distribution teacher data
2. Language models can successfully learn from their own experiences, consolidating successful problem-solving strategies into their parameters through experiential knowledge distillation
3. System prompt distillation enables models to internalize beneficial behaviors without requiring explicit prompts at inference time, reducing computational overhead
4. OPCD maintains strong out-of-distribution generalization while improving in-distribution performance, avoiding the overfitting issues common in traditional distillation
5. Cross-size distillation with OPCD allows smaller models to effectively learn from larger teachers, enabling more efficient deployment without significant performance degradation