Learning to Focus: Causal Attention Distillation via Gradient‐Guided Token Pruning
Authors: Yiju Guo
Wenkai Yang
Zexu Sun
Ning Ding
Zhiyuan Liu
Yankai Lin
Venue: NeurIPS 2025 poster ()
Keywords: Knowledge Distillation
Spurious Correlations
Long-Context Reasoning
Model Interpretability
Large Language Models
Primary area: deep_learning
TL;DR: We propose Learning to Focus (LeaF), which identifies and masks confounding tokens via gradient‐based comparisons, thereby improving long‐context reasoning accuracy and interpretability in large language models.

Abstract: Large language models (LLMs) have demonstrated significant improvements in contextual understanding. However, their ability to attend to truly critical information during long-context reasoning and generation still falls behind the pace. Specifically, our preliminary experiments reveal that certain distracting patterns can misdirect the model’s attention during inference, and removing these patterns substantially improves reasoning accuracy and generation quality. We attribute this phenomenon to spurious correlations in the training data, which obstruct the model’s capacity to infer authentic causal instruction–response relationships. This phenomenon may induce redundant reasoning processes, potentially resulting in significant inference overhead and, more critically, the generation of erroneous or suboptimal responses. To mitigate this, we introduce a two-stage framework called Learning to Focus (LeaF) leveraging intervention-based inference to disentangle confounding factors. In the first stage, LeaF employs gradient-based comparisons with an advanced teacher to automatically identify confounding tokens based on causal relationships in the training corpus. Then, in the second stage, it prunes these tokens during distillation to enact intervention, aligning the student’s attention with the teacher’s focus distribution on truly critical context tokens. Experimental results demonstrate that LeaF not only achieves an absolute improvement in various mathematical reasoning, code generation and multi-hop question answering benchmarks but also effectively suppresses attention to confounding tokens during inference, yielding a more interpretable and reliable reasoning model.

