SRA: Selective Residual Adaptation for LLM Reasoning Enhancement

Zhenxing Dou; Kun He; Mingyang Yi; Yong Liu

cs.LGOct 25, 2026

SRA: Selective Residual Adaptation for LLM Reasoning Enhancement

Authors: Zhenxing Dou, Kun He, Mingyang Yi, Yong Liu

Organizations: Gaoling School of Artificial Intelligence, Renmin University of China, Beijing, China · School of Information, Renmin University of China, Beijing, China

Abstract

Knowledge distillation transfers reasoning capabilities from large language models to lightweight students. However, distilled small models still struggle with multi-step reasoning. We observe that the distribution divergences between teacher and student models along reasoning paths are highly sparse, and target tokens almost always fall within the student top- $K$ candidate set. Based on this finding, unlike approaches updating student backbones or requiring runtime teacher computation, we propose Selective Residual Adaptation (SRA), a plug-and-play framework that internalizes localized corrections. Concretely, SRA employs a Discrepancy Gate to identify divergence positions alongside a Candidate Residual Adapter for top- $K$ logit calibration to correct the mismatch. Theoretically, we establish error bounds ensuring strict sequence level performance gains. Across mathematical benchmarks on Qwen3 models from 0.6B to 4B, SRA consistently enhances distilled student models, raising average Pass@16 by up to 4.37 percentage points while preserving general capabilities.

Figures & tables

Figure 1: (a) Answer accuracy with and without teacher sampling triggered by top-1 disagreement. The annotations show trigger rates, and the dashed line indicates independent teacher accuracy. (b) The fraction of the teacher’s top-

K

probability mass covered by the student’s top-

K

set at positions where their top-1 predictions differ. (c) The fraction of top-1 disagreements within the student’s top-

K

candidate set (

K=16

) missed by TV, forward KL, reverse KL, and student entropy when each measure selects the same number of positions. We select 256 questions from the GSM8K test split ( Cobbe et al., 2021 ) and compare Qwen3-0.6B-Base and 1.7B-Base and OPD students ( Yang et al., 2025 ) with Qwen3-8B as the teacher. Appendix A provides training and evaluation details.

Figure 2: Overview of SRA. The upper panel shows inference. DG reads the final hidden state $h_{t}$ of the frozen student and decides whether to correct the current step. If $g_{\phi}(s_{t})=1$ , CRA predicts logit residuals over the student’s top- $K$ candidates and the next token is sampled from $q_{\theta}$ , otherwise it is sampled from $\pi_{S}$ . The lower panel shows training, where the teacher provides $q_{T}$ and $\ell(s_{t})$ to supervise CRA and DG in two stages.

	AIME24		AIME25		AMC23		MATH-500		Average
Method	Avg@16	Pass@16	Avg@16	Pass@16	Avg@16	Pass@16	Avg@16	Pass@16	Avg@16	Pass@16
Qwen3-0.6B-Base
KD	1.46	13.33	0.21	3.33	16.72	62.50	36.91	76.40	13.82	38.89
w/ SRA	1.46	10.00	0.21	3.33	17.97	62.50	37.41	77.20	14.26 +0.44	38.26 -0.63
OPD	1.67	10.00	0.00	0.00	25.31	67.50	43.24	76.20	17.55	38.43
w/ SRA	2.08	10.00	0.83	6.67	26.56	67.50	47.70	80.80	19.29 +1.74	41.24 +2.81

Table 1: Main results of distilled Qwen3 models on mathematical reasoning benchmarks. Shaded rows use SRA. Superscripts show changes from the corresponding baseline in percentage points, with gains in green and drops in red. Boldface and underlining indicate the best and second best results among student models of the same size.

Benchmark	Method	Student		w/ SRA
		Avg@16	Pass@16	Avg@16	Pass@16
AIME24	Base	3.54	26.67	10.42	23.33
	KD	12.92	30.00	15.00	36.67
	OPD	15.83	36.67	16.04	40.00
	EOPD	16.67	40.00	17.71	43.33
