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
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 ht of the frozen student and decides whether to correct the current step. If gϕ(st)=1 , CRA predicts logit residuals over the student’s top- K candidates and the next token is sampled from qθ , otherwise it is sampled from πS . The lower panel shows training, where the teacher provides qT and ℓ(st) 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, λ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
Hyperparameter
KD
OPD
EOPD
SRA
Training configuration
Optimizer
AdamW
AdamW
AdamW
AdamW
Learning rate
10−5
3×10−6
3×10−6
DG: 10−4 , CRA: 5×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 τ=0.1 .
Setting
Student
DG width
dc
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 η=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, η 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 transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.
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.
Ryan Brown, Chris Russell
Oxford Internet Institute, University of Oxford, Oxford, United Kingdom.
Distilling reasoning traces from strong large language models into smaller ones is a promising route to improve intelligence in resource-constrained settings. Existing approaches face a fundamental trade-off: offline distillation from teacher-generated traces provides high-quality, sample-efficient supervision but suffers from distributional drift: during training, the student model conditions on teacher-generated prefixes, whereas during inference the student autoregresses on self-generated prefixes, leading to compounding errors over long reasoning trajectories. Meanwhile, on-policy or self-distillation methods better match the student's inference-time distribution, but require costly online sampling and often produce low-quality traces in early training. We propose a principled offline reasoning distillation framework that preserves the efficiency and supervision quality of offline teacher-generated data while correcting teacher-student distribution drift. It adaptively emphasizes teacher supervision that is better aligned with the student's on-policy distribution. Evaluations on mathematical reasoning benchmarks of GSM8K, MATH, MATH500, and harder held-out competition-style tasks, including AMC, AIME, and OlympiadBench, show that our method improves reasoning accuracy over prior offline distillation algorithms and yields more stable reasoning traces while preserving instruction-following capabilities. Our work shows that lightweight, distribution-correction-aware training can substantially strengthen offline reasoning distillation without online rollouts.
Reasoning capabilities are critical for advancing Large Language Models, yet current approaches either require massive computational budgets or struggle to effectively distill reasoning to smaller models. Standard distillation methods rely on outcome-based rewards, failing to distinguish between sound reasoning and lucky guesses. We propose Collaborative Reasoning Distillation (CRD), a framework that enhances reasoning in compact models through three innovations: (1) interactive cross-feedback where teachers iteratively critique each other's reasoning, (2) fine-grained step-wise quality assessment capturing logical validity independent of final answers, and (3) coherence-aware step stitching that synthesizes complementary strengths. Students are trained via Reasoning Quality Optimization (RQO) with budget constraints. Our model, CRD-4B, achieves 97.3% on MATH-500 and 70.3% on AIME'25, surpassing baselines while using only 50K training examples, up to 12 times smaller than the datasets of comparable models.