Learning to Focus: Causal Attention Distillation via Gradient‐Guided Token Pruning

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.

A1: We address the concern by explicitly isolating the impact of misleading tokens from other confounding factors: First, 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.

Second, 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). CAD consistently delivers additional improvements over CoT-KD on both Math and Code tasks, highlighting that our approach provides value beyond these established techniques.

Finally, 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.

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.

A2: We appreciate the opportunity to clarify our method of defining misleading tokens, as there seems to be a misunderstanding. We 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.

More 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.

Crucially, 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.

Q3(a)：Superficial application of causal framework: While citing Pearl's causal framework, it essentially amounts to simple token masking.

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.

Our 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. Thus, CAD is not merely token masking—it is a principled causal intervention designed explicitly to eliminate spurious correlations while preserving essential information.

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.

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. In 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. The 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. By 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.

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.

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.

Specifically, 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.

We 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.

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?

A5: The process of masking misleading tokens consists of two stages.

First, for each instruction, we compute the normalized gradient difference between the student and teacher for every token individually.

After 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.

In 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.

This approach better reflects the reality that misleading patterns can often appear as phrases or segments, preserving the semantic integrity of the data augmentation process.

Q6: Is random masking in Section 4.1 applied to consecutive token spans or independent tokens, given that misleading token masking targets spans?

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.

Q7: How were misleading tokens identified in Figure 1’s preliminary experiment? Was the same gradient-based method used as in the main approach?

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.

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.

A8: Refer to Reviewer 2JsC Q2(a) and Q2(c)

Q9: [line 58] typo "Causal Attention Distillation"

A9: Thanks for your suggestion, we will update in the next version.

Comment (A1): I appreciate the clarity in A1 regarding the isolation of CAD's benefits from confounding factors such as sequence length or standard CoT distillation. The ablation studies help clarify CAD’s unique contributions.

Response 1: Thank you for your feedback. We're glad the ablation studies in A1 clearly demonstrate CAD’s unique contributions, independent of factors like sequence length and standard CoT distillation. We will retain and highlight these results in the revision.

Comment (A2): The reviewer appreciates the explanation of misleading tokens but requests a clearer description of the data selection process in the main paper and clarification on whether the baselines use the same filtered subset.

Response 2: Thank you for your valuable feedback. We will incorporate a more detailed description of the misleading token selection mechanism into the main paper to improve clarity and transparency, as this component is central to our approach.

Regarding the data used in Table 1, the training configurations are as follows:

• Baseline (CoT KD): 12K Original Training Samples
• CAD: 12K Original Training Samples + Counterfactual Samples (generated by pruning misleading spans)

Importantly, the 12K original samples used in both CoT KD and CAD are unfiltered and include all cases where the student model succeeds or fails.

This setup ensures a fair comparison between CAD and the baseline, as both methods are built upon the same foundational dataset.

Comment (Q8): The reviewer notes that CAD's performance is somewhat sensitive to the choice of the threshold τ, with certain configurations (e.g., τ = 0.05 or 0.15 on MATH) showing performance drops relative to KD w/o Mask. While this effect is limited in most cases, it is more pronounced on SQuAD 2.0. The reviewer suggests including a brief explanation in the main paper of how τ is selected or tuned, even heuristically, to improve clarity and reproducibility.

Response 3: Thank you for the thoughtful suggestion. We agree that clarifying the selection of the threshold τ would improve the clarity and reproducibility of our method.

Based on our threshold sensitivity analysis on MATH, Code, and SQuAD 2.0, we find that τ = 0.15 works well across tasks for LLaMA3.2-1B-Instruct, while τ = 0.10 yields better overall results for LLaMA3.2-3B-Instruct. We also observe that, regardless of whether the supervision is applied at the instruction or response level, lower-capacity models are more susceptible to noise from spurious tokens and thus benefit from a stricter threshold. This insight guides our heuristic selection of τ in practice.

To improve clarity, we will add a dedicated subsection in Section 4 of the main paper detailing the selection strategy for τ, along with a cross-dataset sensitivity analysis that illustrates its impact across different model sizes and tasks.