AIME25	Base	2.71	16.67	7.50	26.67

Table 2: Comparison of SRA and continued distillation on Qwen3-4B. Student extends distillation by two epochs, shaded applies SRA after two epochs, and Base applies SRA directly to the base model.

Setting	Avg@16	Pass@16
Full SRA (L24, $\lambda_{R}=0.1$ )	41.45	67.92
Final hidden state
w/o final hidden state	40.88	66.16
Intermediate hidden state
w/o intermediate state	40.41	63.28
L12	41.03	64.60

Table 3: Ablations on Qwen3-4B OPD across four mathematical reasoning benchmarks against SRA. Full results and configurations are in Appendix D.1 .

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Training configuration
Hyperparameter	KD	OPD	EOPD	SRA
Optimizer	AdamW	AdamW	AdamW	AdamW
Learning rate	$10^{-5}$	$3\times 10^{-6}$	$3\times 10^{-6}$	DG: $10^{-4}$ , CRA: $5\times 10^{-5}$
Learning rate schedule	Cosine	Cosine	Cosine	Constant
Training epochs	2	2	2	2
Batch size	128	128	128	128

Appendix

Table 4: Training hyperparameters for KD, OPD, EOPD, and SRA.

Figure 3: TV threshold sensitivity on GSM8K with $K=16$ . (a) Disagreement recall. (b) Proportion of selected positions. Vertical lines mark $\tau=0.1$ .

Setting	Student	DG width	$d_{c}$	Heads	FFN width	Intermediate layer
Distilled students	0.6B	128	768	6	3,072	19
	1.7B	256	1,280	10	5,120	19
	4B	320	2,048	16	8,192	24
Cross-mode	0.6B	128	1,024	8	4,096	24
	1.7B	256	2,048	16	8,192	24
	4B	320	2,048	16	8,192	24

Appendix

Table 5: SRA module dimensions. DG width is the hidden width of the gate MLP. FFN width is the intermediate dimension of the feedforward network in each CRA Transformer layer.

Figure 4: DG F1 on the DAPO validation split after the first SRA epoch with Qwen3-4B OPD. TV and Top-1 denote the two supervision targets. Each curve evaluates a separately trained gate against its own target, with the highest F1 at

\eta=0.8

Benchmark	Method	Student		w/ SRA
		Avg@16	Pass@16	Avg@16	Pass@16
GPQA- Diamond	KD	31.60	90.40	32.70	91.41
	OPD	31.31	89.39	33.68	90.91
	EOPD	32.23	87.88	32.95	90.40
HumanEval	KD	75.91	94.51	76.49	94.51
	OPD	71.34	95.73	71.65	95.73

Appendix

Table 6: Out-of-domain results of Qwen3-4B distilled students. Shaded columns use SRA.

Figure 5: Forward KL within student top-16 candidates before and after SRA correction on Qwen3-4B. Reported values average DG selected positions across four mathematical benchmarks.

	Avg@16 / Pass@16 (%)				Average (%)
Setting	AIME24	AIME25	AMC23	MATH-500	Avg@16	Pass@16
Final hidden state
w/o final hidden state	15.21 / 40.00	15.00 / 43.33	55.16 / 87.50	78.14 / 93.80	40.88	66.16
SRA	16.04 / 40.00	14.38 / 46.67	56.88 / 90.00	78.49 / 95.00	41.45	67.92
Intermediate hidden state
w/o intermediate state	13.33 / 30.00	16.04 / 43.33	54.69 / 85.00	77.59 / 94.80	40.41	63.28

Appendix

Table 7: Complete ablations and adaptation comparisons on Qwen3-4B OPD. Shaded rows repeat the same SRA reference. Bold and underline indicate the best and second best scores within each group. For candidate size ablations,

\eta

is the DG threshold that achieves the highest F1 score on the validation set.

	Avg@16 / Pass@16 (%)
