On-policy attention distillation improves math reasoning accuracy and efficiency
Teach Yourself Where to Look: On-Policy Attention Self-Distillation for Reasoning
Machine LearningComputation and Language
Summary
Mathematical reasoning models often learn by following guidance from a teacher model that knows the correct answers. The authors found that teaching these models not only what answers to expect but also where to focus their attention in the problem helps them perform better. Their new method, called On-Policy Attention Self-Distillation (OPASD), directs the student model to look at important parts of the problem identified by the teacher, even when those parts are hidden from the student. This leads to more accurate answers, shorter reasoning steps, and faster training compared to previous approaches.
What this means in practice
- •For ai developers: Train math reasoning AI models to improve accuracy and efficiency by guiding their focus during learning.
- •For automated tutoring platforms: Enhance automated math problem solvers used in tutoring systems to produce more precise and faster explanations.
Authors
Safaeid Hossain Arib, Rabeya Akter, Ismam Nur Swapnil, Md. Faiyaz Abdullah Sayeedi, Md Mofijul Islam, Tasnim Mohiuddin
Abstract
On-policy self-distillation trains reasoning models on their own trajectories using dense token distribution guidance from a privileged teacher with access to a verified solution. This supervision transfers what the teacher predicts without directly transferring where it attends within the preceding context. We introduce On-Policy Attention Self-Distillation (OPASD), which complements token-level supervision with solution-conditioned attention distillation. Because the privileged teacher can attend to verified solution tokens unavailable to the student, OPASD projects teacher attention onto student-visible positions and renormalizes the resulting distribution before alignment. Across three model sizes and four competition-level mathematics benchmarks, OPASD consistently outperforms token-only OPSD, improving average accuracy by 4.98 to 8.40 percentage points. OPASD also avoids the response-length inflation and performance degradation observed with token-only distillation, reducing generated rollout tokens by 73.9% and estimated model compute by 72.6% while training 1.53x faster. These results show that solution-conditioned attention provides a complementary supervision signal that makes on-policy self-distillation more accurate, stable, and compute-efficient.