On-Policy Context Distillation for Language Models

TL;DR

OPCD addresses the challenge of transient in-context knowledge in LLMs by learning a student to internalize context-conditioned behavior via on-policy distillation. It minimizes the reverse KL divergence $D_{\mathrm{KL}}(\pi_\theta(\cdot|x,y_{<t}) \| \pi_{\mathrm{teacher}}(\cdot|c,x,y_{<t}))$, enabling encoding of experiential knowledge and optimized system prompts into parameters. The framework is demonstrated on experiential knowledge distillation and system prompt distillation across math reasoning, text-based games, and domain-specific tasks, with consistent gains over off-policy baselines and improved out-of-distribution robustness. These results enable cross-size distillation and suggest practical pathways for continual, context-aware enhancement of LLMs without maintaining large prompts at inference time.

Abstract

Context distillation enables language models to internalize in-context knowledge into their parameters. In our work, we propose On-Policy Context Distillation (OPCD), a framework that bridges on-policy distillation with context distillation by training a student model on its own generated trajectories while minimizing reverse Kullback-Leibler divergence against a context-conditioned teacher. We demonstrate the effectiveness of OPCD on two important applications: experiential knowledge distillation, where models extract and consolidate transferable knowledge from their historical solution traces, and system prompt distillation, where models internalize beneficial behaviors encoded in optimized prompts. Across mathematical reasoning, text-based games, and domain-specific tasks, OPCD consistently outperforms baseline methods, achieving higher task accuracy while better preserving out-of-distribution capabilities. We further show that OPCD enables effective cross-size distillation, where smaller student models can internalize experiential knowledge from larger teachers.

Figures (17)

• Figure 1: Overview of on-policy context distillation (OPCD). Given a context and an input prompt, the student model generates a response without the context. It is then trained to minimize the reverse KL divergence to the teacher model that conditions on the context. The student internalizes the contextual information with on-policy learning.
• Figure 2: System prompt distillation on Medical.
• Figure 3: System prompt distillation on Safety.
• Figure 4: OPCD consistently improves the evaluation results of smaller Qwen3 models using experiential knowledge distilled from a frozen Qwen3-8B teacher. In contrast, directly injecting this knowledge into smaller-model contexts degrades performance.
• Figure 5: Comparison of OPCD and off-policy context distillation on in-distribution (safety) and out-of-distribution (medical) tasks when distilling from safety system prompt. Left: accuracy on the safety test dataset. Right: accuracy on the medical test dataset. OPCD achieves superior in-distribution performance while mitigating forgetting on OOD tasks.