Setting	AIME24	AIME25	AMC23	MATH-500	Average	Recovery
Qwen3-0.6B
non-thinking	1.25 / 3.33	2.50 / 16.67	23.75 / 62.50	43.16 / 81.60	17.67 / 41.02	–
w/ SRA	3.33 / 13.33	6.04 / 23.33	29.22 / 72.50	53.74 / 88.60	23.08 / 49.44	62.94 / 79.05
thinking	11.67 / 43.33	14.58 / 33.33	48.12 / 82.50	72.29 / 91.00	36.67 / 62.54	–
Qwen3-1.7B

Appendix

Table 8: Cross-mode results for Qwen3. Shaded rows apply SRA in non-thinking mode. Bold indicates best performance and underline indicates second best within each size. Recovery is the fraction of thinking performance attained by SRA.

	Avg@16 / Pass@16 (%)
Method	AIME24	AIME25	AMC23	MATH-500	Average
Initial Instruct (3B)	1.25 / 6.67	0.00 / 0.00	9.22 / 55.00	27.91 / 71.80	9.60 / 33.37
w/ SRA	1.67 / 13.33	0.21 / 3.33	10.78 / 55.00	28.31 / 73.40	10.24 / 36.27
KD	0.42 / 6.67	0.00 / 0.00	8.12 / 47.50	24.40 / 63.60	8.24 / 29.44
OPD	0.42 / 6.67	0.21 / 3.33	10.00 / 50.00	32.10 / 69.60	10.68 / 32.40
EOPD	0.00 / 0.00	0.42 / 6.67	8.91 / 42.50	27.09 / 65.00	9.10 / 28.54

Appendix

Table 9: Performance comparison of Llama-3.2-3B-Instruct and distillation baselines with an 8B teacher reference. The shaded row applies SRA to the initial Instruct model. Bold indicates best performance and underline indicates second best among students.

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 514 papers

Task-Specific Knowledge Distillation via Intermediate Probes
Feb 18, 2026Ryan Brown, Chris RussellLanguage Model Distillation Knowledge Distillation
Oxford Internet Institute, University of Oxford, Oxford, United Kingdom.
Distribution Corrected Offline Data Distillation for Large Language Models
May 13, 2026Yumeng Zhang, Zhengbang Yang, Yevin Nikhel Goonatilake +1CoT Distillation Online Knowledge Distillation
George Mason University, Fairfax, VA, USA
Collaborative Reasoning Distillation via Cross-Feedback and Coherent Curation
Oct 7, 2026Taehoon Kim, Seunggeun Cho, Dongsu HanReasoning Distillation LLM Reasoning
KAIST

Feb 18, 2026J/K move · Enter open · S save

Task-Specific Knowledge Distillation via Intermediate Probes

Ryan Brown, Chris Russell

Oxford Internet Institute, University of Oxford, Oxford, United Kingdom.

Language Model Distillation Knowledge Distillation

Open paper AskPDF arXiv

Knowledge distillation from large language models (LLMs) assumes that the teacher's output distribution is a high-quality training signal. On reasoning tasks, this assumption is frequently violated. A model's intermediate representations may encode the correct answer, yet this information is lost or distorted through the vocabulary projection, where prompt formatting and answer-token choices creates brittle, noisy outputs. We introduce \method{}, a distillation framework that bypasses this bottleneck by training lightweight probes on frozen teacher hidden states and using the probe's predictions, rather than output logits, as supervision for student training. This simple change yields consistent improvements across four reasoning benchmarks (AQuA-RAT, ARC Easy/Challenge, and MMLU), with gains most pronounced under limited data. Probes trained on intermediate representations provide cleaner labels than the teacher's own outputs, effectively denoising the distillation signal. \method{} requires no architectural changes to student or teacher, is architecture-agnostic, and adds minimal compute since probe training is cheap and teacher representations can be cached. By exploiting internal representations, \method{} enables practitioners to extract more value from large teacher models without additional training data or architectural complexity.

Abstract

Language Model Distillation

Explore similar work

Task-Specific Knowledge Distillation via Intermediate Probes

Distribution Corrected Offline Data Distillation for Large Language Models

Collaborative Reasoning Distillation via Cross-Feedback and Coherent Curation