Search papers, labs, and topics across Lattice.
The paper introduces On-Policy Context Distillation (OPCD), a method for distilling in-context knowledge into language models by training a student model on its own generated trajectories. OPCD minimizes the reverse Kullback-Leibler divergence between the student's output and a context-conditioned teacher model, effectively bridging on-policy and context distillation. Experiments across mathematical reasoning, text-based games, and domain-specific tasks demonstrate that OPCD outperforms baselines in task accuracy and out-of-distribution generalization, while also enabling effective cross-size distillation.
Language models can now internalize experiential knowledge and system prompts more effectively through on-policy context distillation, leading to better task accuracy and out-of-distribution generalization.
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.
Microsoft Research