Teach Yourself Where to Look: On-Policy Attention Self-Distillation for Reasoning
cs.LG, cs.CL
Submitted: 2026-09-27
Updated: 2026-09-29
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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.
Terminology
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.
Sources
- HDPO: Hybrid Distillation Policy Optimization via Privileged Self-Distillation
- Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
- DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
- Learning to Focus: Causal Attention Distillation via Gradient-Guided Token Pruning
- Privileged, but Biased: How PI-Conditioned Teachers Break Self-Distillation
- Self-Distillation Zero: Self-Revision Turns Binary Rewards into Dense Supervision
- Distilling the Knowledge in a Neural Network
- Uni-OPD: Unifying On-Policy Distillation with a Dual-Perspective Recipe
- LoRA: Low-Rank Adaptation of Large Language Models
- Reinforcement Learning via Self-Distillation
- Asymmetric On-Policy Distillation: Bridging Exploitation and Imitation at the Token Level
- Entropy-Aware On-Policy Distillation of Language Models
- Why Does Self-Distillation (Sometimes) Degrade the Reasoning Capability of LLMs?
- Scaling Reasoning Efficiently via Relaxed On-Policy Distillation
- Reinforced Attention Learning
- Unifying Group-Relative and Self-Distillation Policy Optimization via Sample Routing
- RISE: Recursive Improvement via Self-Extrapolating Policy Distillation
- Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe
- OISD: On-Policy Internal Self-Distillation of Language Models
- Privileged Information Distillation for Language Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks