MASS: Multi-mask Selective Self-Distillation for Leakage-Resilient Reasoning

Xiuqi Liu; Oveys Delafrooz Noroozi; Upamanyu Madhow

cs.LGOct 28, 2026

MASS: Multi-mask Selective Self-Distillation for Leakage-Resilient Reasoning

Authors: Xiuqi Liu, Oveys Delafrooz Noroozi, Upamanyu Madhow

Organizations: Chongqing Bashu Secondary School · University of California, Santa Barbara

Abstract

On-policy self-distillation enables the same model to act as student and teacher, with dense token-level supervision obtained by the teacher being given access to privileged information. However, conditioning the teacher on a single reference can concentrate its predictions on one reasoning trajectory and limit transfer to the student. We propose Multi-Mask Selective Self-Distillation (MASS), which approximates the unavailable marginal over diverse reference realizations using multiple independently masked views of one verified solution. MASS aggregates these complementary teacher predictions, selectively transfers reliable token-level corrections, and adaptively focuses masking on influential reference spans. Across mathematical reasoning benchmarks and model families, MASS consistently outperforms strong self-distillation baselines.

Figures & tables

Figure 1: PI Bias Scores for the in-context solution and a distinct correct solution under the full-reference teacher, a single masked view, and the arithmetic mean of multiple masked views.

Figure 2: Overview of MASS. Multiple masked reference provide complementary teacher predictions over the same student trajectory and are averaged into a marginal target. MASS selectively distills this target using reliability-aware weighting. In parallel, anchor-relative coalition utilities are decomposed into span-level attributions, which adapt future masking probabilities.

Model	Method	AIME24	AIME25	HMMT25	Avg.
Qwen3-1.7B	Base	51.39	36.67	23.06	37.04
	OPSD	55.56	41.11	26.94	41.20
	AVSD	59.17	45.28	29.72	44.72
	RLSD	58.89	45.83	30.00	44.91
	MASS	60.83	47.50	30.83	46.39
Qwen3-4B	Base	75.00	66.39	42.22	61.20

Table 1: Performance comparison on mathematical reasoning benchmarks. We report Avg@12 accuracy on AIME 2024, AIME 2025, and HMMT 2025. MASS consistently achieves the strongest average performance across all evaluated model settings.

Figure 3: Performance Dynamics over 250 Training Steps. We compare MASS with OPSD and RLSD on Qwen3-4B reporting Avg@12 accuracy on AIME25 within 250 training steps. OPSD peaks early, while RLSD and MASS sustain improvement throughout training.

Model	Target	Acc.
Qwen3-4B	Masked avg. (Base)	73.89
	Masked avg. $+$ anchor	72.50
OLMo-7B	Masked avg. (Base)	75.00
	Masked avg. $+$ anchor	73.89

Table 2: Ablation on the anchor teacher. AIME 2025 Avg@12 accuracy is reported.

Panel of Figure 4 — Figure 4: Reasoning behavior during self-distillation. We measured average response length and epistemic token usage for (a) Qwen3-8B and (b) OLMo-7B on AIME25 evaluation.

Figure 5: Sensitivity to the number of teacher routes. Increasing the number of routes improves AIME 2025 Avg@12 accuracy on Qwen3-1.7B.

Figure 6: Sensitivity to masking strength. On Qwen3-4B, a moderate masking ratio gives the best AIME 2025 Avg@12 accuracy.

Variant	Weighting	Adaptive Masking	AIME25
Neither	✗	✗	70.83
Weighting only	✓	✗	72.22
Adaptive only	✗	✓	72.78
Full MASS	✓	✓	73.89

Table 3: Ablation of reliability weighting and adaptive masking. Both components improve AIME 2025 performance on Qwen3-4B, with the full configuration achieving the strongest result.

Model	Method	GPQA	BBH
Qwen3-8B	Base	57.7	84.4
	OPSD	54.8	82.9
	AVSD	56.6	83.7
	MASS	58.3	84.6
OLMo3-7B-Think	Base	46.2	86.6
	OPSD	43.9	85.0

Table 4: OOD evaluation. GPQA Avg@8 and BBH Avg@1 accuracy are reported.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Model	$\epsilon=0.01$	$\epsilon=0.03$	$\epsilon=0.05$	$\epsilon=0.10$
Qwen3-8B	79.08	79.44	79.17	78.61
OLMo-7B	74.62	75.00	74.72	73.89

Appendix

Table 5: Sensitivity analysis of the masking probability clipping We report Avg@12 performance on AIME 2025. All experiments use the same MASS training configuration except for the clipping threshold.

Distillation Objective	AIME24	AIME25
Backward KL	79.28	72.50
Forward KL (MASS)	81.11	73.89

Appendix

Table 6: Effect of KL divergence direction on Qwen3-4B. Results are reported as Avg@12 accuracy on AIME 2024 and AIME 2025.

Figure 7: Teacher prompts in OPSD and MASS. OPSD uses the complete reference solution, while MASS constructs multiple teacher views from independently masked versions of the same reference; one example is shown. Labels [S1] – [S5] and the surrounding braces are added only for visualization to clarify the span boundaries and are not included in the actual prompts. During training, selected spans are directly replaced by [MASK] .

Method	FLOPs / Step	Relative to OPSD
OPSD	$1.31\times 10^{15}$	$1.00\times$
MASS	$3.41\times 10^{15}$	$2.60\times$
RLSD	$8.39\times 10^{15}$	$6.40\times$

Appendix

Table 7: Estimated training computation per optimization step. FLOPs are normalized by OPSD.

Figure 8: Normalized span-level masking attribution Heatmap

Hyperparameter	Value
Number of teacher routes ( $K$ )	8
Masking ratio ( $\rho$ )	0.3
Anchor coefficient ( $\lambda$ )	0.5
Weight regularization ( $\varepsilon$ )	$1\times 10^{-6}$
Shapley regression regularization ( $\tau$ )	0.1
EMA smoothing ( $\gamma_{\mathrm{EMA}}$ )	0.9

Appendix

Table 8: Method hyperparameters used in MASS.

Training Parameter	Value
Learning rate	$5\times 10^{-6}$
Effective batch size	32
Maximum sequence length	1024 tokens
Training steps	250
Number of Generations Per Prompt	1
Sampling temperature	1.1

Appendix

Table 9: Training configuration used in MASS.

Hyperparameter	Value
Number of teacher routes ( $K$ )	8
Masking ratio ( $\rho$ )	0.3
Anchor coefficient ( $\lambda$ )	0.5
Weight regularization ( $\varepsilon$ )	$1\times 10^{-6}$
Shapley regression regularization ( $\tau$ )	0.1
EMA smoothing ( $\gamma_{\mathrm{EMA}}$ )	0.9

Appendix

Table 10: Method hyperparameters used in MASS.

Evaluation Parameter	Value
Maximum new tokens	38912
Thinking mode	Enabled
Sampling temperature	1.0
Top- $p$	0.95
Top- $k$	-1
Min- $p$	0.0

Appendix

Table 11: Evaluation configuration used in MASS.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 522 papers

Teach Yourself Where to Look: On-Policy Attention Self-Distillation for Reasoning
Sep 27, 2026Safaeid Hossain Arib, Rabeya Akter, Ismam Nur Swapnil +3Numerical Reasoning in Language Models Attention-Guided Distillation
ACI PLC, Bangladesh · University of Dhaka, Bangladesh · BRAC University, Bangladesh +2
Anti-Self-Distillation for Reasoning RL via Pointwise Mutual Information
May 12, 2026Guobin Shen, Xiang Cheng, Chenxiao Zhao +4On-Policy Self-Distillation RL for Language Model Reasoning
1Xiaohongshu Inc. · Institute of Automation, Chinese Academy of Sciences
E $^2$ -OPSD: Taming Entropy Overshoot in On-Policy Self-Distillation
Oct 4, 2026Yifei Liu, Minghao Fang, Xinyu Gu +6On-Policy Self-Distillation Mathematical Reasoning
The Chinese University of Hong Kong · Zhejiang University · Shanghai Artificial Intelligence Laboratory +5

Sep 27, 2026J/K move · Enter open · S save

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

Safaeid Hossain Arib, Rabeya Akter, Ismam Nur Swapnil, Md. Faiyaz Abdullah Sayeedi+2

ACI PLC, Bangladesh · University of Dhaka, Bangladesh · BRAC University, Bangladesh+2

Numerical Reasoning in Language Models Attention-Guided Distillation On-Policy Self-Distillation Knowledge Distillation

Open paper AskPDF arXiv

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.

Abstract

Language Model Distillation

Explore similar work

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

Anti-Self-Distillation for Reasoning RL via Pointwise Mutual Information

E2^22-OPSD: Taming Entropy Overshoot in On-Policy Self-Distillation

E $^2$ -OPSD: Taming Entropy Overshoot in On-Policy Self-Distillation