PaperlayerView exact revision on arXiv

Source-linked AI summary

Teach Yourself Where to Look: On-Policy Attention Self-Distillation for Reasoning

Safaeid Hossain Arib, Rabeya Akter, Ismam Nur Swapnil, Md. Faiyaz Abdullah Sayeedi, Tasnim Mohiuddin, Md Mofijul Islam

arXiv:2609.33200v2cs.LGcs.CL

TL;DR

On-policy self-distillation provides dense, solution-conditioned token guidance on a student's own reasoning, but leaves the teacher's attention guidance unused and faces a mismatch between teacher-visible and student-visible context. OPASD adds projected attention supervision to address that mismatch. Across three model scales and four mathematics benchmarks, it improves accuracy over token-only OPSD while also reducing training cost.

  • Problem

    OPSD guides what the student predicts with solution-conditioned token distributions, but does not directly transfer where the teacher attends, and teacher attention includes solution positions the student cannot observe.

    Source: §1 · pp. 1–2

  • Method

    OPASD adds solution-conditioned attention supervision by projecting teacher attention onto student-visible positions and renormalizing before aligning it alongside token distributions.

    Source: §1 · p. 2

  • Results

    Across three Qwen3 model sizes and four competition-level mathematics benchmarks, OPASD improves average accuracy over token-only OPSD by 4.98 to 8.40 percentage points.

    Source: §1 · p. 2

  • Takeaways & Limitations

    OPASD provides a complementary supervision signal and is reported to improve accuracy, training stability, and efficiency over token-only OPSD in the evaluated setting.

    Source: §1 · p. 2, §5 · p. 9

    Evaluation covers three Qwen3 model sizes, and the privileged teacher remains fixed; other model families and teacher-update strategies remain to be studied.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

OPSD provides dense, solution-conditioned next-token guidance on student-generated reasoning but does not directly guide attention to earlier context. OPASD adds projected attention supervision, and the paper reports higher accuracy alongside lower training cost.

1 INTRODUCTION

  • OPSD transfers privileged guidance through next-token distributions, but does not directly guide where the student attends in multi-step reasoning.Later reasoning steps can depend on constraints and intermediate results from earlier steps.Source: §1 · p. 1
  • Directly aligning teacher and student attention is invalid because the teacher can attend to verified-solution positions absent from the student's context.Simply discarding those positions leaves the remaining teacher attention unnormalized.Source: §1 · pp. 1–2
  • OPASD projects teacher attention onto student-visible positions and renormalizes it, retaining solution-conditioned guidance while making the target comparable to student attention.It combines this attention supervision with the original token-level objective.Source: §1 · p. 2, Figure 1 · p. 2
  • Across three Qwen3 model scales and four competition-level mathematics benchmarks, OPASD improves average accuracy over token-only OPSD by 4.98 to 8.40 percentage points.It also reduces generated rollout tokens by 73.9% and estimated model compute by 72.6%, while training 1.53× faster.Source: §1 · p. 2

2 RELATED WORK

Related work develops on-policy distillation, self-distillation, representation alignment, and attention supervision. OPASD differs by aligning solution-conditioned teacher attention over student-generated reasoning despite the teacher and student having different context supports.

On-Policy Distillation

  • On-policy distillation queries a teacher at student-generated prefixes, while OPSD uses the same model as student and privileged teacher under different conditioning contexts.The student samples from the problem alone; the teacher also conditions on the verified solution.Source: §2 · pp. 2–3

Representation Distillation

  • Representation distillation aligns teacher and student features or states, but OPASD explicitly aligns attention distributions over preceding context.Some recent representation work aligns hidden states along student-generated trajectories.Source: §2 · p. 3

Attention Distillation

  • Prior attention-distillation work guides reasoning or explores multimodal and internal attention supervision, but does not address privileged on-policy self-distillation with different context supports.OPASD projects solution-conditioned teacher attention onto student-visible positions before alignment.Source: §2 · p. 3

3 ON-POLICY ATTENTION SELF-DISTILLATION

OPASD extends on-policy self-distillation with solution-conditioned attention supervision alongside token-level guidance. Its projection makes the teacher's attention target comparable to student attention, and the joint objective trains the student using both signals.

3 ON-POLICY ATTENTION SELF-DISTILLATION

  • OPASD combines the original token-level OPSD objective with an attention-level objective over the same student-generated reasoning trajectory.The privileged teacher evaluates that trajectory under its solution-conditioned context.Source: §3 · p. 3, §3.4 · p. 5

