Projected distillation improves learning from mistakes on hard problems
Projected Distillation: On Learning to Solve Hard Problems from Close, Correct Targets
Machine LearningArtificial Intelligence
Summary
When AI models try to solve difficult tasks, they often fail and don't get useful feedback to learn from their mistakes. The authors propose projected distillation, a method that gently corrects the model's wrong attempts instead of replacing them completely with right answers. This approach helps the model learn better and keep its previous skills at the same time. Their tests show this method works well across different types of problems like math and coding.
What this means in practice
- •For ai software engineers: Improve training AI models to learn from their own mistakes on complex tasks without losing existing abilities.
- •For automation developers: Develop automated systems that adapt to hard problem solving by gently correcting errors during learning.
Authors
Shripad V. Deshmukh, Christian Bock, Kangrui Ruan, Abhijeet Awasthi, Vijay Lingam, Rein Houthooft, Manish Prajapat
Abstract
On-policy learning is central to post-training: training on a model's own generations limits distribution shift and helps preserve existing capabilities. On hard problems, however, on-policy RL provides no reward-based learning signal when all attempts fail. On-policy distillation provides next-token supervision from a teacher, but on hard problems, if a rollout goes wrong early, subsequent supervision is conditioned on prefixes that already contain the mistake. Correct behavior can instead come from a gold or teacher solution, but such off-policy targets can differ substantially from the model's generations and, as we show, may erode existing capabilities. We therefore treat on-policyness as a continuum and propose projected distillation, which minimally edits the student's failed attempts into correct targets. The method uses contrastive updates to favor corrections over failed attempts and refreshes these corrections as the student evolves, enabling near-on-policy learning even when all student attempts fail. Theoretically, we prove that the student's learning regret grows with the edit size: it scales linearly with the horizon under minimal correction, as in DAgger, and quadratically under full-answer replacement, as in behavior cloning. Empirically, across two students (Qwen3-8B and OLMo-2-7B) and five domains spanning math, code, science, and tool use, projected distillation gives the highest area under the specialization-retention frontier (best on 8 of 9 cells against GRPO and OPD, and 7 of 9 against SPOT and SODA). Our ablations highlight the importance of online corrections, the contrastive loss, and on-policyness: moving targets farther off-policy steadily worsens the specialization-retention trade-off, with retention falling from 0.916 to 0.744 when minimal corrections are replaced by student-agnostic gold solutions.