Source-linked AI summary
On-Policy Context Distillation for Language Models
Tianzhu Ye, Li Dong, Xun Wu, Shaohan Huang, Furu Wei
TL;DR
Existing context distillation methods are limited by off-policy training and forward KL, while in-context knowledge disappears when prompts are removed. OPCD trains a context-free student on its own trajectories using reverse KL against a context-conditioned teacher. Across reasoning, games, and domain-specific prompt tasks, it outperforms baselines, preserves out-of-distribution capabilities, and supports distillation from larger teachers.
Problem
In-context knowledge is transient, and prior off-policy forward-KL context distillation suffers exposure bias and mode-covering behavior.
Method
OPCD trains a context-free student on its own generated trajectories while minimizing reverse KL against a context-conditioned teacher.
Results
Across mathematical reasoning, text-based games, and domain-specific tasks, OPCD outperforms baselines while preserving out-of-distribution capabilities and enabling smaller students to internalize knowledge from larger teachers.
Takeaways & Limitations
OPCD consolidates experiential knowledge and optimized system-prompt behaviors into model parameters without requiring the original context at inference.
Takeaways & Limitations
Self-distillation can be unstable, although using an exponential moving average of student parameters as teacher can alleviate this instability.
Abstract
from arXiv · showhide
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.
1 Introduction
OPCD addresses the transient nature of in-context knowledge and the exposure-bias and mode-covering drawbacks of prior off-policy context distillation. It trains students on their own trajectories against a context-conditioned teacher and applies this framework to experiential knowledge and system prompt distillation.
- In-context knowledge guides behavior without parameter updates but is lost when the context resets.
- Context distillation compresses context-aware teacher behavior into student weights, allowing inference without the original context.
- Prior methods use off-policy forward KL on fixed data, creating exposure bias between teacher-guided training and autoregressive inference.
- Forward KL can encourage mode-covering distributions that produce hallucinations or overly broad behavior when students cannot model the teacher fully.
- OPCD samples student responses without context and minimizes reverse KL to a context-conditioned teacher along those trajectories.
- OPCD targets experiential knowledge distillation and system prompt distillation across reasoning, games, medical question answering, and safety classification.
- Across mathematical reasoning, text-based games, and domain-specific tasks, OPCD outperforms baselines while preserving out-of-distribution capabilities and reducing catastrophic forgetting.
2 Related Work
Related work frames OPCD as the intersection of context distillation, on-policy distillation, and self-distillation. Its distinguishing focus is internalizing contextual information while addressing the training–deployment mismatch of off-policy methods.
- Context Distillation: Context distillation compresses in-context knowledge into parameters, but prior forward-KL methods suffer exposure bias from teacher-guided training.
- On-Policy Distillation: On-policy distillation trains students on their own trajectories and uses reverse KL to promote mode-seeking behavior around high-likelihood teacher regions.
- On-Policy Distillation: OPCD adapts on-policy distillation specifically to context internalization, bridging training and deployment distributions.
- Self-Distillation: Self-distillation improves models from their own outputs or conditioned versions, including on-policy training with privileged information.
3 Method
OPCD internalizes context by training a context-free student on its own generated responses while matching a context-conditioned teacher with token-level reverse KL. The framework supports larger-teacher and self-distillation configurations.
- Method: OPCD compresses in-context knowledge c into student parameters so the student can replicate a context-aware teacher without c at inference.
- Method: Given input x, OPCD uses on-policy sampling to minimize reverse KL between the context-free student and context-conditioned teacher distributions.
- Objective: Sequence-level divergence is decomposed into a sum of token-level divergences for optimization.
- Objective: The loss averages token-level reverse KL over responses y sampled from the student on training data D, with c denoting the knowledge to internalize.
- Objective: The implementation approximates analytic KL by summing over the top-k tokens with highest student probability.
- Objective: Reverse KL encourages mode-seeking behavior by focusing the student on teacher high-probability tokens rather than the long tail.
- Training Procedure: The student generates complete trajectories without context, after which the teacher evaluates the concatenated sequence [c; x; y].
- Teacher Configurations: Teacher–student distillation can use a larger or more capable teacher, while self-distillation uses shared weights with context supplied only to the teacher.
4 Experiments
The experiments evaluate experiential knowledge and system prompt distillation across mathematical reasoning, text-based games, and domain-specific tasks. OPCD generally improves accuracy, preserves out-of-distribution performance, and supports distillation across model sizes.
- Experiential Knowledge Distillation: Experiential knowledge distillation extracts transferable knowledge from self-generated solution traces, accumulates it across problems, and consolidates it into student weights.The consolidation stage removes the need for extended context at inference time.
- Experiential Knowledge Distillation: OPCD outperforms off-policy context distillation on math and text-game test accuracy, and can surpass the original model with experiential knowledge supplied in context.The consolidation data came from validation-extracted experiential knowledge that the original model did not access, providing an additional learning signal.
- System Prompt Distillation: OPCD outperforms off-policy context distillation in test accuracy across most medical and safety system-prompt configurations.The teacher and student use the same model size, with a frozen teacher; on-policy training also produces more stable training improvements.
- Effect of Model Size: OPCD improves test accuracy across Qwen3-1.7B, Qwen3-4B, and Qwen3-8B student scales when experiential knowledge is generated by Qwen3-8B.Directly injecting teacher-generated knowledge into smaller-model contexts can degrade performance instead.
- OOD Performance: Approximately 2% higher IF-Eval scores than context distillation are achieved by OPCD on Frozen Lake, indicating better OOD performance after experiential knowledge distillation.The method samples from the student distribution, which the passage identifies as mitigating forgetting on OOD tasks.
- OOD Performance: OPCD achieves higher in-distribution safety performance than off-policy distillation while maintaining OOD medical performance and exceeding the off-policy baseline by approximately 4 points.This comparison uses Qwen2.5-3B-Instruct as student and Qwen2.5-7B-Instruct as frozen teacher with a safety system prompt.
5 Conclusion
The paper concludes that OPCD internalizes contextual information into model weights through on-policy distillation and outperforms baselines across several tasks while improving out-of-distribution generalization.
- Using raw response traces as experiential context degrades performance on the math validation dataset.
- OPCD internalizes in-context knowledge into language-model parameters through on-policy distillation.
- OPCD outperforms baseline methods across math problem solving and text-based games while enhancing out-of-distribution generalization.
A Experiential Knowledge Distillation Details
The experiential knowledge distillation details specify prompt wrappers for extracting, accumulating, and applying experiential knowledge across math and text-game tasks.
- Extraction: Experiential knowledge is extracted from lines beginning with “– EXPERIENCE ITEM:”.
- Extraction: Text-game extraction uses a prompt wrapper containing the latest interaction history as internal experience.
- Application: New problems receive accumulated experiential knowledge through a dedicated prompt template.
A.2 Dataset Details
The experiments use DAPO-Math-17K and the Frozen Lake and Sokoban text-based game environments, with task-specific interaction settings.
- Datasets: The experimental datasets are DAPO-Math-17K, Frozen Lake, and Sokoban, implemented partly through TextArena.
- Math: DAPO-Math-17K contains approximately 14K verifiable English math problems, each with a numerical answer.
- Text-based games: Frozen Lake requires grid navigation to a goal while avoiding holes, whereas Sokoban requires pushing a box to a designated target.
A.3 Training Details
Training constructs experiential contexts from teacher-generated traces, then distills students on training problems using reverse KL divergence and task-specific response limits.
- Context construction: The procedure samples 30 validation problems, extracts experiential knowledge from teacher traces, and forms 30 contexts, repeated 10 times to yield 300 contexts.
- Distillation: Students are distilled for 50 steps using selected experiential contexts, with batch size 128 and reverse KL computed over the top 256 student-probability tokens.
- Task settings: Math responses allow up to 16384 tokens, while text-game interaction lasts up to 5 rounds with responses capped at 1024 tokens.
A.4 Experiential Knowledge Accumulation
The teacher extracts experiential knowledge by solving sampled validation problems, and accumulating knowledge from different problems improves validation accuracy.
- 30 validation problems are sampled for the teacher to solve and extract experiential knowledge, with the procedure repeated 10 times.
- Validation accuracy improves as experiential knowledge accumulates across different problems.Figure 7 demonstrates the improvement.
A.5 Experiential Knowledge Examples
The paper illustrates experiential knowledge with transferable problem-solving strategies for mathematical reasoning and the Frozen Lake game.
- Math: The math examples include recognizing simpler patterns in interdependent sequences, including Fibonacci-like recurrences.
- Math: Modular arithmetic can expose periodicity and reduce computation by enabling predictions from cycle lengths.
- Frozen Lake: The Frozen Lake examples emphasize systematic movement toward the goal, direct routes, and minimizing unnecessary backtracking.
- Frozen Lake: Strategic use of available actions, such as down or right, is presented as key to reaching the goal in the fewest steps.
B System Prompt Distillation Details
This appendix details system prompt distillation for medical and safety tasks, including datasets, prompts, training settings, and evaluation procedures.
- B.1 System Prompts: System prompts optimized for medical and safety tasks are drawn from MetaSPO, with medical evaluation using 500 held-out MedMCQA samples.
- B.1 System Prompts: Medical distillation trains on both the teacher’s selected answer choice and supporting explanation for each multiple-choice question.
- B.1 System Prompts: Safety evaluation combines Tweet Eval and Hatecheck datasets, with the Tweet Eval instruction requiring offensive-text classification and a yes/no answer format.
- B.2 Training Details: Distillation conditions the student on the corresponding system prompts and runs for 50 steps using reverse KL divergence over the student-ranked top 256 vocabulary tokens.
- B.2 Training Details: Training uses batch size 128 and sweeps learning rates of 1e-6 and 5e-6.
- B.2 Training Details: Generated responses are capped at 512 tokens, checkpoints are saved every 2 steps, and test accuracy averages the three best-performing checkpoints.