3.1 ON-POLICY SELF-DISTILLATION

  • The student samples a trajectory from the problem alone, while the frozen privileged teacher evaluates that same trajectory conditioned on the problem and reference solution.Gradients pass through the student, while the teacher supplies a detached target.Source: §3.1 · pp. 3–4, Figure 2 · p. 4

3.2 SUPPORT MISMATCH IN PRIVILEGED ATTENTION

  • Teacher and student attention distributions have different supports because only the teacher can attend to reference-solution positions.Those teacher-only positions prevent direct attention alignment.Source: §3.2 · p. 4

3.3 STUDENT-SUPPORT PROJECTION

  • OPASD restricts teacher attention to student-visible positions and renormalizes the remaining mass to form the projected target.The projection preserves the teacher's relative attention over visible positions, and the target remains solution-conditioned.Source: §3.3 · p. 5

3.4 JOINT DISTILLATION

  • The joint objective distills projected teacher attention with generalized Jensen-Shannon divergence and combines that loss with token-level OPSD loss.The attention-loss coefficient controls the strength of attention supervision; only the student receives gradients.Source: §3.4 · p. 5

4 EXPERIMENTS

Across three Qwen3 model sizes and four competition-level mathematics benchmarks, OPASD outperforms token-only OPSD. The experiments also examine attention-distillation choices, reasoning behavior, and training efficiency.

4.1 EXPERIMENTAL SETUP

  • Experiments use instruction-tuned Qwen3-1.7B, 4B, and 8B, train on up to 30K OpenThoughts problem-solution pairs, and evaluate on four AIME/HMMT benchmarks.OPSD and OPASD share the dataset, initialization, and training budget; the privileged teacher is fixed during training.Source: §4.1 · p. 6

4.2 COMPARISON OF ON-POLICY SELF-DISTILLATION METHODS

  • 4.98, 8.40, and 5.97 points: OPASD raises average scores over OPSD on Qwen3-1.7B, 4B, and 8B, respectively, and wins all four benchmarks at every scale.The gains remain consistent across sampling budgets, training seeds, and test-time generation budgets.Source: §4.2 · p. 6

4.3.1 ATTENTION DIVERGENCE OBJECTIVE

  • 32.85 average accuracy: JSD exceeds reverse KL at 32.36 and forward KL at 31.73, balancing broader context coverage with focus on relevant positions.The comparison is on Qwen3-1.7B with other settings fixed.Source: Table 2 · p. 7, §4.3.1 · p. 7

4.3.2 STRENGTH OF ATTENTION SUPERVISION

  • 34.02 average score: attention-loss weight λattn = 0.5 outperforms weights 1.0 and 2.0, which score 32.85 and 31.82.The token-loss coefficient is fixed at 1; the authors report larger attention weights reduce the token-level objective’s influence.Source: Table 3 · p. 7, §4.3.2 · p. 7

4.3.3 ATTENTION DISTILLATION DEPTH

  • 34.02 average score: distilling only the final layer outperforms using the final two or three layers, which score 32.54 and 31.04.The authors attribute the advantage to final-layer attention acting on representations refined by preceding layers and being closer to next-token prediction.Source: Table 4 · p. 7, §4.3.3 · p. 7

4.3.4 QUALITATIVE ANALYSIS

  • On an AIME24 number-theory problem, OPASD carries broader candidate analysis through to the correct minimum, unlike the base model and OPSD.The authors describe its attention supervision as guidance on earlier calculations that helps evaluate later candidates more consistently.Source: §4.3.4 · pp. 7–8, Figure 3 · p. 8

4.3.5 PERFORMANCE AND BEHAVIORAL DYNAMICS

  • OPASD maintains stable accuracy and concise responses, while OPSD’s accuracy declines as responses lengthen and uncertainty and reconsideration markers rise.Across ten epistemic markers, mean counts rise 4.4% for OPASD versus 70.9% for OPSD; OPASD’s validation Avg@12 varies by 2.22 points versus OPSD’s 10.76-point span.Source: Figure 4 · p. 8, §4.3.5 · pp. 8–9

4.3.6 EFFICIENCY

  • 1.53× faster: over 200 Qwen3-4B steps, OPASD raises Avg@12 from 42.36 to 50.76 and reduces response tokens by 73.9%, estimated compute by 72.6%, and peak GPU memory by 11.9%.Wall-clock time falls from 8.04 to 5.27 hours; Figure 5 reports Pareto dominance over OPSD across accuracy, time, generated tokens, and peak memory.Source: §4.3.6 · p. 9

5 CONCLUSION

The conclusion presents OPASD as complementary to token-level self-distillation, with consistent accuracy gains and more stable, efficient training. Its evaluation and fixed-teacher setting leave further assessments open.