Comment on Causal Framing: The reviewer raises concerns that natural language involves more complex token-level dependencies, such as mediator structures, where masking a token like A could disrupt valid reasoning paths. Further empirical validation is suggested to support the claim that CAD performs meaningful causal interventions on input tokens.

Response 4: Thank you for the thoughtful and insightful feedback. We fully agree that token-level dependencies in natural language can involve more complex structures, including mediators such as X₁ → A → X₂ → Y. We address your comment from two perspectives:

First, regarding the concern that masking a token like A might inadvertently disrupt a valid reasoning path: in our method, misleading tokens are identified based on cases where the student model initially produces an incorrect answer but becomes correct after the token is masked. This selection criterion ensures that the masked token is not part of a valid reasoning chain—on the contrary, its removal leads to improved model performance. Thus, the masking operation functions as a targeted intervention that eliminates spurious influences without interfering with necessary dependencies.

Second, we appreciate your observation that our current causal framing may implicitly assume a simplified A → Y structure. Inspired by your comment, we will revise our description to better reflect more general causal patterns, including mediator chains such as X₁ → A → X₂ → Y. In these cases, the masked token A may serve as an incorrect condition or irrelevant premise that propagates downstream reasoning errors (e.g., X₂). Removing such tokens can help disrupt not only direct spurious links (as in A → Y), but also more complex erroneous intermediate dependencies (e.g., X₁ → A → X₂ → Y), thereby aligning with the broader goals of causal intervention.

We will clarify these points in the revised paper to improve the conceptual rigor of our causal framing, while keeping the core method empirically grounded.

Dear Reviewer WSL6,

We have diligently addressed all the concerns raised during the rebuttal period, and we would like to kindly ask if any remaining issues have been resolved. As the rebuttal deadline is fast approaching, we would greatly appreciate any further feedback, suggestions, or concerns at your earliest convenience. We are eager to engage in further discussion and clarify any unresolved points to ensure all matters are fully addressed.

Best regards,

Authors of Submission 9306

I thank the authors for their detailed and constructive engagement throughout the rebuttal period. Many of my earlier concerns have been addressed, and I appreciate the additional statistics and clarifications.

Regarding Q2, I agree with the authors’ observation that counterfactual samples tend to be longer; this is an interesting property of the data, and I find their explanation plausible.

However, for Q1 and the associated experiments, I still have several reservations:

1. On the claim of a contrastive learning effect from using both original and counterfactual data I find it difficult to agree with this claim. Since both the original and counterfactual data are trained with a CE loss to maximize log-likelihood, the original data will still encourage the model to generate misleading tokens. The mere use of original and counterfactual data together does not necessarily create a contrastive signal in the standard CE setup, and this connection remains unclear to me.
2. On the claim that the model can learn to ignore misleading patterns when both datasets are used I partially agree. The results suggest that retaining misleading patterns can outperform fully removing them, but this is somewhat counterintuitive—if misleading tokens induce spurious correlations, why should preserving them be beneficial? This point may be outside the immediate scope of the paper, but it warrants further discussion or at least acknowledgement.

From a broader perspective, I personally find the method closer to reasoning distillation with pruning than to a strict causal framework, though the authors’ interpretation is not without merit. The planned revision to better align the causal framing with complex token-level dependencies is welcome.

I also encourage the authors to revise lines 125–126 to explicitly state the composition and intended role of the training data, as this will help readers fully understand the design choices. Including results from the ongoing ablation study that removes original samples would also be valuable if available before the camera-ready deadline.

Overall, I appreciate the diverse experiments and analyses, which have addressed most of my concerns. I will raise my score accordingly. While I personally view the approach as closer to reasoning distillation with pruning than to a pure causal framework, the authors’ claim is also persuasive. I hope the discussed data statistics, the clarification of the training data composition and intention (particularly lines 125–126), and the additional ablation results will be well reflected in the revision.

Further Clarification on Q1

We appreciate your thoughtful comments and would like to clarify two key points regarding potential misunderstandings about misleading patterns.

First, in our setting—whether at the instruction or response level—misleading patterns exist only in the input, not in the output.

Therefore, the concern that “the original data will still encourage the model to generate misleading tokens” does not apply, as no misleading tokens appear in the ground-truth outputs.