Learning to Focus: Causal Attention Distillation via Gradient-Guided Token Pruning
Yiju Guo Wenkai Yang", Zexu Sun Ning Ding', Zhiyuan Liu", Yankai Lin Pe
Gaoling School of Artificial Intelligence, Renmin University of China Department of Computer Science and Technology, Tsinghua University Baidu Inc.
Byyijuguo, yankailin)@ruc eruc.euu. edu. cn
Abstract
Large language models (LLMs) have demonstrated significant improvements contextual understanding. However, their ability to attend to truly critical infor- mation during long-conteds reasoning and generation still falls behind the pace. Specifically, our preliminary experiments reveal that certain distracting patterns can misdirect the model's 's attention during inference, and removing these patterns substantially improves reasoning accuracy and generation quality. We attribute this phenomenon to spurious correlations in the training data, which obstruct the model's capacity infer authentic causal instruction-response relationships. This phenomenon may induce redundant reasoning processes, potentially resulting in significant inference overhead and, more critically, the generation of erroncous suboptimal responses. To mitigate this, we introduce a two-stage framework called Learning to Focus (LeaF) leveraging intervention-basee based inference to disentangle confounding factors. In the first stage, LeaF employs gradient-based comparisons with an advanced teacher to automatically identify confounding tokens based on causal relationships in the training corpus. Then, in the second stage, it prunes these tokens during distillation to enact intervention, aligning the student's attention with the teacher's 's focus distribution on truly critical context tokens. Experimental results demonstrate that LeaF not only achieves an absolute improvement in vari- ous mathematical reasoning, code generation and multi-hop question answering benchmarks but also effectively suppresses attention to confounding tokens during inference, yielding more interpretable and reliable reasoning model.
Introduction
Large language models (LLMs) have achieved remarkable success in natural language process- ing tasks, demonstrating strong capabilities in contextual understanding [57 37 4 and language generation [14 TT.. Despite these advancements, LLMs still struggle to maintain focus truly
critical information, especially long-context reasoning [52 55 20 and complex instruction
tasks [51 31, which adversely impacts reasoning accuracy and generation quality.
To systematically investigate this phenomenon, we first identify distracting patterns via gradient- based comparisons of teacher and student sensitivities, then assess the performance of student models on NuminaMath- CoT [26] and MocCoode-7KK (53]. Surprisingly. as shown in Figure [ simply pruning distracting patterns yields substantial gains improving average accuracy by over 20% the MATH
B Corresponding author: Yankai Lin (yankailin@mmee.duuu
2 Code and data are available httttp:://gtthub.cmm////tt///
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

1LaMA32-BB-Enaee 2-3i--a
LLaMA32 -30 Instrx
25
Quen2.5
Absalute Imnpowenum (95
Qween2.5
Queet: 5-1. 1.33 Insmict
20
3040
16.33
17.5
25.86
35.66
15
Absolate
t0
11.25
13.33
11.00
20
15030
1543
10
Olympinds AMC_AIME
(a) Code Accuracy Improvement
(b) MATH Accuracy Improvement
Figure I: Accuracy improvements achieved by removing confounding tokens from small models on the math and code training corpora. The results demonstrate significant increase in performance, with over 20% improvement on the math corpus and more than 10% on the code corpus. (For further details on these categories, see Appendix A
Distracting Patterns Region
In certain colony of bacteria, the number of bacteria doubles every day. The colony starts with 3 bacteria, and has at the end of day 1, 12 at the end of day 2, and so on. What is the number of the first day which ends with the colony having more than 100 bacteria?
Before Confounding Tokens Pruning
After Confounding Tokens Pruning
Stop 1: Understand the growth pattern the bacteria colony. The colony with bacteria doubles every day: This on the second day, there will be be 6 *2 12 bacteria, and so on. Step 2: Calculate number bacteria
Step 1: Identify the initial number of hacteria and the growth
that
pattern. The colory day.
bacteria doubles every
day. To when the colony has more than 100 bacteria, need to calculate the number bacteria on each day until we reach that point. 3: day 1, colory has harteria.
day
Step
end day
to find when exceeds 100.
Day 1: *2
Day 2: + =
Step 4: On day 2, the colory doubles to 12 bacteria.
Day 3: 12 2 24
Step day
colony doubles again
bacteria.
Day
Step 6: On day 4, the colony doubles again to 48 bacteria.
Day S: 43 2 96
Step 7: day 5, colony doubles again
Day
Step 8: On day 6, the colory doubles again 192 bacteria.
Step Determine the first day the col collory has more than 100 bacteria. From the calculations, we see that the colony has
11:0n day the calony doubles again 1536 bacteria. The final answer fs: 1536
than bacteria the 6tth
The ftpal answer [s:
Figure Comparison reasoning before and after pruning distracting patterns. Blue-shaded regions indicate pruned confounding tokens. Pink highlights mark areas that require focus, while blue highlights show where excessive attention caused errors.
training corpus 26 and more than 10% on Code training corpus 53 Furthermore, we observe greater improvements on AMC AMC_AIME AIME (26 compared GSM8K
suggesting that complex reasoning problems may contain more distracting patterns that interfere with model inference. These findings demonstrate that mitigating the influence distracting patterns essential for improving the robustness and accuracy of L.LM reasoning. We further present representative case in Figure 2 Removing distracting patterns from the instruction, without any additional training, helps the model focus critical information and enhance reasoning. This suggests that improving attention to relevant details is promising direction for advancing model reasoning.
g
Based on our analysis, we adopt a causal perspective (Figure 3) to explain and mitigate the observed phenomenon. We propose Learning Focus (LeaF), a two-stage framework that treats distracting patterns as spurious confounders in LLM reasoning. In the first stage, LeaF identifies confounding tokens through gradient- based comparisons between a high-capacity teacher and a student model. Then, it generates counterfactual samples by span pruning, removing contiguous spans of the detected confounding tokens from each instruction. In the second stage, LeaF introduces a hybrid distillation loss that minimizes two KL divergences: one for original sample (standard distillation) and one for

---

counterfactual sample (counterfactual distillation). This composite objective encourages the student model capture true causal dependencies by contrasting the teacher model's outputs before and after pruning confounding tokens, improving the robustness and interpretability of L.LM reasoning.
We demonstrate through comprehensive experiments that Learning Focus (LeaF) significantly enhances critical token identification and attention consistency during inference, leading to improved performance on downstream tasks such as mathematical reasoning, code generation and multi- hop question answering. Compared to standard knowledge distillation, LeaF yields an average accuracy gain of 2.41% % on GSM8K 8 MATH-500 [7]. and OlympiadBench [T6] for LLaMA- IB/3B-Instruct [34] and Qwen2.5- 1at--MMtthttt 48 In code generation, LeaF achieves an average improvement of 2.48% on HumanEval+ [32. LeetCode [9] and LivecodeBench [22]. On multi- hop question answering, LeaF improves performance by an average of 3. 3.24% on HotpotQA [50], 2WikiMultHHopQA [9, and Musique [44. These results validate our hypothesis that enhancing attention to key information is critical for improving reasoning performance. Attention visualizations in Section 4.4 further demonstrate LeaF"s interpretabilityy
2 Methodology
2.1 Causal Framework
Following Pearl's 's Structural Causal Model [39]. we formulate DAG G' to model the causal relation- ships among different components (refer to Figure B. We assume the distribution of (X, A, Y) is faithful to the DAG. In this framework, X [=1 22, ................................ I,,] represents the input tokens, and Y is the model's output. We define confounding tokens as a subset A that obscure true causal relationships by introducing spurious correlations with both the output Y and the complementary input X \ A. These misleading dependencies distort the model's attention mechanisms and bias its reasoning process, ultimately yielding unreliable predictions.
The desired behavior of the model is to eliminate spurious correlations introduced by confound- ing tokens A. When A influences both X and Y (see dashed arrows in Figure 3), the observed conditional distribution becomes:
X
A
X R-
X
P(Y X; = x)
Σ P(Y 2) A
(1)
which deviates from the interventional distribution P(Y | (x))) and reflects bias introduced by the indirect influence of A on Y through spurious paths.
Y
Figure Causal graph of the reasoning process. X represents the input prompt, and Y denotes the model's output. A subset of tokens in X, identified as confounding tokens (A), introduces spurious correlations that disrupt the reasoning process. Our method detects and masks A, effec- tively eliminating the spurious edge from A to Y and restoring the true causal dependency.
To block these non-causal influences, we propose causal pruning. which removes the effect of prior to distillation. This encourages the student model to learn attention patterns grounded in true causal structure, improving robustness and inter- pretability.
2.2 LeaF: Learning to Focus Framework
To eliminate spurious dependencies, we introduce the Learning to Focus (LeaF) framework (Figure 4). which consists of two main stages: (1) Confounding Token Detection (Section 2.2.1 where LeaF identifies confounding tokens via teacher-student gradient-based comparisons and constructs counterfactual samples by pruning these tokens. (2) Causal Attention Distillation (Section 2.2.2). where LeaF captures causal dependencies through hybrid distillation loss that aligns the student with the teacher on both original and counterfactual samples. We discuss them in detail in the following.
3

---

Stage 1: Confounding Token Detection
PDD
X
Ag
Pruning respectively
P
Tescher Model
Teacher Model
Taiis
Training Samples
Counterfactud Samples
Student Model
A0
8i
G.
every
firstday doubles
Student Model
Confounders
Gradient Comparison
Confounder Pruning
Stage 2: Casual Attention Distillation
Original
somple
Student Model
Teacher Model
Original Training Samples
Copture casual relationshipp
Knowledge
Inference
Learn Focusl
Distillation
Sample
Distribution
Sample
Counterfactudl Samples
Counterfactual Distillation
Student Model
The training pipeline comprises two key stages: (1) Confounding Token Detection: gradient-based comparisons between an advanced teacher model and the student model are used to identify confounding tokens in the training samples and constructs counterfactual samples by pruning these tokens; and (2) Causal Attention Distillation: prune identified confounders respectively during training align the student's 's attention with the teacher's and capture casual rela- tionships. This targeted intervention steers the model toward actual causal dependencies, enhancing both robustness and interpretability.
Figure 4: Method Overview.
2.2.1 Confounding Token Detection
To identify the confounding tokens A that introduce spurious correlations, we adopt a gradient-based
approach 42 43 to quantitatively measure the influence of each token on the model's i output Y. Specifically, we leverage the gradient sensitivity of both the teacher model 0 and the student model 05. We focus on data instances mispredicted by the student but correctly handled by the teacher to isolate confounding tokens. For each token x; X, we compute the gradient of the cross-entropy loss between predicted logits and gold references with respect to the embedding of x; for both models.
Dl(x; 0x
(S)
0f(r; (s)s)
g
8r;
Ox:
These gradients reflect the sensitivity of each model to perturbations in F, To enable token-level comparison between models with differing gradient scales, we apply min-max normalization to sensitivity values.
(S)
9i miny 9j
9, (9)
min, 9j
(T) min, 9j
(T)
(S) min, 9j
S8)
max, 9j
max, 9j
To identify confounding tokens, we capture the difference in token-level attention between the teacher and student models by computing the gradient difference for each token:
(T) (8) Ag. = 0: 9.
(4)
To ensure consistent scaling of gradient discrepancies across instances, we normalize the gradient difference and classify token from the student model but negligible attention from the teacher during inference, as formalized in Equation 5 and (ii) its removal results in correct predictions from both models.
Confounding Token if (i) it receives significant attention
Ag. min, AÚ,
Taonfounders
max; AØÚ min; AØj
where Teonfounder is threshold determined via statistical analysis on validation set. Sensitivity analyses on the threshold are presented in Section 4.3 Intuitively, confounding tokens capture sensitivity discrepancies between the teacher and student models that indicate spurious dependencies.

---

Original Training Corpus
Collective Pruning
Sanple 1 X1A,
Sample 2
Spon Pruning
Sample k x1,4
Valid Token
Confounding Tokens
Distracting Patterns Region
Figure 5: Illustration of Collective Pruning and Span Pruning.
Moreover, we also explored perplexity-bee based detection, but without teacher guidance, it tends to capture tokens indicative of model uncertainty rather than true confounders, especially on challenging tasks like the Olympiads. Accordingly, we adopt gradient-based detection our primary strategy (see Section 4.1 for details).
Pruning Strategies. We study two pruning strategies for removing confounding tokens from instruction X: (1) Collective Pruning, which removes the entire set of identified confounders A, yielding X A. (2) Span Pruning, which removes only one contiguous confounding span A, a time, yielding X A, Preliminary experiments show that span pruning outperforms collective pruning (see Appendix C as pruning all distracting patterns simultaneously disrupts sentence integrity. Hence, we construct the counterfactual samples via the span pruning strategy:
Prruned
{(X Ai v)) )==
where each A; denotes distinct confounding span. This augmentation encourages model to learn reasoning paths invariant to specific confounders.
2.2.2 Causal Attention Distillation
After generating both original and counterfactual samples, we optimize two complementary distilla- tion objectives to steer the student toward true causal dependencies:
Standard Distillation. Align the student's output distribution with the teacher's original instruc-
tions:
Led DKL (yr(y) XX) || ps X)),
where pr and ps denote the teacher and student output distributions, respectively.
Counterfactual Distillation. Align the student's output distribution with the teacher's counterfac-
tual instructions X A (confounders pruned):
Led DKL (pr(y) X A) (| X A)).
we blend these objectives with a weighting factor X:
= L + (1 - X) Ledt
where À [0, 1] controls the trade-off between standard and counterfactual distillation. This composite loss steers the student to preserve semantie knowledge while enforcing genuine causal dependencies.
Response Splitting Strategies. For our methods, we consider two variants: (1) Instruct-level Pruning: We detect and prune confounding tokens only in the instructions, and perform our LeaF on instruction-level pruned samples. (2) Both Instruct- and Response-level Pruning: We can also treat previously generated tokens as contextual input, and prune those that are misleading for subsequent generation, order to help the model produce more accurate continuations, Thus, we detect and prune confounding tokens in both instructions and preceding generations.

---

Language To (training corpus)
[Instruction] [Step 1] [Step 2] [Step 3] [Step K].,[Stpp N-1] [Step N] [Answer]
Instruct-level Pruning (Response without split)
[Instruction pruned distracting patterns] [Step 1] (Step 2] [Step N] [Answer]
Response-level Pruning
[Instruction] [Half of response pruned distracting patterns] [Half of the Response
(2-segment splins)
Response-leve level Pruning (3-segment splits)
[Instruction] [Segment pruned distracting patterns] [Segment 2] [ [Segment
[Instruction] [Segment
[Segment pruned distracting patterns] [Segment 3]
Figure 6: Illustration of Response Splitting Strategies: Language CoT, Instruct-level Pruning, and Response-level Pruning (2-segment and 3-segment splits). Highlighted white areas represent the input, and blue underlined areas represent the outputs used for cross-entropy loss computation.
3 Experiments
3.1 Experiment Settings
Training Datasets. We conduct experiments to evaluate the effectiveness our proposed method Learning Focus Framework (LeaF) on mathematical reasoning. code generation, and multi-hop question answering tasks. For mathematical reasoning, to ensure the model encounters an equal number of confounding tokens across tasks, we randomly select 30k instances from each of the following subsets in NuminaMath-CC T [26]: Olympiads [16]. AMC AIME (26). GSM8K (8) and MATH [17]. For code generation, randomly select subset of 120k instances from the AceCode- 87k [53] dataset. For multi-hop question answering, we construct the training set by merging the
KILT 40 datasets provided Helmet
[52) totaling 3k annotated samples drawn equally from HotpotQA [50], NQ [I, and PopQA (36], where each query is explicitly linked to its corresponding
gold passages containing the answers.
Evaluation Datasets.
For math task evaluation, we select three widely used benchmarks with varying levels of difficulties, including GSM8K (8, MATH-500 [17), and OlympiadBench [16]. In code domain, we conduct evaluations on HumanEval+ 32 LeetCode [9, and LivecodeBench (v4) (22),
providing comprehensive assessment across diverse aspects of coding performance. For multi-hop question answering, we evaluate on three representative benchmarks, including HotpotQA [50], 2WikiMulúHHPpQA [19] and Musique 44 The detailed description of the mathematical reasoning, code generation and multi-hop QA benchmarks is provided in Appendix K
Base Models and Baselines. We conduct comprehensive experiments two different model families,
LLaMA family [8 (34 and Qwen family [48], covering various sizes models. For LLaMA-base
experiments, we employ LLaMA3.2- Instrruct 34 and C.MMA...--B-Insrtuuee
[34] student
models, while their teacher models are set as -L-MM333-770B--...
[34]. For Qwen-based
experiments, as the teacher model. Our baseline for comparison is standard knowledge distillation without pruming. We use the CoT-based variant for math and the vanilla variant for code generation and multi-hop QA.
Qwen2.5-Math-11. (48) the student model, with Owen2.5-72BBInntmee [48]
Training and Evaluation Settings. (1) During training, models are trained using the Alpaca-LoRA framework with nall-parameter logits knowledge distillation and cosine learning rate schedule with maximum learning rate of 10 for three epochs. The batch size is 64 for LLaMA based models and 32 for Qwen- based models. Detailed hyperparameters and platform information are in Appendix [] (2) In evaluation, teacher and student models are evaluated on math, code and QA tasks using greedy decoding. The maximum generation length is 1024 tokens for code and 16,384 tokens for math tasks. Official chat templates are followed during inference. Full evaluation details are Appendix B
3.2 Main Results
The main results are presented in Table [ and Table 2 the results: (1) For open-source baselines, both standard distillation and our approach lead to
We can draw several conclusions from

---

Table 1: Performance MathBench and CodeBench for Instruct models (LLaMA 1B2-1B/BB Instruct) and Base model (Qwen 2.3 2.5-Math-1-33 1.5B) under three pruning schemes (no mask, Instruct- level Mask, and Response-level Mask), where Instr Mask refers to Instruct-level pruning, and Resp Mask refers to Response-level pruning. The best and second-best results are marked in bold and underlined, respectively.
MathBench
CodeBench
Model
Olympiad- Bench
Human- Eval+
Leet- Code
Livecode- Bench
GSM8K
MATH-500
Avg.
Avg.
Teacher Model -LaMA3.3-70B-Inssie
95.60
70.40
36.50
67.50
78.05
53.90
45.02
58,99
Qwen2.5 2een2.5-72B-nnstruue
95.45
73.80
41.25
70.17
81.71
69.40
54.42
68.51
LLaMA3.2-1B-Inss
Instruct Model (Pre-KD)
44.88
24.20
5.79
24.96
29.27
7.22
9.68
15.39
KD w/o Mask
56.79
33.40
8.90
33.03
32.32
6.11
13.74
17.39
LeaF (Instr Mask)
57.70
35.40
10.09
34.40
39.02
6.67
13.60
19.76
LeaF (Instr & Resp Mask)
58.98
35.20
9.94
34.71
39.63
7.22
12.48
19.77
LLaMA3.2-3B-Instte
Instruct Model (Pre-KD)
76.88
42.80
13.20
44.29
48.78
13.89
20.34
27.67
w/o Mask
82.87
49.00
18.99
50.29
54.88
16.67
24.12
31.89
LeaF (Instr Mask)
83.09
51.80
20.77
51.88
55.49
19.44
25.39
33.44
LeaF (Instr & Resp Mask)
84.69
52.40
22.55
53.21
56.
21.67
25.81
34.53
@wen2.5-Matle
Base Model (Pre-KD)
65.20
41.40
21.96
42.85
35.37
6.67
1.26
14.43
KD w/o Mask
82.18
67.80
31.16
60.38
41.46
7.78
10.10
19.78
LeaF (Instr Mask)
84.69
68.60
32.79
62.03
42.68
9.94
10.80
20.97
LeaF (Instr & Resp Mask)
85.29
70.60
31.75
62.54
43.29
9.94
13.04
21.92
Table 2: Performance multi- hop QA benchmarks under different pruming strategies. Given the short response length of multi-hop QA tasks, only the LeaF (Instruct-level Mask) variant is evaluated. Given the long-context requirement of multi-hop QA, validate LeaF using L.aMA-3.22 1...MM.-..--B---mm Innstucttttt
2WikiMultHHpQAA
Musique
HotpotQA
Model
Avg.
EM
F1
SubEM
EM
F1
SubEM
EM
F1 SubEM
LLaMA3.2-3B-Instte
Instruct Model (Pre-KD)
22.50
34.73
45.50
9.50
16.56
18.50
19.50
32.67
36.00
26.16
KD w/o Mask
43.50
51.93
53.00
20.50
28.56
22.50
39,00
49.30
43.00
39.03
Leaf (Instr Mask)
46.50
53.89
55.00
27.00
33.00
28.00
40.50
52.06
44.50
42.27
improvements in the performance of models across nearly all tasks in math, code and multi-hop question answering domains. (2) In terms of performance, LeaF' consistently outperforms standard
knowledge distillation when using the same training corpus. This suggests that LeaF has the potential to enable the model to attend to critical information more effectively, thereby enhancing its reasoning capabilities. (3) Extending from instruction-level to response-level pruming yiclds further performance improvements across most tasks the LLaMA and Qwen series, suggesting that distracting patterns at the response level affect subsequent generations. We hypothesize that instruction-level and response-level distracting patterns differ, and simultaneously learning both can enhance the model's reasoning capabilities. A more detailed analysis based on the response splitting strategy is provided in Section 4.2 Additionally, Appendix [ shows that LeaF more robust than other pruning strategies.
4 Further Analysis
In this section, we conduct masking strategies analysis (Section 4.1
response splitting analysis
(Section (42) and threshold sensitivity analysis (Section 4.3). Finally, we present case study
(Section 4.4 to illustrate the interpretability of our approach. In the appendix, we include robustness

---

Random Manking
Resoose wilhout Splk
PPI-baood Mocking
Requosse Spli into Segments
Absolute Improveent (%)
Gradient based Masking
Requasse Split Segments
2,80
Abwolute Improvement amprovte
3,40 3.40
3.56
2.00 2.20
1773
2.15
-0.07
1.82
1.78
1.59
4,00
-4.45
GSMSK
MATH
OlympisdBench
Avg.
GSM8K
MATH
OlympiadBench
Avg.
Figure 7: Comparison of accuracy improvement with masking strategies over baseline (KD).
Figure 8: Comparison of accuracy improvement with splitting strategies over baseline (KD).
analysis (Appendix E
computational overhead analysis (Appendix D ablation study (Appendix F). and dataset statistics of counterfactual and original samples (Appendix H
threshold stability across tasks (Appendix G
4.1 Masking Strategies Analysis
To demonstrate the effectiveness of our gradient-based masking strategy, we compare LeaF with two other choices of token masking: (1) Random Masking, where the same number of tokens identified by our gradient- based method are randomly masked for each instance; and (2) PPL-based Masking, where tokens with the highest perplexity scores are masked at the same ratio, following the same data filtering process used in our method to generate augmented training data. We evaluate three masking strategies, Random Masking, PPL--based Masking, and Gradient-based Masking, the GSM8K, MATH-500, and OlympiadBench datasets.
Results. We present the results Figure 7 Our observations are as follows: (1) Gradient-based Masking (ours) consistently outperforms both baselines, with the highest accuracy on MATH-500 and OlympiadBench. (2) Random Masking leads performance degradation on GSM8K and Olympiad, despite showing slight improvement on MATH-500. This suggests that naively masking tokens without informed selection can undermine distillation performance. (3) PPL.-based masking provides modest improvements on GSM8K and MATH-500, but performs comparably to random masking on OlympiadBench, indicating limitations on complex tasks. Our analysis is, while perplexity may suffice for detecting confounding tokens in simpler settings, lacks the sensitivity needed for challenging benchmarks. This highlights the necessity an advanced teacher model to guide token selection in complex reasoning scenarios.
4.2 Response Splitting Strategies
We compare three response splitting strategies, considering that distracting patterns exist at both the instruction and response levels. We evaluate response without split (instruction-level pruning), 2-segment splits, and 3-segment splits. diagram illustrating these strategies is shown in Figure
Results. As shown the results in Figure 8 our findings are as follows: (1) Response-level pruning (both 2-segment and 3-segment splits) significantly outperforms instruct-level pruning, demonstrating the importance and benefits of extending the learning of confounding tokens from the instruction level to the response level. We hypothesize that the improvement stems from differences in distracting patterns between the instruction and response levels, suggesting that combining both levels could further enhance model performance. (2) The performance of 3-segment splits is comparable to that of 2-segment splits, suggesting that further segmentation the response level yields diminishing returns. We hypothesize that distracting patterns at the response level exhibit certain regularities, and the data generated by 2-segment splits is sufficient for the model learn these patterns effectively, rendering additional segmentation unnecessary.

---

Average Performance (%)
LLaMAA,2-KD-11
Average Performance (%)
LLaMA3.2-KD-LI
LLaMoO.2-KD-30
LLaMAB.2-KD-3B
LLaM0332-Lea--18
LLaMA3.2-Lea--IB
60
LLaMA3.2-.aa-3388
60
LLMMA3.2-eea-338
51.58
50555
3121
5035
$1 1199
40-
40-
34400
13
5277
33.05
34.71
0.05
0.10
0.15
0.20
0.25
0.025
0.075 Confounder Threshold
0.05
0.10
Confounder Threshold
Figure Instruct-level in MathBench.
Figure 10: Response-level in MathBench.
4.3 Threshold Sensitivity Analysis
We perform a sensitivity analysis on the threshold used for confounding tokens pruning, assessing the performance of LLaMA3.2 IB---trrctttt and LLaMA3.2-3B 2-3B--nnt as student models. As illustrated the Figure 9 and Figure TO we present the average experimental results of these models across MathBench (GSM8K, MATH-500, OlympiadBench) at different threshold levels.
Results. We observe the following results: (1) Instruct-level: -LaMa3.2-Leate performs best at a threshold of 0.10, while LLaMa3. 3-MMa3332Laa-333 performs best at 0.05. (2) Response-level: LLaMa3.22 LeaF-1B performs best 0.15, and L.L.Ma3.2- response-level, I.LMM3.2244 2.-aaa..F- achieves optimal performance at higher misleading token threshold than I.LaMa3.2-Less 2..ea333-- This suggests that smaller models, which are more susceptible to confounding tokens, benefit from higher thresholds that filter out disruptive tokens more effectively. Detailed cross-domain results presented in Appendix further confirm that Tmis = 0.10 yields stable performance across tasks.
at 0.10. (3) For both instruction-level and
4.4
Case Study in an Interpretable Perspective
We conduct a case study from an interpretability perspective show that LeaF, compared standard knowledge distillation (KD), enables the model to focus critical information and avoid confounding tokens. The interpretability case study shown Figure TT illustrates that LeaF allows the model to focus more on the critical information, such as "real number", "all", "are real". By identifying the requirement that roots be real, LeaF guides the model through a coherent reasoning chain: recognizing = as an evident real root, and applying the discriminant condition to ensure the quadratic factor also yields real solutions. In contrast, misapplies the AM- GM inequality potentially negative values, resulting in an incorrect final answer. This demonstrates that LeaF enables the model to focus on critical information better.
KD model overlooks this constraint and
Related Work
Chain-of-Thoughh Knowledge Distillation. Our work falls within the knowledge distillation paradigm [18 [4], a widely adopted technique for boosting the performance of open-source models complex tasks such as mathematical
Subbemm-solving [58] 59] and code generation B 21. CoT-
KD [41] takes the first step towards this line to employ chain-of-thoughl explanations generated by high-capacity teacher fine-tune student model, thereby equipping with advanced reasoning capabilities. Recent advancements can be broadly categorized into two complementary directions:
(1) Data-focused approaches aim enhance distillation by improving the quality and diversity of training data (e.g., CD [13], SCORE [56], and Skintern (28). (2) Model-focused approaches enhance model architectures and inference strategies to improve efficiency and reasoning capability (e.g., PRR 41 and ATM (7).. However, existing methods primarily focus on output imitation, while our approach explicitly distills the teacher's ability capture critical information during inference into the student, enabling context-aware reasoning in the student model.

---

Let
be
positive
real
number
such
that
the
roots
of
ax
Find
smallest
LlaMa3.2 2-33-- Instruct (Distilled by KD)
LlaMa3 2 1--3--I-nie Instruct (Distilled by LeaFF)
Step given cubic equation
axx2
Step We are given a cubic equation x we need to find the smallest possible value of such that all routs of equation are real.
and we need to find the smallest possible value of approach problem, can start by applying the AM-GM inequality to the roots of the given cubic equation.
Step given equation = (x + 1)(x"2 + (a- [a-1)x+ + using the sum of cubes formula..
fastored as + ax ax
Step 2: Let the roots of the equation be and t. By Vieta's s Applying the AM-GM inequality to the roots, get [sqrt[3][sst)
formulas,
Stap 3: For the roots of the equation to be real, the discriminant of quadratic factor x (a-1)x 1)x + must nonnegative.
Step 3: Substituting the values of + and rss AM- GM inequality we have is \sgrtf3][-1] (333 This gives inequality i (sqrt[3][3).
Step Setting the discriminant greater or equal zero, we get the inequality (a 1y"2 4 is U.
Solving incquality will give satisfy the condition of having real roots.
range
for
Step 4:
the smallest possible determine when equality occurs the AM GM inequality. Equality occurs when = = t= t- ((
to
Step S: Solving the inquullit (a- 1yn2 44 D, find that (a
1)22 > 4, which gives a- Wlee ora- This simplifies to We or
Step S: Substituting = t= (grr))) into the equation x*3 + ax"2 + we can solve for a to find the smallest possible value.
Step However, since positive real number, the only valid solution from the inequality a ze 3.
This gives \sgrt[3](3).
This means that the smallest possible value of Sa$ that satisfies the condition of having all real roots is 3.
The final answer is:
[3
The final answer
Figure 11: Case study comparing LeaF and knowledge distillation (KD) performance the MATH- 500. Top:
Heatmap showing the attention score differences between KD and LeaF on instruction tokens. Dark blue indicates higher KD attention, and deep pink indicates higher LeaF attention. Left: Response of LLaMA-3.2-333 3-33----ttttt after standard knowledge distillation, with blue text marking incorrect steps by the KD model. Right: Response of I.L.MA-..-LLaa--3BB with pink text highlighting correct steps by our model.
Critic Token Identification. Existing works have explored various strategies to identify and mitigate redundant steps [10
33 35 or less-important tokens [15] 47 during language model reasoning.
Methods like LLMLingua (33 rely on models' self-assessment to determine token importance, which may introduce biases due to the model's limited reasoning capacity. More recent approaches further refine token selection across different training stages. RHO-1 29 introduces Selective Language Modeling (SLM) during the Supervised Fine-Tuning phase, prioritizing informative tokens to enhance efficiency and reasoning accuracy. TokenSkip [47] focuses on ofain-of-Thhught stages, selectively skipping low-impact tokens compress reasoning paths without sacrificing performance. Meanwhile, cDPO [30] enhances Direct Preference Optimization by isolating critical tokens through contrastive learning. However, these methods primarily concentrate output tokens while overlooking the influence of prior contextual information from the instruction phase. In contrast, our method introduces advanced teacher model to identify confounding tokens based gradient differences, establishing stronger connections between instruction and output, and achieving more holistic and effective token filtering.
Reasoning Consistency. detailed discussion of reasoning consistency is provided in Appendix
6 Conclusion
In this paper, we proposed Learning to Focus Framework (LeaF), novel strategy to enhance the consistency of large language models. By leveraging causal analysis and gradient- basee pruning, our method effectively identifies and eliminates confounding tokens, enabling the student model to cap- ture more reliable causal dependencies. Experimental results show substantial improvements both accuracy and robustness across mathematical reasoning, code generation and multi-hop QA bench- marks. In future work, how to achieve model consistency through self-improvement mechanisms without relying advanced models remains promising direction for further exploration.
10

---

Acknowledgements
This work was supported by The National Natural Science Foundation of China (No. 62376273 and No.U2436209) and Beijing Nova Program (No. 20240484568).
We sincerely thank all the anonymous reviewers and (S)ACs for their constructive comments, as well as Zhen Tian, Tingchen Fu, Wentong Chen, Yongda Yu, Hongjia Liu and Qinghui Wang for their valuable discussions on both the development of the ideas and the writing of this paper.
References
[1] David Ifeoluwa Adelani. Jade Abbott, Graham Neubig, Daniel D'suzza, Julia Kreutzer, Con- stantine Lignos, Chester Palen- Miice,, Happy Buzaaba. Shruti Rijhwani, Sebastian Ruder, et al. Masakhaner: Named entity recognition for african languages. Transactions of the Association for Computational Linguistics, 9:1116-1131, 2021.
[2] Pranjal Aggarwal, Aman Madaan, Yiming Yang. et Let's sample step by step: Adaptive- consistency for efficient reasoning and coding with llms. In The 2023 Conference on Empirical Methods Natural Language Processing, 2023.
[3] Wasi Uddin Ahmad, Sean Narenthiran, Somshubra Majumdar, Aleksander Ficek, Siddhartha Jain, Jocelyn Huang, Vahid Noroozi, and Boris Ginsburg. Opencodereasoningg: Advancing data distillation for competitive coding. arXiv preprint arXiv:2504.0194. 2025.
[4] Shengnan An, Zexiong Ma. Zeqi Lin, Nanning Zheng, Jian Muan-uang Lou, and Weizhu Chen. Make your IIm fully utilize the context. Advances in Neural Information Processing Systems, 37:62160-622
2024.
[5] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. 2023.
[6] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph. Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 2021.
[7] Xiaoshu Chen, Sihang Zhou, Ke Liang, and Xinwang Liu. Distilling reasoning ability from large language models with adaptive thinking. arXiv preprint arXiv:2404.09177, 2024.
[8] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jenyy Tworek, Jacob Hilton, Reicchio Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14163 2021.
[9] Tristan Coignion, Clément Quinton, and Romain Rouvoy. A performance study of llm-generated code on leetcode. In Proceedings of the 28th International Conference on Evaluation and Assessment in Software Engineering, pages 79-89, 2024.
[10] Mengru Ding, Hanmeng Liu, Zhizhang Fu, Jian Song, Wenbo Xie, and Yue Zhang. Break the chain: Large language models can be shortcut reasoners. arXiv preprint S.Kiv:2406.068888 2024.
[11] Chenhe Dong, Yinghui Li, Haifan Gong, Miaoxin Chen, Junxin Li, Ying Shen, and Min Yang. A survey of natural language generation. ACM Computing Surveys, 55(8):1-38, 2022.
[12] Shaohua Duan, Xinze Li. Zhenghao Liu, Xiaoyuan Yi, Yukun Yan, Shuo Wang, Yu Gu, Ge Yu, and Maosong Sun. Chunks arms: Multi-ameel bandit-guided sampling for long-context I1m preference optimization. arXiv preprint arXiv:2508. 13993, 2025.
[13] Tao Feng, Yicheng Li, Li Chenglin, Hao Chen, Fei Yu, and Yin Zhang. Teaching small language models reasoning through counterfactual distillation. Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 583 5831-5842, 2024.
11

---

[14] Albert Gatt and Emiel Krahmer. Survey of the state of the art in natural language generation: Core tasks, applications and evaluation. Journal of Artificial Intelligence Research, 61:65-170, 2018.
[15] Tingxu Han, Zhenting Wang, Chunrong Fang, Shiyu Zhao, Shiqing Ma, and Zhenyu Chen. Token-budget-aware Ilm reasoning. arXiv preprint arXiv:2412. 185447777 2024.
[16] Chaoqun He, Renjie Luo, Yuzhuo Bai, Shengding Hu, Zhen Leng Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, et al. Olympiadbench: A challenging benchmark for promoting agi with olympiad-level bilingual multimodal scientific problems. In ACL (1), 2024.
[17] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874, 2021.
[18] Geoffrey Hinton. Distilling the knowledge in a neural network. arXiv preprint arXiv: 2015.
[19] Xanh Ho. Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing a multi-hop qa dataset for comprehensive evaluation of reasoning steps. arXiv preprint arXiv:2011.01060, 2020.
[20] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size your long-context language models? arXiv preprint arXiv:2404.06654 2024.
[21] Yufan Huang, Mengnan Qi, Yongqiang Yao, Maoquan Wang, Bin Gu, Colin B Clement, and Neel Sundaresan. Program translation via code distillation. In EMNLP, 2023.
[22] Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang, Ar- mando Solar-Lezama, Koushik Sen, and lon Stoica. Livecodebench: Holistic and contamination free evaluation of large language models for code. arXiv preprint NKiv:2403.07977 2024.
[23] Huiqiang Jiang, Qianhui Wu, Chin- Yew Lin, Yuqing Yang. and Lili Qiu. LImlingua: Compress- ing prompts for accelerated inference of large language models. In The 2023 Conference Empirical Methods in Natural Language Processing, 2023.
[24] Xin Lai, Zhuotao Tian, Yukang Chen, Senqiao Yang. Xiangru Peng, and Jiaya Jia. Step-dpo: Step-wise preference optimization for long-chain reasoning of llms. urXiv:2406.18622 2024.
[25] Dacheng Li, Shiyi Cao, Tyler Griggs, Shu Liu, Xiangxi Mo, Eric Tang, Sumanth Hegde, Kourosh Hakhamaneshi, Shishir G Patil, Matei Zaharia, et al. Llms can easily learn to reason from demonstrations structure, not content, what matters! arXiv preprint arXiv:2502.07374, 2025.
[26] Jia II, Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Costa Huang, Kashif Rasul, Longhui Yu, Albert Jiang, Ziju Shen, Zihan Qin, Bin Dong, Li Zhou, Yann Fleureau, Guillaume Lample, and Stanislas Polu. Numinamath. http:////gagaaaaa...ooooo///////
(https://github com/ wttpa/aammmmooooooo http:///www....pp.pppppp///////a www.aaunnnnnnnnnnntnnnaeau pdf)
2024.
[27] Yiwei Li, Peiwen Yuan, Shaoxiong Feng, Boyuan Pan, Xinglin Wang, Bin Sun, Heda Wang, and Kan Li. Escape sky-high cost: Early-stopping self-consistency for multi--tep reasoning. In The Twelfth International Conference on Learning Representations, 2024.
[28] Huanxuan Liao, Shizhu He, Yupu Hao, Xiang Li, Yuanzhe Zhang, Jun Zhao, and Kang Liu. Skintem: Internalizing symbolic knowledge for distilling better cot capabilities into small language models. In Proceedings of the 31st International Conference on Computational Linguistics, pages 3203-32022 2025.
[29] Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, et al. Rho-1: Not tokens are what you need. arXiv preprint arXiv:2404.0700 2405..9... 2024.
12

---

[30] Zicheng Lin, Tian Liang, Jiahao Xu, Xing Wang, Ruilin Luo, Chufan Shi, Siheng Li, Yujiu Yang, and Zhaopeng Tu. Critical tokens matter: Token-level contrastive estimation enhence Ilm's 8s reasoning capability. arXiv preprini arXiv:241 1..244/111993... 19943,, 2024.
[31] Zhan Ling, Kang Liu, Kai Yan, Yifan Yang, Weijian Lin, Ting-Han Fan, Lingfeng Shen, Zhengyin Du, and Jiecao Chen. Longreason: A synthetic long-context reasoning benchmark via context expansion. arXiv preprint arXiv:2501. 1509, 2025.
[32] Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. Is your code generated by chatGPT really correct? rigorous evaluation of large language models for code generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[33] Tengxiao Liu, Qipeng Guo, Xiangkun Hu, Cheng Jiayang, Yue Zhang, Xipeng Qiu, and Zheng Zhang. Can language models learn to skip steps? arXiv preprint arXiv:2411.0883 2024.
[34] Zechun Liu, Changsheng Zhao, Igor Fedorov, Bilge Soran, Dhruv Choudhary, Raghuraman Krishnamoorthi, Vikas Chandra, Yuandong Tian, and Tijmen Blankevoort. Spinquant: Lim quantization with learned rotations. arXiv preprint arXiv:24055 164006 2024.
[35] Xinyin Ma, Guangnian Wan, Runpeng Yu, Gongfan Fang, and Xinchao Wang. Cot-valve: Length-compressible chain-of-thought tuning. arXiv preprint arXiv:2502.066/ 2025.
[36] Alex Mallen, Akari Asai, Victor Zhong, Rajarshi Das, Daniel Khashabi, and Hannaneh Ha- jishirzi. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. arXiv preprint arXiv:2212.00511 2022.
[37] Daye Nam, Andrew Macvean, Vincent Hellendoorn, Bogdan Vasilescu, and Brad Myers. Using an Ilm to help with code understanding. Proceedings of the IEEF/ACM 46th International Conference Software Engineering, pages 1-13, 2024.
[38] David Patterson, Joseph Gonzalez, UIs Hölzle, Quoc Le, Chen Liang, Lluis-Miquel Munguia, Daniel Rothchild, David R So, Maud Texier, and Jeff Dean. The carbon footprint of machine learning training will plateau, then shrink. Computer, 55(7))18-28, 2022.
[39] Judea Pearl. Causal inference. Causality: objectives and assessment, pages 39-58, 2010.
[40] Fabio Petroni, Aleksandra Piktus, Angela Fan, Patrick Lewis, Majid Yazdani, Nicola De Cao, James Thorne, Yacine Jernite, Vladimir Karpukhin, Jean Maillard, et al. Kilt: a benchmark for knowledge intensive language tasks. arXiv preprint arXiv:2009.02255 20220000.0222.... 2020.
[41] Kumar Shridhar, Alessandro Stolfo, and Mrinmaya Sachan. Distilling reasoning capabilities into smaller language models. In Findings of the Association for Computational Linguistics: ACL 2023, pages 7059-7073, 2023.
[42] Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps. arXiv preprint arXiv:1312.000 2013.
[43] Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In International conference on machine learning, pages 3319-33222 PMLR, 2017.
[44] Harsh Trivedi, Niranjan Balasubramaniann Tushar Khot, and Ashish Sabharwal. Musique: Multihop questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 0554,--554, 2022.
[45] Guangya Wan, Yuqi Wu, Jie Chen, and Sheng Li. Dynamic self-consistency: Leveraging reasoning paths for efficient Ilm sampling. arXiv preprint arXiv:2408. 2024.
[46] Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed H Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2022.
[47] Heming Xia, Yongqi Li, Chak Tou Leong, Wenjie Wang, and Wenjie Li. Tokenskip: Controllable chain-of-thoughe compression in 1lms. arXiv preprint 20X2 2025.
13

---

[48] An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu. Jianhong Tu, Jingren Zhou, Junyang Lin, et al. Qwen2. 5-math technical report: Toward mathematical expert model via self seff-improvement, arXiv preprint arXiv:2409. /2/22, 2024.
[49] Wenkai Yang, Yankai Lin, Jie Zhou, and Ji-Rong Wen. Distilling rule-based knowledge into large language models. In Proceedings of the 3 Interaatinnal Conference on Computational Linguistics, pages 913 113-932. 2025.
[50] Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhut- dinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. arXiv preprint arXiv:/809.09600, 2018.
[51] Xi Ye, Fangcong Yin, Yinghui He, Joie Zhang, Howard Yen, Tianyu Gao, Greg Durrett, and Danqi Chen. Longproc: Benchmarking long-context language models on long procedural generation. Xrviv preprint arXiv:2501.05.4 05444, 2025.
[52] Howard Yen, Tianyu Gao, Minmin Hou, Ke Ding, Daniel Fleischer, Peter Izsak, Moshe Wasserblat, and Danqi Chen. Helmet: How to evaluate long-context language models effectively and thoroughly. arXiv preprint arXiv:2410.0266 026944, 2024.
[53] Huaye Zeng, Dongfu Jiang, Haozhe Wang, Ping Nic, Xiaotong Chen, and Wenhu Chen. Acecoder: Acing coder rl via automated test-case synthesis. ArXiv, abs/2207.01780, 2025.
[54] Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q Weinberger, and Yoav Artzi. Bertscore: Evaluating text generation with bert. arXiv preprint XXX4::9994009755 2019.
[55] Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang. Zhiyuan Liu, et al. o00 bench: Extending long context evaluation beyond 100k tokens. arXiv preprini arXiv:24022 13718, 2024.
[56] Yunxiang Zhang, Muhammad Khalifa, Lajanugen Logeswaran, Jaekyeom Kim, Moontae Lee, Honglak Lee, and Lu Wang. Small language models need strong verifiers to self-correct reasoning. In ACL (CC((Fini (Findings). 2024.
[57] Zheng Zhao, Emilio Monti, Jens Lehmann, and Haytham Assem. Enhancing contex- tual understanding in large language models through contrastive decoding. arXiv preprint 24X5v:240..02750 2024.
[58] Xunyu Zhu, Jian Li, Yong Liu, Can Ma, and Weiping Wang. Distilling mathematical reasoning capabilities into small language models. Neural Networks, 179:106594 2024.
[59] Xunyu Zhu, Jian Li, Can Ma, and Weiping Wang. Improving mathematical reasoning capabilities of small language models via feedhack-driven distillation. arXiv preprint arXiv:2411. 14698, 2024.
14

---

A Preliminary Experiment
A.1 Gradient Heatmap Comparison
We compare the gradient heatmaps of the small model and the big model to analyze the differences in their attention and focus on critical tokens during the inference process.
Each
of
Natalie
's
blue
berry
bushes
yields
eight
containers of blue
berries
she
can
trade
five
containers of blue
berries for two z uC chin Is
how
many
bushes
does
Natalie
need
to pick in order
to
have
forty
-eight
chin
is
?
LlaMa3 2-1b-Instruct
LIaMa3.3- 3-70b-Instrcct
Step Determine number of containers of blueberries needed for 48 zucchinis.
establish the exchange containers blueberries and zucchinis We'ne told that 5 contalners of blueberries can be traded for 2 zucchinis. how many containers blueberriis needed for zucehini, we divide the number of contalners by the number michinis the trade:
Step
To find out how many containers of blueberries Natalie needs to trade 48 zucchinis, need divide the number zucrhinis by the number of containers she can trade for one zucchini. This will give us the total number containers nceded.
Step
containers 2 zucchinis 2.5 containers per zucchini.
Stop 3: Since Natalie wants zucchinis, muktiply number of zuechinis she wants by the number of containers needed per zucchini: zucchinis container's //zucchini 120 containers.
Step 2:
Calculate the number containers needed
48 gucchinis / zucchinis per container 24 containers
Deter the number of bushos needed yield 24 containers blueberries Since each bush yields containers of blueherries, we need to number of containers needed by the number of containers per bush to find out how many bushes are needed.
Step 3:
Step Now, we know that each Natalie's blueberry bushes yields a containers of blueberries. To find out how many bushes she needs yield 120 containers, divide the total number containers needed by the number containers per bush: contalners / contalners/but /buch bushes
Step 4: Calculate the number of bushes needed 24 containers containers per bush bushes
Therefore, Natalie needs pick 15 bushes have enough blueberries to trade for 48 zucchinis.
The final
The final answer is:
case study comparing the performance of the student model (LlaMa3.2 1 I-Instruct) and the teacher model (laMaa33-700-Instruue on the MATH task.
Figure 12:
Results. Figure 12 illustrates that Llama3,3 3.0-a....-7------ 708--nruuttt successfully captures a key contextual
relation: "Five containers of blueberries can be traded for two zucchinis." " Gradient heatmaps show that the teacher model aligns its attention closely with the relevant tokens, while the student model's attention is more dispersed. This observation motivates patterns, we can guide the student model to better focus on salient information, enhancing its reasoning capabilities. To evaluate this, conduct pilot studies in settings: (1) assessing accuracy gains on math and code benchmarks after pruning distracting patterns (Appendix (.2) and (2) evaluating improvements in response quality(Appendi)) A.3)
hypothesis: by pruning distracting
A.2 Performance Gains from Token Pruning
For the LLaMA series, we use LLaMA3.2 1B/3B-BBBBttttt as the student models and LLaMA3.3- 70B-Instruct as the teacher model. For the Qwen series, use 0wen2.5-Math-11.5 as the student model and Siven2.5-72B-Insstte
as the teacher model for preliminary experiments.
For mathematical problem solving, select filtered subset from the Numina-CoT dataset (26). where the teacher models generate correct inferences, consisting of 12K examples (3K each from GSM8K, Olympiads, AMC AMC_AIME, and MATH). We then select the subset where the student models produce incorrect results. We compute each sample's 's gradient difference between the student and teacher models to identify potential distracting patterns. Finally, these distracting patterns removed, and the student models are re-evaluated to assess the resulting improvement in accuracy. The results presented in Figure ((b)
For code generation, we construct filtered subset of 12K examples from the AceCode dataset 53 where teacher models generate correct inferences. This subset comprises 6K samples from the Evol
15

---

subset and 6K from the Oss subset. We apply the same distracting pattern identification and pruning procedure as used in the mathematical domain. Subsequently, we re-evaluate student models assess the impact of this intervention on code generation accuracy. The results are presented in Figure T(a)
A.3 Generation Quality Improvements from Token Pruning
We analyze the alignment between the student model's outputs and the teacher model's outputs under two conditions: (1) Original Instruction: The student model generates output based on the original
Original Instruation Fit
0.05
Original Instraction Fit
Middesing Patters Fit
Minlending Pattems 1
Original Instrution
Original Instractim
Mirkesdig Pattems Prined
Minlending Pateems Pruned
Probability
0.03
Probabilitt
0.03
0.02
0.02
0.01
0.01
0.00
0.0
0.2
0.4
0.6
0.8
1.0
0.00
0.0
0.4
0.6
0.8
1.0
Jaccard Similarity
Jaceard Similurity
(a) Jaceard Similarity in Code Corpus
(b) Jaceard Similarity in Math Corpus
Figure 13: Jaccard similarity distribution between student model responses (original vs. instruction pruned distracting patterns) and ground-truth responses on math and code datasets.
instruction. (2) Instruction pruned distracting patterns: The student model generates output from an instruction where confounding tokens are masked. To quantify the similarity between the model outputs,
widely adopted metrics: Jaccard Similarity (54). This metric enables a practical evaluation of how well the student model's 's output aligns with the contextual meaning conveyed by the teacher model.
Results. Figure 13 show a shift in Jaccard Similarity distribution for responses generated by the student model on both code and math tasks after removing confounding tokens. This indicates that, by ignoring distracting patterns, the student model not only improves reasoning accuracy (Figure I] but also generates responses more aligned with the teacher model, thereby enhancing output quality.
B Detailed Evaluation Settings
For mathematical problem solving, we evaluate on GSMSK, MATH, and OlympiadBench using
the Step-DPO framework [24] with modifications address its data extraction inaccuracies. For code generation, we report pass @ 1 on HumanEval(+)) [6] and LeetCode [9. using EvalPlus [32).
and pass @ 10 on LiveCodeBench [22], using the Skythought-Evw question answering, we evaluate on HotpotQA (50). 2WikiMultiHopQ) [19] and Musique [44
framework 25). For multi-hop
using the LongMab-PO framework 122 with all datasets sourced from LongBench [
C Collective Masking vs. Span Masking
We compare two pruning strategies (Figure5): Collective Pruning, which simultaneously masks all identified confounding tokens, and Span Pruning, which masks only contiguous spans of tokens.
Results. Table presents results for LLaMA3.2 1B/3B B-2211//BBnsttuue models trained on Math corpus (86K) and evaluated on MATH-500 under two pruning schemes. (1) We find that span pruning substantially outperforms both collective Pruning and native knowledge distillation. (2) Collective Pruning not only fails yield improvements but even degrades performance on the L.a.M33.--3B. Instruct model. We suspect that pruning all confounding tokens at once disrupts the training data's 'ss
16

---

semantic coherence, undermining the model's learning Consequently, we adopt the Span Pruning strategy for all subsequent experiments.
Table 3: Comparison two pruning strategies: Collective Pruning vs. Span Pruning on MATH-500 for respectively.
The best and second-best results marked in bold and underlined.
Model
MATH-500
1.-st33.-/BBImstruuc
Instruct Model (Pre-KD)
24.20
KD w/o Mask
34.00
Collective Pruning
34.20
Span Pruning
37.40
L..MA3.2-3B-Insmt
Instruct Model (Pre-KD)
42.80
KD wo Mask
50.00
Collective Pruning
49,20
Span Pruning
54.40
D Computational Overhead Analysis
We report detailed runtime and memory analyses of LeaF compared with standard knowledge distillation (KD) under identical hardware configurations.
Gradient Computation. Gradient difference computation in LeaF a one-time offline process on 8xNVIDIA A A00 (80GB) GPUs, jointly computing gradients from the teacher and student models. Processing around 7K samples takes about 3 hours.
Counterfactual Generation. Counterfactual responses are generated offline using vLLM on 4xNVIDIA A100 (80GB) GPUs. Processing 26K samples requires approximately 50 minutes for smaller models (e.g.. LLaMA-3.2 1B-Isttttttt and up to 2.85 hours for larger ones (e.g.. LLaMA- 3.3-70B-Instruct (wwen2.5-72B-Innttree This step is parallelizable and executed only once.
Training Overhead. We further measure the end-to-end training overhead over epochs 4xNVIDIA A100 (80GB) GPUs. Results are reported in Table LeaF incurs an additional 10-13% training time compared standard KD.
Table 4: End-to-end runtime comparison between standard KD and LeaF.
Model
KD Runtime (h)
LeaF Runtime (h)
Overhead (%)
(LaMA-3.221B 3-2-1B-Instruc
26.2
29.2
+11.5
LLaMA-3.2-3BB 3.-MA--.--3B-Intrruc
23.7
26.2
+10.6
Qwen2.5-Mee -.t.5----tt-158
27.3
30.9
+13.3
Results. LeaF introduces a moderate training overhead of 10-13% compared to standard KD, yet consistently yields 2-3% absolute accuracy gains on mathematics, coding, and multi-hop QA benchmarks. This improvement indicates that the additional computation is effectively utilized to enhance reasoning performance. Its auxiliary procedures, including gradient normalization, span pruning, and counterfactual generation, are one-time, offline, and fully parallelizable, incurring no extra cost during training or inference. These properties make LeaF practical and scalable for large-scale applications.
17

---

E Robustness Analysis
To evaluate test-time robustness, we assess LeaF on perturbed versions of the MathBench benchmarks (GSM8K, MATH-500, and OlympiadBench). Perturbations generated through standard data augmentation technique for producing realistic paraphrastic variations.
Table 5: Performance perturbed MathBench datasets for Instruct models (LLaMA-3.2-IB/BB Instruct) across different pruning strategies.
Model
GSM8K
MATH-500
OlympiadBench
Avg.
LLaMA 3.2-1B-Instrucl
Instruct (Pre-KDD)
39.65
24.80
3.86
22.77
KD (no mask)
50.42
31.80
5.19
29.14
LeaF (Instr Mask)
51.10
32.00
6.53
29.88
LeaF (Instr & Resp Mask)
51.71
34.20
6.97
30.96
LLaMA 3.2-3B-Instruct
Instruct (Pre-KD)
70.43
40.80
9.64
40.29
KD (no mask)
74.00
48.20
14.09
45.43
LeaF (Instr Mask)
74.07
50.20
15.58
46.62
LeaF (Instr & Resp Mask)
76.12
51.20
19.88
49.07
Results. Table the results for LLaMA-3.2-IBB3B- Instrrcct models evaluated on perturbed MathBench datasets. (1) LeaF consistently outperforms standard KD across noisy conditions, indicating that its pruning and attribution mechanisms are robust to moderate linguistic perturbations. (2) The performance degradation under noise remains within 2-3%, suggesting that LeaF effectively preserves reasoning capability even when input distributions shift. These findings highlight LeaF's robustness for real-world scenarios involving input perturbations or distributional noise.
F
Ablation Study: Importance of Contrastive Pairs
We include an ablation setting (w/o Student- Student-Wrong Originals) to examine the impact of excluding original samples that the student model fails to solve due confounding tokens. In contrast, LeaF leverages both student-correct and student-wrong originals together with counterfactual samples.
Table 6: Ablation results MathBench comparing LeaF and w/o Student- Wrong Originals. The ablated variant uses only student-correct originals and counterfactuals, excluding student- student-wrong wrong ones.
Model
GSM8K
MATH-500
OlympiadBench
Avg.
LLaMA 3.2-IB-
LeaF
57.70
35.40
10.09
34.40
w/o Student- Student-Wrong Wrong Originals
58.15
34.80
7.42
33.46
LLaMA 3.2-3B-Instruct
LeaF
83.09
51.80
20.77
51.88
w/o Student- Student-Wrong Wrong Originals
84.08
47.80
16.02
49.30
Results. (1) LeaF outperforms the ablated variant on MATH-500 and OlympiadBench, demonstrat- ing that retaining original samples containing misleading patterns alongside counterfactual samples provides a stronger contrastive signal. This helps the student model downweight spurious correlations and attend causally relevant information. (2) Excluding student-wrong originals reduces exposure to misleading cases, limiting the model's 's ability to generalize under confounding conditions. (3) The slight performance gain of the ablated variant on GSM8K stems from dataset composition. Compared with MATH-500 and OlynpiadBench, GSM8K is simpler dataset, and after filtering, its student-correct samples constitute a larger share of the training data. This distributional bias favors casier problems, leading to higher accuracy GSM8K.
18

---

G Threshold Stability Across Tasks
We further examine the stability of the threshold (Tmis) across domains, including math (GSM8K, MATH and OlympiadBenchh, code (HumanEval+, LeetCode and LivecodeBench( V4)), and question answering (SQuAD 2.0). For each task, model accuracy is evaluated at three threshold settings (Tmis 0.05, 0.10, and 0.15) using LLaMA-1B 1AMMA--B-Instruct and LLaMA-3B B--ssruut as student models. Table 7 summarizes the results.
Table 7: Accuracy under different Tiis thresholds across tasks.
Task
Model
7=0.05
7=0.10
r=0.15
LLaMA-332-IBB-nsttn
32.43
34.40
33.17
MathBench
LLaMA-3.2-3B-Insste
51.88
51.07
50.55
LLaMA-3,2-IB-Instte
17.85
18.27
19.76
CodeBench
LLaMMA-32-9B-IIssmmcc
32.83
33.55
32.95
LLaMA-3.2-1B-IIntte
72.12
73.33
74.78
SQuAD 2.0
LLaMA-3.2-3B-Instne
88.67
89.44
83.66
Results. (1) Tmis 0.10 yields the best or near-best performance, indicating that this threshold
CaAMMA--BInnsrrce (B-Instruct) exhibit mild sensitivity to value is robust across domains. (2) Smaller models (LLaMA- larger thresholds, while larger models (LLaMA-33 BB-A---r-nnt-tt remain stable around = 0.10. (3) These findings indicate that Tmis 0.10 achieves stable trade-off between filtering noisy gradients and preserving useful learning signals.
H Statistics of Counterfactual and Original Samples
We report token-length statistics and sample counts for the original and counterfactual datasets used in training. Results are presented separately for the math and code tasks, covering both LLaMA-3.2. IB- Instructt and LLaMA-3.2-33 1.A-MA-..2-3B-Instii Insstructt For each subset. include the minimum, maximum, and average token lengths, along with total sample counts and subtask-level breakdowns.
Table 8: Statistics of original and CF (counterfactual) samples for the math task.
Model
Min
Max
Avg.
Total
AIME
MATH
GSM8K
OlympiadBench
Original Samples
23
2711
98.50
12000
3000
3000
3000
3000
LLaMA-3.2-1B CF
24
2710
160.78
4576
1620
1029
907
1020
LLMMA-3.2-BB CF
23
2711
148.79
2843
1228
709
172
734
Table 9: Statistics of original and CF (counterfactual) samples for the code task.
Model
Min
Max
Avg.
Total
Original Samples
34
527
140.17
12000
LLaMA-3.2- IB CF
47
507
158.61
6210
I.AMMA3.2--BR CF
46
446
159.72
4059
Results. Table and Table summarize the distributional properties of the training data used in the main experiments. We observe that the average token length of counterfactual samples slightly higher than that of the original samples in both the math and code tasks. This pattern likely arises because longer and more complex problems tend include misleading patterns for student models, increasing the likelihood of counterfactual generation.
19

---

I Extended Related Work
Reasoning Consistency. A series of existing works have focused on decoding-phass consistency. Self-Consistency [46] stabilizes final answers by sampling multiple reasoning chains and applying majority voting. To further reduce inference costs, Adaptíve Consistency (2] and Early Scoring
Self-Consist 27 introduce Dirichlet-baser and window-based stopping criteria, respectively, Reasoning Aware Self-Consistency [45] enhances answer consistency by weighting sample quality and reasoning- importance. However, these methods primarily focus on stabilizing the final answer without addressing the model's ability to maintain and leverage key contextual information throughout the generation process. We instead define consistency as both answer stability and context adherence, and introduce distillation strategy that systematically suppresses the influence of misleading signals, enhancing the student's focus on key contextual information and promoting more robust context adherence during inference.
Detailed Training Settings
The complete training hyper-paramen parameter in knowledge distillation are put in Table T0
Table 10: Training hyper-parameters in Knowledge Distillation.
Model
Hyper-parameter
Value
LR
10
LR Scheduler
cosine
Batch Size
64
LLaMA3.2- IB- Instrucct
Epochs
3
Maximum Sequence Length
4096
Warmup Steps
5
Distill Loss Type
KL
Validation Set Size (Math)
1035
Validation Set Size (Code)
2000
LR
x 10
LR Scheduler
cosine
Batch Size
64
LLaMA3.2-3B-Inss
Epochs
Maximum Sequence Length
3000
Warmup Steps
Distill Loss Type
KL
Validation Set Size (Math)
1035
Validation Set Size (Code)
2000
LR
10-5
LR Scheduler
cosine
Batch Size
32
Owen2.5-Math-11.
Epochs
3
Maximum Sequence Length
4096
Warmup Steps
5
Distill Loss Type
KL
Validation Set Size (Math)
1200
Validation Set Size (Code)
2000
K
Evaluation Benchmarks
GSM8K g comprises 8500 grade-school-level word problems, each requiring 2-8 steps of basic arithmetic. natural-language diversity and multi-step structure make a standard measure for chain-of-thought prompting.
20

---

MATH [7 contains 12500 competition-styye problems grouped into seven topics (Prealgebra, Algebra, Number Theory, Counting Probability, Geometry, Intermediate Algebra, Precalculus). Every question is accompanied by detailed solution.
OlympiadBench [T6] originally bilingual, multimodal collection of 8476 Olympiad-level math and physics problems. We filter out proof-basel and image-based questions to obtain 674 pure-text tasks, enabling focused evaluation of advanced symbolic reasoning.
HumanEval+ [32] extends the original HumanEval with additional Python programming tasks and augmented unit tests, targeting functional correctness across diverse code patterns.
LeetCode samples real-world algorithmic challenges from the LeetCode platform arrays, trees, dynamic programming, etc. to assess models' ability to generate correct and efficient solutions.
LiveCodeBench 22 provides large-scale suite of real-world coding tasks with comprehensive unit tests and human preference annotations, allowing evaluation of both functional accuracy and coding style.
HotpotQA 500 consists of question-answer pairs that require reasoning over multiple Wikipedia paragraphs. It evaluates models' ability to retrieve and combine evidence from different sources to answer complex questions.
2WikiMultiHopQA [9] includes multi-hop questions automatically generated from Wikipedia, where each question involves two entities and demands reasoning across at least two documents.
Musique provides multi-hop questions with fine-grained decomposition and supporting evidence. It is designed to test compositional reasoning and factual consistency across multiple textual contexts.
L Open-source Instruct Models
The download links for the four open-source models are provided below:
Laama-3.--IB-nnstun
Hlmma-3.23B-Instrrcc
http://tanggiigaaeccccocmmmmmmmm Instruct
1imma--33-70-Insstrue
http://hwnggnnggeeccooooommmmmmm -t
Qwen2.5-Math -aaa255-
http://tnggiigggccccoooo///////2
Qwen2.5-72BB-nstrruct https//hhggiiiiicccooo/////////2
M Limitations
Our work has the following limitations: (1) Dependence on an advanced teacher model: Our method relies on teacher model to identify confounding tokens. Exploring self-improse ielf-improvement that enables models refine their attention critical tokens and boost reasoning can be an interesting and important future work. (2) Limited scalability to long-form generation: Due to the inherent limitations of the student model, we have only validated our approach on math and code tasks, leaving its applicability to long-text generation and other domains for future investigation.
21

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA] is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
In Abstract section and Section 1.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer:[Yee
Justification: In the Appendix.
22

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't I acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: Our paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
In Section 1.
Guidelines:
The answer NA means that the paper does not include experiments.
23

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. w
For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g.. large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
In our supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ /ttppp.cc/pbllc/guudeeee http://oeeuuiiisssoooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
24

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer:
[Yes]
Justification: In Section 3 and Section 4.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: In the Section 3 and Section 4.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: In Section
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
25

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer:
[Yes]
Justification: The authors have reviewed the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Here is no societal impact of the work performed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: Our paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
26

---

Released models that have high risk for misuse or dual use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We cite existing assets in Section 3.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing nor research with human subjects.
27

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the L.M used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [Yes]
Justification: In Section 3.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (ttr::e///errrppccccooofrrrreerp for what should or should not be described.
28


Reviews:
[{"review_id": "Oyiq8I4d80", "reviewer": "Reviewer_2JsC", "summary": "This paper describes Causal Attention Distillation (CAD), a novel two-stage framework that enhances reasoning in Large Language Models (LLMs) by eliminating attention to spurious markers. Through gradient-guided comparisons between teacher and student models (e.g., LLaMA3.3-70B → LLaMA3.2-1B), CAD first identifies misleading tokens that interfere with causal reasoning, and then masks them during distillation to enhance attention to critical contexts. Based on Perl's causal model, CAD achieves d-separation of confounding variables for consistent gains. This work combines causal theory with practical distillation techniques to provide a scalable solution for achieving robust attention adjustment in LLM.", "questions": "Q1. The paper states that spurious confounders are a general issue for LLMs, but Figures 1 and 2 only demonstrate this problem in a 1B model—not in the 70B variant. Could the authors clarify whether this issue is specific to low-parameter models or distillation settings, rather than a universal phenomenon across model scales?\n\nQ2. The paper lacks analysis of the threshold parameter τ, which plays a key role in identifying misleading tokens. Could the authors provide insights into how τ was selected, whether negative values were considered, and how sensitive the model’s performance is to different τ values across datasets?\n\nQ3. Sections 4.1 and 4.2 present experimental results, but it is unclear which model(s) these results are based on. Could the authors specify the model architecture and size used in each subsection to improve reproducibility?\n\nQ4. In Table 1, the proposed methods fall behind some other methods in some datasets—for example, the Instruct Model or KD w/o Mask occasionally outperforms the masking-based methods in LeetCode and Livecode-Bench. Could the authors analyze or explain the possible reasons behind these anomalies?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n- The paper effectively identifies and addresses LLMs’ susceptibility to spurious tokens during reasoning by introducing the Causal Attention Distillation (CAD) framework. This method improves attention robustness and aligns well with causal principles.\n\n- The paper is clearly structured, with causal graphs, heatmaps, and a well-designed case study that highlights how CAD focuses on critical information and avoids misleading tokens. These visual and narrative aids make the method accessible.\n\n- CAD achieves consistent improvements across models and tasks, showing strong practical utility. Its compatibility with existing distillation workflows allows for easy integration and enhances interpretability.\n\nWeaknesses\n\n-  While the paper claims spurious confounders are common across LLMs, Figures 1 and 2 only show this issue in a 1B model—not in the 70B model. This suggests the problem may stem from insufficient model capacity rather than a universal issue, and should be framed within low-parameter or distillation contexts.\n- The paper lacks parameter analysis, especially regarding the threshold τ. τ critically affects performance and generalization; intuitively, it should be negative, as τ = 0 implies equal token sensitivity between student and teacher, contradicting the masking purpose. Its optimal value may also vary across datasets.\n- Some experimental details are unclear, such as which model is used in Sections 4.1 and 4.2. This ambiguity affects the reproducibility and clarity of the experimental setup.", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: The paper states that spurious confounders are a general issue for LLMs, but Figures 1 and 2 only demonstrate this problem in a 1B model—not in the 70B variant. Could the authors clarify whether this issue is specific to low-parameter models or distillation settings, rather than a universal phenomenon across model scales?**\n\n**A1:** We would like to clarify that spurious confounders are not specific to low‑parameter models or distillation settings, but are a universal phenomenon across model scales. To demonstrate this, we extend the analytical experiments from the Introduction (Appendix A.2.2) to larger models, such as **Llama‑3.3‑70B‑Instruct** and **Qwen2.5‑72B‑Instruct**, using the same 12 K training examples. This enables the models to identify misleading tokens without relying on a teacher model. The specific experimental procedure is as follows:\nFirst, we select the subset where the large model produces incorrect results. Then, we compute each sample’s gradient of the model to identify potential misleading patterns. Finally, we remove these misleading patterns from the instruction and re‑evaluate the model’s prediction.\n\n| Model                         | Threshold | Originally Wrong | Corrected After Removing Confounders | Improvement |\n|-------------------------------|-----------|-----------------:|-------------------------------------:|------------:|\n| Llama‑3.3‑70B‑Instruct        | 0.10      |            2,112 |                                  453 |      21.45% |\n| Qwen2.5‑72B‑Instruct          | 0.10      |            2,604 |                                  613 |      23.54% |\n\n\nThe observed **21.45%~23.54%** improvement on 70 B models after confounder removal **demonstrates that spurious confounders also degrade the reasoning accuracy of large-scale models**. Therefore confirming that spurious confounders are a universal LLM phenomenon, not an artifact of small‑scale or distillation‑only settings.\n\n**Q2(a)：How was the threshold parameter τ selected?**\n\n**A2(a)**: Due to space limitations in the main text, a detailed analysis of the threshold parameter τ is provided in the supplementary material, specifically in Appendix A.1, titled \"Threshold Sensitivity Analysis.\"\n\nIn this analysis, we considered both instruction-level and response-level performance. Our findings are as follows:\n\n**1.Instruction-level**: LLaMa3.2-CAD-1B performs best at a threshold of 0.10, while LLaMa3.2-CAD-3B performs best at 0.05.\n\n**2.Response-level**: LLaMa3.2-CAD-1B performs best at 0.15, and LLaMa3.2-CAD-3B at 0.10.\n\nFor both instruction-level and response-level, LLaMa3.2-CAD-1B achieves optimal performance at a higher misleading token threshold than LLaMa3.2-CAD-3B. This suggests that smaller models, which are more sensitive to misleading tokens, benefit from a higher threshold that more effectively filters out disruptive tokens.\n\n**Q2(b)：Were negative values of τ considered?**\n\n**A2(b):** Since the threshold parameter τ represents the percentage of misleading tokens relative to instruction tokens, there is no need to consider negative values.\n\n**Q2(c)：How sensitive is the model's performance to different values of τ across datasets?**\n  \n**A2(c):** To evaluate how the choice of τ affects performance, we measured each model’s accuracy on three benchmarks (Math, Code, and SQuAD 2.0) at τ = 0.05, 0.10, and 0.15:\n| Task       | Model                | τ=0.05  | τ=0.10  | τ=0.15  |\n|------------|----------------------|:-------:|:-------:|:-------:|\n| **Math**   | LLaMa‑1B‑Instruct    | 32.43%  | **34.40%**  | 33.17%  |\n|            | LLaMa‑3B‑Instruct    | **51.88%**  | 51.07%  | 50.55%  |\n| **Code**   | LLaMa‑1B‑Instruct    | 17.85%  | 18.27%  | **19.76%**  |\n|            | LLaMa‑3B‑Instruct    | 32.83%  | **33.55%**  | 32.95%  |\n| **SQuAD 2.0** | LLaMa‑1B‑Instruct  | 72.12% | 73.33%  | **74.78%**  |\n|             | LLaMa‑3B‑Instruct   | 88.67%  | **89.44%**  | 83.66%  |\n\nWe observe that LLaMa‑1B‑Instruct generally achieves its best performance at \\( τ = 0.15\\), while LLaMa‑3B‑Instruct consistently performs best at \\( τ = 0.10\\) across these tasks. This indicates that the threshold \\( τ_mis \\) is relatively stable in different datasets.\n\n**Q3: Sections 4.1 and 4.2 present experimental results, but it is unclear which model(s) these results are based on. Could the authors specify the model architecture and size used in each subsection to improve reproducibility?**\n\n**A3:** Thank you for the suggestion. The base model used in both Section 4.1 and Section 4.2 is **Llama‑3.2‑3B‑Instruct**, and we will clarify these settings more explicitly in the future version.\n\n**Q4：In Table 1, the proposed methods fall behind some other methods in some datasets—for example, the Instruct Model or KD w/o Mask occasionally outperforms the masking-based methods in LeetCode and Livecode-Bench. Could the authors analyze or explain the possible reasons behind these anomalies?**\n\n**A4:** We attribute these anomalies primarily to **the limited coding capability of LLaMa3.2‑1B‑Instruct**.  \nTo validate this, we conducted a preliminary experiment using Kodcode’s 5K synthetic LeetCode dataset.  After removing misleading tokens, LLaMa3.2‑1B‑Instruct solved only **264 additional problems (+5.28%)**, whereas **Evol** and **OSS** achieved gains of  **909 (+15.2%)** and **992 (+16.5%)**, respectively—a performance gap of around **10%**.  This clearly suggests that most of the failures by LLaMa3.2‑1B‑Instruct on LeetCode and Livecode-Bench are **due to inherent limitations in its coding and reasoning capabilities**, rather than misleading patterns.  As a result, even improved focus on key information cannot significantly boost its performance.\n\nIn contrast, larger models such as LLaMA‑3.2‑3B‑Instruct and Qwen2.5‑Math‑1.5B exhibit **stronger generalization and reasoning capabilities**, which enables them to more effectively learn to disregard misleading patterns and focus on truly informative tokens. This improved attention behavior can then be **successfully transferred to more complex tasks like LeetCode and Livecode-Bench**, leading to more consistent performance gains."}, {"review_id": "puSkaqOh5Q", "reviewer": "Reviewer_qdQ9", "summary": "This paper proposes Causal Attention Distillation (CAD), a two-stage framework for enhancing the reasoning ability and interpretability of student large language models (LLMs) by distilling causal attention from a more advanced teacher. In the first stage, distracting or confounding tokens in training data are identified through gradient-based comparisons between teacher and student models. In the second stage, the student model is distilled not only on original inputs but also on counterfactual samples created by pruning confounders, using a hybrid loss. Experimental results on mathematical reasoning and code benchmarks indicate that CAD improves student model accuracy and shifts attention toward more causally relevant tokens, enhancing robustness and interpretability.", "questions": "- How sensitive is confounding token identification in CAD to potential errors or misattributions in the teacher model? If the teacher misfocuses or mispredicts, could CAD propagate harmful attention patterns?\n- The method depends on clear differences between teacher and student gradient attributions. Have the authors tested robustness under adversarially inserted confounders or in noisy real text, especially for models not tuned for math/code?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n- The paper is motivated pretty well and carefully written.\n- Experiments span multiple settings and benchmarks. The main results show consistent accuracy gains over standard knowledge distillation (KD), especially when both instruction and response-level masking are used.\n\n\n**Weaknesses:**\n\n- CAD’s reliance on a large and advanced teacher to guide the token importance. This limits its application when such teacher models are unavailable or computationally prohibitive.\n- The proposed approach requires teacher-student training and extra steps to compute and normalize gradients, generate counterfactuals, and perform span pruning. This could add extra computational overhead that can't be omitted when scaling up?\n- It seems the magnitude of improvement over standard KD (about 2-3% absolute) is rather moderate. For the largest/strongest teacher-student pairs, the relative improvement is smaller. \n- The efficacy of the proposed method is assessed mainly on mathematical and code-generation tasks where the reasoning is relatively structured. I wonder if other scenarios like open-domain QA and long-text summarization (I later noticed the authors mentioned this as well in the limitation section).", "was_revised": true, "final_justification": "The initial score still reflects my evaluation to this paper after carefully checking out the authors' rebuttal.", "rebuttal": "**Q1: CAD’s reliance on a large and advanced teacher to guide the token importance. This limits its application when such teacher models are unavailable or computationally prohibitive.**\n\n**A1:** We acknowledge in **Appendix A.8 Limitation** that CAD’s reliance on a large, advanced teacher model is indeed a limitation. However, this dependency does not diminish the value of our approach. Beyond the well-known gap in reasoning capabilities, smaller models also lack a critical ability to focus attention on key information. CAD enhances standard distillation by guiding small models to focus on the critical information large models use during reasoning. Such guidance is especially crucial for enabling smaller models to effectively handle challenging reasoning tasks.\n\nFurthermore, **our experiments in Section 4.2 validate the importance of teacher models**. Specifically, CAD is compared against random masking and PPL-masking on Olympiad problems. The results show that PPL-masking performs similarly to or worse than random masking and standard distillation, **indicating that smaller models without external guidance often fail to identify genuinely important tokens in complex reasoning tasks**. By leveraging teacher models to explicitly highlight these essential tokens, CAD substantially improves the reasoning performance of smaller models, thereby unlocking their greater potential.\n\n**Q2: The proposed approach requires teacher-student training and extra steps to compute and normalize gradients, generate counterfactuals, and perform span pruning. This could add extra computational overhead that can't be omitted when scaling up?**\n\n**A2:** We quantified the additional computational overhead of CAD compared to standard KD by measuring end-to-end runtime using identical hardware:\n\n**Gradient computation (one-time, 8× NVIDIA A100 80GB):**\nComputing gradients for 7K samples consistently takes about 3 hours, regardless of the student model size (LLaMA‑3.2‑1B, LLaMA‑3.2‑3B, or Qwen2.5‑Math‑1.5B), when using larger teacher models (LLaMA‑3.3‑70B or Qwen2.5‑72B-Instruct).\n\n**Counterfactual generation (one-time, offline via vLLM, 4×NVIDIA A100 80GB):**\nGenerating counterfactuals for 26K samples requires around 50 minutes for smaller models (e.g., LLaMA‑1B) and approximately 2.85 hours for larger models (e.g., LLaMA‑70B or Qwen2.5‑72B). This process is easily parallelizable.\n\n**Training overhead (3 epochs 4×NVIDIA A100 80GB):**\n| Model                  | KD Runtime | CAD Runtime | Overhead  |\n|------------------------|-----------:|------------:|----------:|\n| LLaMA‑3.2‑1B           |     26.2 h |      29.2 h |   +11.5%  |\n| LLaMA‑3.2‑3B           |     23.7 h |      26.2 h |   +10.6%  |\n| Qwen2.5‑Math‑1.5B      |     27.3 h |      30.9 h |   +13.3%  |\n\nOverall, CAD incurs only **10–13%** additional training overhead compared to standard KD, while achieving 2–3% absolute accuracy improvements in mathematics and coding tasks. Importantly, the extra data-processing steps—gradient normalization, span pruning, and counterfactual generation—are one-time, offline operations, fully parallelizable, and do not recur during inference or downstream fine-tuning, making CAD practical to scale.\n\n**Q3: It seems the magnitude of improvement over standard KD (about 2-3% absolute) is rather moderate. For the largest/strongest teacher-student pairs, the relative improvement is smaller.**\n\n**A3:** We believe that the observed 2–3% absolute improvement over standard KD is substantial given the strong performance of **the CoT KD baselines**. \nSpecifically, **for the Code task**, standard KD itself achieves an average improvement of **3.86%** over the base models, while our CAD approach further adds an average improvement of **2.39%** on top of KD—indicating that CAD provides gains nearly on par with those of KD itself. \n\n**For Math tasks**, COT KD is already a strong baseline, yet CAD consistently yields an additional 2–3% absolute improvement, clearly demonstrating meaningful incremental value. \n\nMoreover, as shown in our supplementary results on **SQuAD v2** (see Q4), **CAD achieves an average 4% improvement** on QA tasks—highlighting its **robustness and general applicability** across Math, Code, and QA domains.\n\n**Q4: The efficacy of the proposed method is assessed mainly on mathematical and code-generation tasks where the reasoning is relatively structured. I wonder if other scenarios like open-domain QA and long-text summarization (I later noticed the authors mentioned this as well in the limitation section).**\n\n**A4:** We acknowledge in Appendix A.8 that, due to small models’ current limitations with long texts, CAD has constraints in the long‑document domain.\nTo test CAD beyond math and code, we evaluated on the Stanford Question Answering Dataset 2.0 (SQuAD 2.0), where each answer is a text span from the passage (or the question may be unanswerable). We sampled 3.5 K training examples and measured exact‑match (EM) on 900 test items.\n\n**Results (EM):**\n\n**Llama‑1b‑Instruct:** 41.11% → KD: 72.56% → CAD: 74.78%\n\n**Llama‑3b‑Instruct:** 73.67% → KD: 83.67% → CAD: 89.44%\n\nThese results show that, with sufficient training data, CAD not only boosts performance on math and code tasks but also generalizes effectively to open‑domain QA.\n\n**Q5: The method depends on clear differences between teacher and student gradient attributions. Have the authors tested robustness under adversarially inserted confounders or in noisy real text, especially for models not tuned for math/code?**\n \n**A5:** To evaluate the robustness of our method, we conducted analyses from both the training and testing perspectives.\n\nFirst, regarding **training data**, we used standard and widely adopted datasets—NuminaMath for math, AceCoder for code, and SQuAD 2.0 for question answering. Our approach constructs distillation data from these clean corpora and **does not introduce noise during training**. Moreover, due to the **inherent capability gap between teacher and student models**, clear differences in gradient attributions naturally arise on these standard datasets, ensuring the generality of our method without requiring adversarial noise.\n\nSecond, to assess **test-time robustness**, we introduced noise into evaluation datasets—including Math benchmarks (GSM8K, MATH-500, OlympiadBench) and SQuAD 2.0 (900 samples)—via back-translation, a common data augmentation technique for generating realistic input perturbations. The following table summarizes the results, demonstrating that our method maintains strong performance improvements even under noisy test conditions:\n\n| Dataset             | Model                  | Instruct | KD     | CAD(Instr Mask)    | CAD (Instr & Resp Mask) |\n|---------------------|------------------------|----------|--------|--------|------------|\n| **SQuAD 2.0**       | LLaMA‑3.2‑1B‑Instruct  | 45.89%   | 66.33% | 69.11% | –          |\n|                     | LLaMA‑3.2‑3B‑Instruct  | 69.22%   | 77.00% | 79.44% | –          |\n| **MathBench GSM8K** | LLaMA‑3.2‑1B‑Instruct  | 39.65%   | 50.42% | 51.10% | 51.71%     |\n|                     | LLaMA‑3.2‑3B‑Instruct  | 70.43%   | 74.00% | 74.07% | 76.12%     |\n| **MathBench MATH-500** | LLaMA‑3.2‑1B‑Instruct | 24.80%  | 31.80% | 32.00% | 34.20%     |\n|                     | LLaMA‑3.2‑3B‑Instruct  | 40.80%   | 48.20% | 50.20% | 51.20%     |\n| **MathBench Olympiad** | LLaMA‑3.2‑1B‑Instruct | 3.86%   | 5.19%  | 6.53%  | 6.97%      |\n|                     | LLaMA‑3.2‑3B‑Instruct  | 9.64%    | 14.09% | 15.58% | 19.88%     |\n\nThese results demonstrate that CAD remains robust even when noisy inputs are introduced at test time, maintaining and often improving upon standard KD performance. We thank the reviewer for highlighting this important aspect of robustness.\n\n\n**Q6: How sensitive is confounding token identification in CAD to potential errors or misattributions in the teacher model? If the teacher misfocuses or mispredicts, could CAD propagate harmful attention patterns?**\n\n**A6:** Our method includes** a filtering step** to ensure the reliability of misleading token identification. Specifically, we only retain cases **where the teacher model predicts the correct answer both before and after the identified tokens are removed**. This ensures that the removed tokens are not essential for correct reasoning, and that the resulting patterns do not reflect harmful or erroneous attention."}, {"review_id": "OObRrNt0ab", "reviewer": "Reviewer_ZMF9", "summary": "This paper introduces Causal Attention Distillation (CAD), a novel two-stage distillation framework that aims to improve reasoning consistency in large language models (LLMs). The authors adopt a causal perspective and treat misleading tokens as spurious confounders. CAD first identifies such misleading tokens through gradient-based sensitivity comparisons between a teacher and a student model. Then, these tokens are masked during training, enabling the student model to better align its attention with the teacher's on causally relevant parts of the input. The method is evaluated on both mathematical reasoning and code generation task, showing consistent performance gains over standard knowledge distillation and other masking strategies (random, PPL-based).", "questions": "1.How to ensure the stability of the gradient difference threshold (τ_mis) in cross-task scenarios?\n\n2.How does CAD compare to standard KD in terms of training time and memory usage? Are the additional costs of gradient comparison practical for large-scale applications? Could the gradient comparison be approximated more efficiently?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths: 1.The author combines causal inference with gradient analysis to effectively separate false correlations, providing a new approach for addressing attention bias in LLM reasoning.\n\n2.The paper systematically compares different masking strategies (random, PPL-based, gradient-based), as well as varying the number of response segments. These analyses strengthen the empirical claims and justify the design choices.\n\n3.The attention heatmap comparison between CAD and traditional KD adds interpretability, demonstrating how CAD reshapes the student’s focus toward causally relevant information，the case study (Section 4.3) clearly demonstrates the attention shift toward critical tokens after CAD training, helping to validate the underlying causal intuition.\n\nWeaknesses: 1.The method was mainly verified on the LLaMA and Qwen series models, and the generalization of other architectures (such as GPT, Claude) was not tested.\n\n2.The mathematical proof of the causal model is relatively brief and does not fully expound the theoretical connection between the DAG structure and attention optimization.\n\n3.The proposed method relies on computing token-level gradients for both teacher and student, which may incur high computational overhead. No time/memory cost analysis is provided.\n\n4.The decision to mask individual misleading patterns instead of all at once is motivated by empirical performance. However, the semantic trade-off (between removing noise and retaining context) could be formalized further, perhaps via information-theoretic analysis.", "was_revised": true, "final_justification": "After reading the authors' response and other reviewers' comments, my concerns have been addressed, thus I change my score from 3 to 4.", "rebuttal": "**Q1: The method was mainly verified on the LLaMA and Qwen series models, and the generalization of other architectures (such as GPT, Claude) was not tested.**\n\n**A1:** We acknowledge that our method was primarily validated on the LLaMA and Qwen model families. \n\n**Our approach aligns with the prevailing trend in recent model capability enhancement research, which mainly focuses on open-source models**. Closed-source models like GPT and Claude, however, restrict access to gradient information, which is crucial for identifying misleading tokens in our method. This practical constraint prevents direct application of our approach to these architectures.\n\nFurthermore, using teacher and student models from the same family (e.g., LLaMA-3.2-1B-Instruct and LLaMA-3.3-70B-Instruct) provides significant advantages. **Their shared architecture and similar pretraining data create compatible internal representations, enabling the student model to more effectively learn from the teacher through distillation**. This enhances the student’s ability to focus attention on critical information and ultimately improves reasoning performance.\n\n**Q2: The mathematical proof of the causal model is relatively brief and does not fully expound the theoretical connection between the DAG structure and attention optimization.**\n \n**A2:** Thank you for this suggestion. To make the link between our DAG formulation and attention optimization clearer, we will expand the theoretical derivation as follows:\n\nOur goal is to **eliminate spurious correlations** introduced by confounding tokens $\\(A\\)$. In the presence of a confounder $\\(A\\)$ that influences both $\\(X_i\\)$ and the output $\\(Y\\)$ (dashed arrows in Figure 3), the standard conditional distribution becomes:\n\n\n$P(Y | X_i = x) = \\sum_{a} P(Y \\mid X_i = x, A = a) \\, P(A = a \\mid X_i = x)$\n\n\nwhich departs from the true interventional distribution $\\(P(Y \\mid \\mathrm{do}(X_i = x))\\)$. The discrepancy arises because the term \\$(P(A \\mid X_i)\\)$ carries spurious influence from the path $\\(A \\to Y\\)$.\n\nWe perform **causal pruning** by intervening with $\\(\\mathrm{do}(A=\\emptyset)\\)$—i.e., masking out tokens in \\(A\\). In Pearl’s framework, this blocks the non-causal paths through \\(A\\), yielding:\n\n\n$P(Y \\mid \\mathrm{do}(X_i = x)) = P(Y \\mid X_i = x, A=\\emptyset)$,\n\n\nwhich restores the genuine causal effect of $\\(X_i\\)$ on $\\(Y\\)$.\n\nIn transformer models, **attention weights** serve as a soft selection mechanism over input tokens: higher attention on \\(X_i\\) implies a stronger dependency of the output on that token. By training on these interventional (counterfactual) inputs, the student’s attention distribution is **optimized** to focus on true causal parents—and away from confounders—thereby aligning its internal reasoning with the DAG’s structure.\n\nThis causal pruning both improves robustness by removing misleading cues and enhances interpretability: as demonstrated in our case study, CAD’s gradient distributions focus more sharply on the instruction’s key tokens compared to standard KD, providing clearer, more explainable reasoning behavior.\n\n\n**Q3: The proposed method relies on computing token-level gradients for both teacher and student, which may incur high computational overhead. No time/memory cost analysis is provided. How does CAD compare to standard KD in terms of training time and memory usage? Are the additional costs of gradient comparison practical for large-scale applications? Could the gradient comparison be approximated more efficiently?**\n\nWe quantified the additional computational overhead of CAD compared to standard KD by measuring end-to-end runtime using identical hardware:\n\n**Gradient computation (one-time, 8× NVIDIA A100 80GB):**\nComputing gradients for 7K samples consistently takes about 3 hours, regardless of the student model size (LLaMA‑3.2‑1B, LLaMA‑3.2‑3B, or Qwen2.5‑Math‑1.5B), when using larger teacher models (LLaMA‑3.3‑70B or Qwen2.5‑72B-Instruct).\n\n**Counterfactual generation (one-time, offline via vLLM, 4×NVIDIA A100 80GB):**\nGenerating counterfactuals for 26K samples requires around 50 minutes for smaller models (e.g., LLaMA‑1B) and approximately 2.85 hours for larger models (e.g., LLaMA‑70B or Qwen2.5‑72B). This process is easily parallelizable.\n\n**Training overhead (3 epochs 4×NVIDIA A100 80GB):**\n| Model                  | KD Runtime | CAD Runtime | Overhead  |\n|------------------------|-----------:|------------:|----------:|\n| LLaMA‑3.2‑1B           |     26.2 h |      29.2 h |   +11.5%  |\n| LLaMA‑3.2‑3B           |     23.7 h |      26.2 h |   +10.6%  |\n| Qwen2.5‑Math‑1.5B      |     27.3 h |      30.9 h |   +13.3%  |\n\nOverall, CAD incurs only 10–13% additional training overhead compared to standard KD, while achieving 2–3% absolute accuracy improvements in mathematics and coding tasks. Importantly, the extra data-processing steps—gradient normalization, span pruning, and counterfactual generation—are one-time, offline operations, fully parallelizable, and do not recur during inference or downstream fine-tuning, making CAD practical to scale.\n\n**Q4:The decision to mask individual misleading patterns instead of all at once is motivated by empirical performance. However, the semantic trade-off (between removing noise and retaining context) could be formalized further, perhaps via information-theoretic analysis.**\n\n**A4:** We appreciate the suggestion and agree that the trade-off between removing noise and preserving semantic context is an important aspect of our method.\n\nOur choice to mask misleading patterns incrementally—rather than masking all identified tokens at once—was guided by empirical observations: masking too aggressively, especially in instructions with limited length, often harms performance by disrupting essential context or breaking syntactic structures. In contrast, span masking allows the model to selectively suppress misleading cues while retaining enough context to preserve semantic coherence.\n\nWe fully agree that this semantic trade-off could benefit from formalization through information-theoretic approaches. For example, one could quantify the information gain associated with varying masking proportions to identify an optimal point that balances noise suppression and context retention. Additionally, metrics such as predictive entropy or the KL divergence between teacher and student attention distributions could serve as principled indicators for optimal masking strategies.\n\nAlthough our current strategy relies primarily on empirical evidence, we view a formal, information-theoretic analysis of this trade-off as an important and promising direction for future research. We greatly appreciate the reviewer highlighting this valuable perspective.\n\n**Q5: How to ensure the stability of the gradient difference threshold (τ_mis) in cross-task scenarios?**\n\n**A5:** To evaluate the stability of the gradient difference threshold(τ_mis) across different tasks, we measured each model’s accuracy on three benchmarks—Math, Code, and SQuAD 2.0—at (τ_mis)values of 0.05, 0.10, and 0.15:\n\n| Task        | Model               | τ=0.05  | τ=0.10  | τ=0.15  |\n|-------------|---------------------|:-------:|:-------:|:-------:|\n| **Math**    | LLaMa‑1B‑Instruct   | 32.43%  | **34.40%**  | 33.17%  |\n|             | LLaMa‑3B‑Instruct   | **51.88%**  | 51.07% | 50.55%  |\n| **Code**    | LLaMa‑1B‑Instruct   | 17.85%  | 18.27%  | **19.76%**  |\n|             | LLaMa‑3B‑Instruct   | 32.83%  | **33.55%**  | 32.95%  |\n| **SQuAD 2.0** | LLaMa‑1B‑Instruct  | 72.12%  | 73.33%  | **74.78%**  |\n|             | LLaMa‑3B‑Instruct   | 88.67%  | **89.44%**  | 83.66%  |\n\nWe observe that LLaMa‑1B‑Instruct generally achieves its best performance at \\( τ= 0.15\\), while LLaMa‑3B‑Instruct consistently performs best at \\(τ = 0.10\\) across these tasks. This indicates that the threshold \\(τ_mis\\) is relatively stable in cross-task scenarios."}, {"review_id": "EAWeExg1KW", "reviewer": "Reviewer_WSL6", "summary": "This paper proposes a knowledge distillation methodology that detects misleading tokens in inputs and outputs during LLM reasoning and makes the model robust to them. The authors argue that LLMs have tokens with spurious correlations unrelated to actual correct answer generation, which hinder accurate and consistent response generation during reasoning processes. To address this, the study defines misleading tokens as those with small gradient magnitude differences between teacher-student models, and performs knowledge distillation using augmented datasets that exclude these tokens. Through various experimental results, they demonstrate that this masking-distillation strategy improves the model's reasoning performance and show that the detected tokens indeed need to be removed for performance improvement.", "questions": "1. A detailed explanation of Token masking and Data Augmentation in Section 2.2 seems necessary. Misleading tokens appear to be detected independently, but looking at Section 4.2 experiments, it seems like consecutive spans are masked. How is the process of converting individual masking tokens to consecutive spans handled?\n2. From my understanding, misleading token masking is performed on consecutive spans. In Subsection 4.1, is random masking performed on arbitrary consecutive tokens? Or on independent tokens?\n3. By what criteria were misleading tokens identified in the preliminary experiment in Figure 1? Did this experiment use the same method as the proposed methodology's gradient-based approach?\n4. How is the threshold $\\tau_\\text{mis}$ in Equation 6 determined? While sensitivity analysis results are said to be in the Appendix, they're currently difficult to verify and more detailed explanation in the main text is needed.\n5. What is the process of defining misleading tokens as those with small gradient magnitude differences between teacher-student in Equations 5 and 6? An intuitive explanation seems necessary.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 4, "originality": 3, "strengths_and_weaknesses": "### Strengths:\n1. Addresses a frequently occurring phenomenon directly. Reasoning models indeed show phenomena where they incorrectly focus on very trivial elements during answer generation, and weaker models especially fail to properly exploit reasoning paths by focusing on wrong patterns. The motivation of this research is very intuitive regarding this issue, and resolving spurious correlations appears important.\n2. Attempted to interpret spurious tokens through a causal framework, trying to capture relationships not only in instructions but also within reasoning steps.\n3. Shows consistent performance improvements. As confirmed through Table 1, performance consistently improves across different model types and sizes compared to baselines.\n\n### Weaknesses:\n\n1. Lack of validity in core assumptions: While Table 1 shows that the proposed methodology brings consistent performance improvements, it's difficult to clearly distinguish whether this is truly due to \"misleading\" tokens or other factors such as sequence length reduction, noise reduction, or data augmentation.\n2. Limitations in misleading token identification method: There's insufficient theoretical/empirical evidence that gradient-based comparison (Equation 5) truly finds \"misleading\" tokens. Equations 5 and 6 appear to simply select tokens with small gradient differences between teacher-student models, and it's difficult to find experiments and evidence in the paper showing that these tokens have spurious correlations.\n3. Superficial application of causal framework: While citing Pearl's causal framework, it essentially amounts to simple token masking. Particularly, as seen in Figure 2's example, tokens detected as misleading actually appear to be necessary information for answer generation. This makes it difficult to argue that these tokens have spurious correlations and that masking resolves this.\n4. Weak baselines: Random masking baseline is too simple, and comparison with more advanced distillation methods is lacking. Direct comparison with CoT knowledge distillation methodologies mentioned in the paper's related works seems necessary. This could demonstrate that it's not simply about distilling CoT reasoning well, but that measures against spurious correlations are needed.", "was_revised": true, "final_justification": "After carefully considering the authors’ rebuttal and additional clarifications, I have decided to raise my score from 2 (Reject) to 4 (Borderline Accept). The authors have addressed most of my initial concerns regarding the technical soundness and experimental validation. In particular, they provided:\n- Additional experiments and ablation studies that strengthen the empirical evidence for the proposed method.\n- Clarifications on the design choices and assumptions, which resolve the ambiguities I previously noted.\n- More detailed description of the dataset composition and experimental setup, improving the reproducibility of the work.\n\nWhile certain aspects, such as the rigor of the causal framework, could still be improved, the overall contribution is now technically solid, and the reasons to accept outweigh those to reject. Therefore, I consider the paper a borderline accept.", "rebuttal": "**Q1: Lack of validity in core assumptions: While Table 1 shows that the proposed methodology brings consistent performance improvements, it's difficult to clearly distinguish whether this is truly due to \"misleading\" tokens or other factors such as sequence length reduction, noise reduction, or data augmentation.**\n\n**A1:** We address the concern by explicitly isolating the impact of misleading tokens from other confounding factors:\nFirst, we rule out **sequence length reduction** and **data augmentation** effects. In Section 4.2 (Random Masking setting), we applied masking with exactly the same proportion as CAD, ensuring identical amounts of sequence length reduction and data augmentation across methods. The results show that random masking provides no improvement—and in fact slightly degrades performance on GSM8K and Olympiad tasks compared to standard distillation—confirming that CAD’s gains cannot be attributed simply to reduced sequence lengths or increased data diversity.\n|                | GSM8K |  MATH | OlympiadBench |    Avg.   |\n|----------------|:-----:|:-----:|:-------------:|:---------:|\n| KD w/o Mask    | 82.87 | 49.00 |     18.99     |   50.29   |\n| Random Masking | 82.41 | 51.00 |   **14.54**   |   49.32   |\n| CAD            | 83.09 | 51.80 |   **20.77**   | **51.88** |\n\nSecond, we eliminate the confounding benefits from **distillation with large models and chain-of-thought (CoT)** by comparing our approach directly against a strong CoT-KD baseline (Table 1). \nCAD consistently delivers additional improvements over CoT-KD on both Math and Code tasks, highlighting that our approach provides value beyond these established techniques.\n\nFinally, we validate that CAD’s improved reasoning ability arises from better focus on critical information through interpretability analysis (In Section 4.3). As shown in Figure 8, CAD allows the model to attend more effectively to key problem elements such as “multiple” and “5” (e.g., combinations like (1,5), (2,5)), which directly relate to the condition “product of two numbers is a multiple of 5.” In contrast, the KD baseline fails to sufficiently focus on these tokens and instead considers invalid pairs (e.g., 6 and 3), leading to incorrect final answers.\n\n**Q2: What is the process of defining misleading tokens as those with small gradient magnitude differences between teacher-student in Equations 5 and 6? An intuitive explanation seems necessary.**\n\n**A2:** We appreciate the opportunity to clarify our method of defining misleading tokens, as there seems to be a misunderstanding.\nWe do **not simply select tokens with small gradient magnitude differences** between teacher and student models. Instead, our approach specifically targets tokens on which the **student model excessively relies** (high gradient magnitude) but which the **teacher model largely ignores** (low gradient magnitude), particularly on samples where the student makes an incorrect prediction while the teacher correctly solves the task.\n\nMore concretely, we first select examples for which the student model is incorrect and the teacher model correct. Next, we employ Equations 5 and 6 to compute gradients with respect to tokens, and identify tokens exhibiting a significant discrepancy—namely, large student gradients versus small teacher gradients—as candidate misleading tokens.\n\nCrucially, after removing these candidate tokens, the student’s predictions often change from incorrect to correct, while the teacher’s correct reasoning remains unaffected. This indicates that the student initially paid undue attention to these tokens at the expense of genuinely informative tokens. Consequently, we designate these tokens as **misleading tokens** and use them to generate counterfactual training samples, guiding the student model to better align its attention with truly relevant information.\n\n**Q3(a)：Superficial application of causal framework: While citing Pearl's causal framework, it essentially amounts to simple token masking.**\n \n**A3(a):** We appreciate the opportunity to clarify why CAD represents more than simple token masking and how it aligns with Pearl’s causal framework.\n\nOur preliminary experiments showed that directly removing certain \"misleading\" tokens could flip a model’s prediction from incorrect to correct—even **without retraining**. This aligns closely with Pearl’s causal perspective: a subset of input tokens ($A$) can act as confounders, creating a spurious correlation ($A \\to Y$). By explicitly performing a causal intervention—masking these misleading tokens—we break the spurious pathway, restoring the true causal relationship ($X \\to Y$). CAD thus leverages these carefully constructed counterfactual examples during distillation, teaching models to ignore spurious signals. \nThus, CAD is not merely token masking—it is a principled causal intervention designed explicitly to eliminate spurious correlations while preserving essential information.\n\n**Q3(b)：Particularly, as seen in Figure 2's example, tokens detected as misleading actually appear to be necessary information for answer generation. This makes it difficult to argue that these tokens have spurious correlations and that masking resolves this.**\n\n**A3(b):** While the tokens identified as misleading (e.g., in Figure 2) may appear necessary at first glance, the information they carry is often redundant and can be inferred from other parts of the prompt.\nIn Figure 2, the phrase **“the number of bacteria doubles every day”** is marked as misleading patterns. Although it appears informative, this detail is already implied by the earlier sentence: **“The colony starts with 3 bacteria, and has 6 at the end of day 1.”** From this, it is evident that the bacteria count doubles daily. Therefore, explicitly stating the doubling rule is not strictly necessary for reasoning.\nThe issue arises when the model over-attends to such redundant information. In this example, excessive focus on the doubling phrase may lead the model to overlook truly critical conditions, such as **“ends with the colony having more than 100 bacteria.”** This misalignment in attention can cause the model to produce incorrect answers, despite having access to all necessary information. \nBy training on counterfactual examples where only the misleading tokens are masked, CAD encourages models to redistribute their attention across all genuinely informative tokens, thereby improving overall reasoning robustness.\n\n**Q4：Weak baselines: The random masking baseline is too simple, and comparisons to advanced methods like CoT knowledge distillation are missing. A direct comparison is needed to show that improvements come from addressing spurious correlations, not just better CoT distillation.**\n\n**A4:** The **KD w/o Mask baselines** reported in Table 1 of Section 3.2 (including LLaMA-1B/3B-Instruct and Qwen2.5-Math-1.5B), as well as the **Random Masking** and **PPL-based Masking** baselines discussed in Section 4.1, are essentially implementations of **Chain-of-Thought (CoT) knowledge distillation** as referenced in the related work.\n\nSpecifically, the comparisons you mention—between Random Masking and CoT knowledge distillation, as well as between our proposed CAD method and CoT knowledge distillation—are directly addressed **in Sections 4.1 and 3.2 (Table 1)** respectively. Therefore, we believe our baselines are both strong and appropriate.\n\nWe acknowledge that the terminology used to describe “KD” in the manuscript may have caused confusion. To clarify, the baselines employed correspond to CoT knowledge distillation rather than simple answer-based knowledge distillation. We will emphasize this distinction more clearly in future versions of the paper to avoid misunderstandings.\n\n**Q5: Section 2.2 describes token masking and data augmentation with misleading tokens detected individually, but Section 4.2 shows consecutive spans being masked. How are individual misleading tokens combined into spans for masking?**\n\n**A5:** The process of masking misleading tokens consists of two stages. \n\nFirst, for each instruction, we **compute the normalized gradient difference** between the student and teacher for every token **individually**. \n\nAfter obtaining these per-token scores, we set a threshold to identify which tokens are considered misleading. If multiple consecutive tokens within the same sample all exceed this threshold, they are grouped and masked together as **a contiguous span, rather than as isolated tokens**. \n\nIn this way, while the detection of misleading tokens is performed at the individual token level, the actual masking operation is applied at the span level when adjacent tokens are identified as misleading. \n\nThis approach better reflects the reality that misleading patterns can often appear as phrases or segments, preserving the semantic integrity of the data augmentation process.\n\n**Q6: Is random masking in Section 4.1 applied to consecutive token spans or independent tokens, given that misleading token masking targets spans?**\n\n**A6:** In Subsection 4.1, the random‑masking setting mirrors CAD’s span lengths: for any example where CAD masks n consecutive tokens, the baseline also masks n consecutive tokens, with the start position chosen uniformly at random.\n\n**Q7: How were misleading tokens identified in Figure 1’s preliminary experiment? Was the same gradient-based method used as in the main approach?**\n\n**A7:** The misleading tokens in Figure 1 were identified using the gradient‑based attribution technique detailed in the Methods section. This experiment use the same method as the proposed methodology's gradient-based approach.\n\n**Q8: How was the threshold in Equation 6 chosen? The appendix sensitivity analysis is hard to verify; more explanation in the main text is needed.**\n\n**A8:** Refer to **Reviewer 2JsC Q2(a) and Q2(c)**\n\n**Q9: [line 58] typo \"Causal Attention Distillation\"**\n\n**A9:** Thanks for your suggestion, we will update in the next version."}]

Author rebuttal:

Meta-review: This paper introduces Causal Attention Distillation (CAD), a two-stage framework that enhances LLMs’ reasoning by identifying and masking spurious tokens through gradient-guided teacher-student comparisons, grounded in causal theory to reduce attention bias toward misleading markers. The man contributions of this work are: 1) it addresses a critical issue of spurious correlations in LLM reasoning, 2) it combines causal principles with practical distillation techniques, 3) it shows consistent performance gains across models/tasks, and 4) it enhances interpretability via clear visualizations. The authors are advised to effectively integrate the constructive suggestions from reviewers into the final version.
Decision: Accept (poster)