5 CONCLUSION

  • Across three Qwen3 sizes, OPASD improves accuracy over OPSD and is reported to stabilize training, shorten responses, and lower training cost.The paper concludes that attention supervision complements token-level self-distillation.Source: §5 · p. 9, Figure 5 · p. 9
  • Evaluation beyond Qwen3 and attention supervision under other teacher-update strategies remain to be studied; the teacher was fixed during training.Source: §5 · p. 9

AI USE STATEMENT

The authors report using generative AI tools for experimental-design feedback, analysis, literature searches, code, figures, and manuscript drafting and editing.

AI USE STATEMENT

  • The authors say they reviewed all AI-assisted material and take responsibility for the paper’s methods, results, claims, code, and figures.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

The paper evaluates OPASD on mathematical reasoning using existing datasets and benchmarks, without human subjects or newly collected personal data.

ETHICS STATEMENT

  • The authors state that deployment should assess reliability and misuse; applying OPASD elsewhere would inherit risks from the underlying models and training data.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

OPASD provides stronger individual responses and higher validation accuracy while using fewer generated tokens than OPSD. The section also outlines reproducibility materials and examines performance across sampling budgets.

REPRODUCIBILITY STATEMENT

  • The reproducibility statement says Section 3 describes OPASD’s objective and attention projection, while Section 4 reports experimental setups and ablations.It also states that code and experiment configurations will be publicly released.Source: REPRODUCIBILITY STATEMENT · p. 10

A APPENDIX

  • The appendix provides additional analyses, ablations, implementation details, and theoretical results supporting the main paper.Source: A APPENDIX · p. 14

B PERFORMANCE ACROSS SAMPLING BUDGETS

  • At the final checkpoint, OPASD achieved average Pass@1 of 48.33 versus 30.21 for OPSD, with higher Pass@1 across all four benchmarks.OPASD also improved final average Pass@12 from 56.51 to 63.82; gains appeared on AIME24, AIME25, and HMMT25, while performance was comparable on AIME26.Source: B PERFORMANCE ACROSS SAMPLING BUDGETS · p. 14, Figure 6 · p. 14

C TOKEN LEARNING EFFICIENCY

  • The supplied passage introduces evaluation of whether OPASD obtains stronger validation performance from the same generation budget, but does not include the complete results.Figure 7 reports Avg@12 accuracy against cumulative response tokens generated by training rollouts.Source: C TOKEN LEARNING EFFICIENCY · p. 14
  • 15.58 million response tokens versus 59.73 million, a 73.9% reduction, accompanied higher validation accuracy for OPASD across all four benchmarks.Across the 200-step runs, OPASD maintained higher validation accuracy while generating shorter, more productive reasoning trajectories.Source: C TOKEN LEARNING EFFICIENCY · p. 15, Figure 7 · p. 15

D EPISTEMIC VERBALIZATION DYNAMICS

OPASD maintains lower and steadier epistemic-marker density than OPSD, including within correct and incorrect responses. Its attention-alignment loss also declines as training proceeds.

D EPISTEMIC VERBALIZATION DYNAMICS

  • Epistemic verbalization denotes uncertainty, reconsideration, or self-verification expressions in a reasoning trace, measured here using ten markers per 1,000 tokens.The markers are treated as a behavioral proxy, not a direct measure of internal uncertainty.Source: D EPISTEMIC VERBALIZATION DYNAMICS · p. 15
  • 12–14 markers per 1,000 tokens characterized OPASD throughout training, while OPSD ranged from 14 to 23 and generally increased.OPSD’s responses also grew longer as validation accuracy declined, whereas OPASD remained stable on both measures.Source: D EPISTEMIC VERBALIZATION DYNAMICS · p. 15, Figure 8 · p. 15
  • For correct responses, OPASD produced 12.1 markers per 1,000 tokens versus OPSD’s 16.4; for incorrect responses, the values were 13.7 versus 20.3.Incorrect responses had more markers under both methods, but OPASD’s lower density within each correctness group was not simply a consequence of higher accuracy.Source: D EPISTEMIC VERBALIZATION DYNAMICS · p. 16, Figure 9 · p. 16
  • At the final checkpoint, OPSD used “wait” 7.42 times per 1,000 tokens versus OPASD’s 3.58, while OPASD used “check” 1.21 times versus OPSD’s 0.46.OPASD’s lower use of uncertainty and alternative-search markers did not uniformly reduce epistemic language.Source: D EPISTEMIC VERBALIZATION DYNAMICS · p. 16

