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

• URL: http://arxiv.org/abs/2506.07851v1
• Date: Mon, 09 Jun 2025 15:16:39 GMT
• Title: Learning to Focus: Causal Attention Distillation via Gradient-Guided Token Pruning
• Authors: Yiju Guo, Wenkai Yang, Zexu Sun, Ning Ding, Zhiyuan Liu, Yankai Lin,
• Abstract summary: Large language models (LLMs) have demonstrated significant improvements in contextual understanding.<n>However, their ability to attend to truly critical information during long-context reasoning and generation still falls behind the pace.<n>We introduce a two-stage framework called Learning to Focus (LeaF) to mitigate confounding factors.
• Score: 47.764552063499046
• License: http://creativecommons.org/licenses/by/4.0/
• 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 and code generation benchmarks but also effectively suppresses attention to confounding tokens during inference, yielding a more interpretable and reliable reasoning model.

Related papers

• A Closer Look at Bias and Chain-of-Thought Faithfulness of Large (Vision) Language Models [53.18562650350898]
  Chain-of-thought (CoT) reasoning enhances performance of large language models.<n>We present the first comprehensive study of CoT faithfulness in large vision-language models.
  arXiv  Detail & Related papers  (2025-05-29T18:55:05Z)
• When Thinking Fails: The Pitfalls of Reasoning for Instruction-Following in LLMs [16.659986373052217]
  Chain-of-thought reasoning can significantly degrade instruction-following accuracy.<n>This is the first work to systematically expose reasoning-induced failures in instruction-following.
  arXiv  Detail & Related papers  (2025-05-16T16:36:00Z)
• Don't Take Things Out of Context: Attention Intervention for Enhancing Chain-of-Thought Reasoning in Large Language Models [32.71672086718058]
  Few-shot Chain-of-Thought (CoT) significantly enhances the reasoning capabilities of large language models (LLMs)<n>We observe that isolated segments, words, or tokens within CoT demonstrations can unexpectedly disrupt the generation process of LLMs.<n>We propose a Few-shot Attention Intervention method (FAI) that dynamically analyzes the attention patterns of demonstrations to accurately identify these tokens.
  arXiv  Detail & Related papers  (2025-03-14T07:46:33Z)
• Aggregation Artifacts in Subjective Tasks Collapse Large Language Models' Posteriors [74.04775677110179]
  In-context Learning (ICL) has become the primary method for performing natural language tasks with Large Language Models (LLMs)<n>In this work, we examine whether this is the result of the aggregation used in corresponding datasets, where trying to combine low-agreement, disparate annotations might lead to annotation artifacts that create detrimental noise in the prompt.<n>Our results indicate that aggregation is a confounding factor in the modeling of subjective tasks, and advocate focusing on modeling individuals instead.
  arXiv  Detail & Related papers  (2024-10-17T17:16:00Z)
• Towards Context-Aware Emotion Recognition Debiasing from a Causal Demystification Perspective via De-confounded Training [14.450673163785094]
  Context-Aware Emotion Recognition (CAER) provides valuable semantic cues for recognizing the emotions of target persons. Current approaches invariably focus on designing sophisticated structures to extract perceptually critical representations from contexts. We present a Contextual Causal Intervention Module (CCIM) to de-confound the confounder.
  arXiv  Detail & Related papers  (2024-07-06T05:29:02Z)
• Evaluating Human Alignment and Model Faithfulness of LLM Rationale [66.75309523854476]
  We study how well large language models (LLMs) explain their generations through rationales. We show that prompting-based methods are less "faithful" than attribution-based explanations.
  arXiv  Detail & Related papers  (2024-06-28T20:06:30Z)
• Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers [49.80959223722325]
  We study the distinction between feed-forward and attention layers in large language models.<n>We find that feed-forward layers tend to learn simple distributional associations such as bigrams, while attention layers focus on in-context reasoning.
  arXiv  Detail & Related papers  (2024-06-05T08:51:08Z)
• Causal Concept Graph Models: Beyond Causal Opacity in Deep Learning [11.13665894783481]
  Causal opacity denotes the difficulty in understanding the "hidden" causal structure underlying the decisions of deep neural network (DNN) models.<n>This work introduces Causal Concept Graph Models (Causal CGMs), a class of interpretable models whose decision-making process is causally transparent by design.<n>Our experiments show that Causal CGMs can: (i) match the generalisation performance of causally opaque models, (ii) enable human-in-the-loop corrections to mispredicted intermediate reasoning steps, and (iii) support the analysis of interventional and counterfactual scenarios.
  arXiv  Detail & Related papers  (2024-05-26T10:15:20Z)
• Identifying Semantic Induction Heads to Understand In-Context Learning [103.00463655766066]
  We investigate whether attention heads encode two types of relationships between tokens present in natural languages. We find that certain attention heads exhibit a pattern where, when attending to head tokens, they recall tail tokens and increase the output logits of those tail tokens.
  arXiv  Detail & Related papers  (2024-02-20T14:43:39Z)
• Unveiling the Magic: Investigating Attention Distillation in Retrieval-augmented Generation [8.363702038073814]
  Retrieval-augmented generation framework can address the limitations of large language models by enabling real-time knowledge updates for more accurate answers. An efficient way in the training phase of retrieval-augmented models is attention distillation, which uses attention scores as a supervision signal instead of manually annotated query-document pairs.
  arXiv  Detail & Related papers  (2024-02-19T02:48:44Z)
• Using Early Readouts to Mediate Featural Bias in Distillation [30.5299408494168]
  Deep networks tend to learn spurious feature-label correlations in real-world supervised learning tasks. We propose a novel early readout mechanism whereby we attempt to predict the label using representations from earlier network layers.
  arXiv  Detail & Related papers  (2023-10-28T04:58:15Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.