Get the 2026 ML Training Cookbook | 52 recipes — GRPO, Flow Matching, World Models, and everything in between Download Now →
On-Policy Distillation
A knowledge distillation technique where the student model generates its own predictions and the teacher provides corrections on those predictions, reducing train-serve distribution mismatch.
On-policy distillation is a knowledge distillation technique where the student model generates its own predictions and the teacher provides corrections on those predictions, as opposed to traditional off-policy distillation where the student learns from the teacher's fixed outputs. In standard (off-policy) distillation, the teacher produces a set of soft labels on the training data, and the student is trained to match those labels. The student never sees its own outputs during training, which means it may encounter a different distribution of inputs at inference time than it was trained on — a phenomenon known as train-serve distribution mismatch.
On-policy distillation addresses this by having the student generate its own forward passes during training. The teacher then evaluates the student's outputs and provides soft-label corrections (typically using temperature scaling) on the student's own predictions. Because the student is trained on the distribution it actually produces, the resulting model is better aligned with its deployment behavior, leading to more robust and consistent performance.
This approach is particularly valuable in sequence generation tasks such as machine translation, summarization, and speech recognition, where autoregressive decoding means the student's input distribution at inference depends on its own prior outputs. By training on-policy, the student learns to recover from its own mistakes rather than merely imitating the teacher's behavior on gold-standard inputs.
- ■Student generates its own predictions; teacher provides corrections on those predictions
- ■Reduces train-serve distribution mismatch compared to off-policy distillation
- ■Teacher uses temperature scaling to produce soft labels on student forward passes
- ■Particularly effective for autoregressive sequence generation tasks
- ■Higher compute cost than off-policy distillation due to student inference during training
Distilling a large speech recognition teacher into a compact student model that will run on-device, where the student's own decoding errors during training teach it to self-correct.