New method reduces training cost by skipping some distillation steps

Not Every Rollout Needs Distillation: Correctness-Filtered On-Policy Self-Distillation

Artificial Intelligence

Summary

Training AI models often involves teaching them by reviewing many example answers, but not every example adds new helpful information. The authors found that some incorrect answers carry redundant signals and can be skipped to save effort without hurting learning. They devised a way to filter which examples to use during training, cutting computation time significantly while keeping performance the same. This approach helps make AI training more efficient by focusing on the most useful teaching examples.

What this means in practice

  • For machine learning engineers: Train large AI models more efficiently by selectively distilling only certain incorrect responses to cut computation time without losing accuracy.
  • For ai infrastructure teams: Reduce expensive GPU or TPU training costs by implementing correctness-filtered self-distillation in existing training pipelines for large language models.

Authors

Harrison Zhao, Lei Huang, Fei Liu, Wen Shi, Zijian Jin

Abstract

On-policy self-distillation (OPSD) supervises student-generated trajectories with a reference-conditioned version of the same base model. We ask whether every trajectory supplies distinct supervision that warrants its distillation cost. Gradient measurements on Qwen3-1.7B at initialization reveal that correct and incorrect responses have similar individual gradient magnitudes, but incorrect responses have approximately three times the coherent gradient share. Their estimated population-mean directions also largely overlap. A signal--noise model fitted at a small group size tracks alignment at larger group sizes without refitting. Motivated by these measurements, we introduce correctness-filtered OPSD (C-OPSD), which generates one response per problem and distills only a subset of incorrect responses. Removing unselected rows before both distillation forwards avoids their teacher/student forward and student backward work. Across Qwen3-1.7B, 4B, and 8B on AIME24, AIME25, and HMMT25, C-OPSD achieves on-par performance with full OPSD. In a measured micro-batch-size-four implementation, executed rows decrease by approximately 68%, and with retained rows packed into full micro-batches, training takes about 40% less time while rollout generation remains unchanged. Together, the measurements and training results connect outcome-conditioned supervision redundancy to more compute-efficient distillation.