Likewise, the question “if misleading tokens induce spurious correlations, why should preserving them be beneficial?” assumes that CAD preserves spurious correlations. In fact, CAD is designed to reduce them.

By jointly training on original samples and their counterfactual counterparts, the student model learns to downweight attention to misleading patterns—a conclusion supported not only by CAD’s consistently higher performance than CoT KD, but also by the qualitative evidence in Section 4.3 (Case Study). Without this dual exposure, spurious correlations in the student model would remain unmitigated.

Second, misleading patterns do not cause spurious correlations in the teacher model’s reasoning.

These patterns are features that the student over-attends to, but that the teacher ignores, and the teacher predicts correctly regardless of their presence (by definition, all original samples are correctly solved by the teacher). Since the student is distilled from the teacher via logits alignment, the training signal itself does not introduce spurious correlations.

Finally, while the CE loss itself is unchanged, the contrastive learning effect arises from the paired input construction with identical labels: each original–counterfactual pair presents the same task—with and without misleading patterns—forcing the student to adjust its attention distribution. Attending to misleading patterns harms performance on counterfactuals, while ignoring them benefits both, creating a functional contrastive signal in the input space.

In short, misleading patterns do not cause spurious correlations in the teacher model’s reasoning. Instead, the dual exposure to original and counterfactual inputs enables the student to recognize and suppress such correlations in its own reasoning process.

We sincerely thank you for the positive feedback and for raising the score.

We also appreciate your recognition of our efforts in conducting diverse experiments and analyses, as well as your constructive suggestions.

We will ensure that the clarification of the training data composition and intention (particularly lines 125–126), the discussed data statistics, and the additional ablation results are clearly and thoroughly incorporated in the revised version of the paper.

Best regards,

Authors of Submission 9306

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.

A1: We acknowledge that our method was primarily validated on the LLaMA and Qwen model families.

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.

Furthermore, 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.

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.

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:

Our 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:

$P(Y | X_i = x) = \sum_{a} P(Y \mid X_i = x, A = a) \, P(A = a \mid X_i = x)$

which 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$.

We 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:

$P(Y \mid \mathrm{do}(X_i = x)) = P(Y \mid X_i = x, A=\emptyset)$,

which restores the genuine causal effect of $X_i$ on $Y$.

In 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.

This 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.

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?

We quantified the additional computational overhead of CAD compared to standard KD by measuring end-to-end runtime using identical hardware:

Gradient computation (one-time, 8× NVIDIA A100 80GB): Computing 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).

Counterfactual generation (one-time, offline via vLLM, 4×NVIDIA A100 80GB): Generating 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.

Training overhead (3 epochs 4×NVIDIA A100 80GB):

Overall, 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.

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.

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.

Our 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.

We 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.

Although 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.

Q5: How to ensure the stability of the gradient difference threshold (τ_mis) in cross-task scenarios?

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:

We 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.

Dear Reviewer ZMF9:

Thank you again for the great efforts and valuable comments. We have carefully addressed the main concerns in detail. We hope you might find the response satisfactory (similar as the other reviewers). As the discussion phase is about to close, we are very much looking forward to hearing from you about any further feedback. We will be very happy to clarify any further concerns (if any).

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.

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.

Furthermore, 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.

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?

A2: We quantified the additional computational overhead of CAD compared to standard KD by measuring end-to-end runtime using identical hardware:

Gradient computation (one-time, 8× NVIDIA A100 80GB): Computing 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).

Counterfactual generation (one-time, offline via vLLM, 4×NVIDIA A100 80GB): Generating 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.

Training overhead (3 epochs 4×NVIDIA A100 80GB):

Overall, 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.

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.

A3: We believe that the observed 2–3% absolute improvement over standard KD is substantial given the strong performance of the CoT KD baselines. Specifically, 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.

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.

Moreover, 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.

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).

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. To 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.

Results (EM):

Llama‑1b‑Instruct: 41.11% → KD: 72.56% → CAD: 74.78%

Llama‑3b‑Instruct: 73.67% → KD: 83.67% → CAD: 89.44%

These 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.

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?

A5: To evaluate the robustness of our method, we conducted analyses from both the training and testing perspectives.

First, 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.

Second, 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:

These 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.

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?

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.