E ATTENTION ALIGNMENT DYNAMICS

  • The JSD between student attention and projected solution-conditioned teacher attention fell from 0.0055 over the first 20 steps to 0.0029 over the final 20.The teacher target is recomputed for each new student-generated trajectory.Source: E ATTENTION ALIGNMENT DYNAMICS · p. 16

F ROBUSTNESS ACROSS RANDOM SEEDS

OPASD’s advantage persists across test-time generation budgets and alternative training controls. The supplied results also report seed stability and favor student-support projection, per-token clipping, and combined token-and-attention distillation.

F ROBUSTNESS ACROSS RANDOM SEEDS

  • Across three Qwen3-4B runs, OPASD achieved an average score of 50.88 ± 0.14, with a 95% confidence interval of [50.53, 51.24].The runs used seeds 1, 42, and 123 with the same data, hyperparameters, compute budget, and configuration.Source: F ROBUSTNESS ACROSS RANDOM SEEDS · p. 16, Table 5 · p. 18

G TEST-TIME GENERATION BUDGET

H HANDLING PRIVILEGED ATTENTION

  • Student-support projection scored 34.02 on average, versus 31.39 for redistribution and 32.15 for shuffled redistribution.Projection restricts teacher attention to student-visible positions and renormalizes it; the passage attributes its advantage to preserving relative attention weights among usable positions.Source: Table 6 · p. 18, H HANDLING PRIVILEGED ATTENTION · p. 18

I TOKEN-LOSS CLIPPING

  • In the token-loss clipping comparison, per-token clipping achieved the highest average score of 32.85, versus 32.22 without clipping and 30.83 with point-wise clipping.It preserves the full signal while limiting positions with very large divergence.Source: Table 7 · p. 19, I TOKEN-LOSS CLIPPING · p. 19

J COMPARISON OF DISTILLATION OBJECTIVES

K CONTROLLING FOR TOKEN-LOSS CLIPPING

  • With per-token clipping held fixed, OPASD scored 34.02 and outperformed token-only OPSD on all four benchmarks; the token-only control scored 29.09 versus 29.04 for original OPSD.The reported advantage therefore persists when clipping is matched.Source: K CONTROLLING FOR TOKEN-LOSS CLIPPING · p. 19, Table 9 · p. 19

L ADDITIONAL QUALITATIVE RESULTS

Two worked mathematics examples contrast the base model, token-only OPSD, and OPASD, illustrating whether setup constraints carry through to the final calculation. The supplied passages also specify a shared token-level objective and an additional OPASD attention objective.

L ADDITIONAL QUALITATIVE RESULTS

  • In the AIME25 probability and AIME26 geometry examples, OPASD carries setup constraints into the final calculation where the base model and OPSD make errors.OPASD accounts for a shared divisor in the probability problem and preserves a nonconvex configuration in the geometry problem.Source: L ADDITIONAL QUALITATIVE RESULTS · p. 20

M HYPERPARAMETERS

  • Both methods use a Thinking-Mode-off student and Thinking-Mode-on teacher; OPASD additionally applies attention distillation with λattn = 0.5.OPASD otherwise uses the same token-level distillation objective as OPSD.Source: M HYPERPARAMETERS · p. 20

N NOTATION SUMMARY

O THEORETICAL ANALYSIS

The appendix characterizes OPASD’s student-support projection and proves it is the unique best approximation to privileged-teacher attention under the stated divergence. It also specifies the attention supports and divergence used in that analysis, and illustrates a geometry problem where OPASD reaches the correct count while the base model and OPSD do not.

O THEORETICAL ANALYSIS

  • On the AIME26 geometry problem, the base model and OPSD derive incorrect constraints, while OPASD reaches the correct count.Highlighted excerpts mark the decisive reasoning steps.Source: Figure 14 · p. 21

O.1 PRELIMINARIES

  • The analysis defines the projected target as teacher attention conditioned on attending to student-visible positions, preserving relative attention among those positions.With a nonempty student-visible support and softmax attention, the projection is well defined.Source: §O.1 · p. 23
  • The preliminaries define the generalized Jensen–Shannon divergence and establish that it is nonnegative, with equality only for identical distributions.At β = 1/2, the divergence reduces to Jensen–Shannon divergence.Source: §O.1 · p. 23

O.2 OPTIMALITY OF STUDENT-SUPPORT PROJECTION

  • Among student-visible attention distributions, the projected teacher attention uniquely minimizes generalized Jensen–Shannon divergence to the full teacher distribution.The minimum is zero only when the teacher assigns no attention to unavailable positions, and it increases with the teacher’s mass on those positions.Source: §O.2 · pp. 23–24
Loading 2609.33200v2…