Title: Learning to Focus: Causal Attention Distillation via Gradient‐Guided Token Pruning
Venue: NeurIPS 2025
Accepted: yes
Keywords: Knowledge Distillation

Spurious Correlations

Long-Context Reasoning

Model Interpretability

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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: 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.\n\nQUESTIONS:\n1.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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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).\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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.\n\nQUESTIONS:\nQ1. 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-07"}