On-Policy Context Distillation for Language Models
This work addresses the problem of improving knowledge internalization in language models for researchers and practitioners, though it appears incremental as it builds on existing context distillation methods.
The paper tackles the problem of enabling language models to internalize in-context knowledge into their parameters by proposing On-Policy Context Distillation (OPCD), which bridges on-policy distillation with context distillation, and demonstrates effectiveness in applications like experiential knowledge distillation and system prompt distillation, achieving higher task accuracy and better out-of-distribution capabilities across tasks such as mathematical reasoning and text-based games.
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.