Thank you again for the great efforts and valuable comments. We have carefully addressed the main concerns in detail. We hope you might find the response satisfactory. As the discussion phase is about to close, we are very much looking forward to hearing from you about any further feedback. We will be very happy to clarify any further concerns (if any).

Dear Reviewer,

As the Area Chair handling this submission, I would like to express my appreciation for the thorough and insightful comments you have provided. The authors have now submitted a comprehensive rebuttal addressing each of the points raised in your reviews. To facilitate a fair and rigorous decision-making process, I kindly urge each reviewer to carefully assess whether the authors' rebuttal adequately addresses your specific concerns and technical questions. Your active participation and constructive input during the upcoming discussion phase will be crucial in arriving at a well-informed decision.

Please feel free to share any additional thoughts or clarifications you may have as we proceed with our deliberations.

Best regards,

Your AC

Dear Reviewer qdQ9,

We have diligently addressed all concerns raised during the rebuttal period, particularly by validating the effectiveness of the CAD method in the QA domain (A4), conducting extensive robustness evaluations of CAD (A5), and providing a detailed time-efficiency comparison of CAD versus baseline approaches (A2). As it has now been nearly six days since our rebuttal, we kindly ask you to review the updated manuscript and let us know if any issues remain. We greatly appreciate any further comments or suggestions at your earliest convenience and are happy to clarify any outstanding points.

Best regards,

The authors of Submission 9306

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?

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: First, 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.

The 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.

Q2(a)：How was the threshold parameter τ selected?

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."

In this analysis, we considered both instruction-level and response-level performance. Our findings are as follows:

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.

2.Response-level: LLaMa3.2-CAD-1B performs best at 0.15, and LLaMa3.2-CAD-3B at 0.10.

For 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.

Q2(b)：Were negative values of τ considered?

A2(b): Since the threshold parameter τ represents the percentage of misleading tokens relative to instruction tokens, there is no need to consider negative values.

Q2(c)：How sensitive is the model's performance to different values of τ across datasets?

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:

We 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.

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?

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.

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?

A4: We attribute these anomalies primarily to the limited coding capability of LLaMa3.2‑1B‑Instruct.
To 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.

In 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.

Thank you again for the great efforts and valuable comments. We have carefully addressed the main concerns in detail. We hope you might find the response satisfactory. As the discussion phase is about to close, we are very much looking forward to hearing from you about any further feedback. We will be very happy to clarify any further concerns (if any).

Dear Reviewer,

As the Area Chair handling this submission, I would like to express my appreciation for the thorough and insightful comments you have provided. The authors have now submitted a comprehensive rebuttal addressing each of the points raised in your reviews. To facilitate a fair and rigorous decision-making process, I kindly urge each reviewer to carefully assess whether the authors' rebuttal adequately addresses your specific concerns and technical questions. Your active participation and constructive input during the upcoming discussion phase will be crucial in arriving at a well-informed decision.

Please feel free to share any additional thoughts or clarifications you may have as we proceed with our deliberations.

Best regards,

Your AC

Dear Reviewer 2JsC,

We have diligently addressed every concern raised, especially our spurious-confounder analysis on larger-scale models and our detailed threshold analyses on the Math, Code, and SQuAD 2.0 benchmarks. As we have been awaiting your feedback for nearly six days, we kindly request your review of the rebuttal, we would greatly appreciate any further feedback, suggestions, or concerns at your earliest convenience. We are eager to engage in further discussion and clarify any unresolved points to ensure all matters are fully addressed.

Best regards,

Authors of Submission 9306

Dear Reviewer 2JsC,

We sincerely thank you for your earlier comments and the opportunity to address them. During the discussion phase, we have carefully and thoroughly responded to all points raised, particularly our spurious-confounder analysis on larger-scale models and the detailed threshold analyses for the Math, Code, and SQuAD 2.0 benchmarks.

As the discussion deadline is now less than 12 hours away, we would greatly value your feedback on our responses. If you feel that your concerns have been addressed, we would be grateful if this could be reflected in your updated evaluation. If there are still outstanding questions, we are more than happy to provide further clarification within the remaining time to ensure a constructive exchange before the rebuttal period concludes.

Best regards,

Authors of Submission 9306