Learning to Attribute with Attention

• URL: http://arxiv.org/abs/2504.13752v1
• Date: Fri, 18 Apr 2025 15:36:28 GMT
• Title: Learning to Attribute with Attention
• Authors: Benjamin Cohen-Wang, Yung-Sung Chuang, Aleksander Madry,
• Abstract summary: We propose treating attention weights of different attention heads as features.<n>This way, we can learn how to effectively leverage attention weights for attribution.<n>Our resulting method, Attribution with Attention (AT2), reliably performs on par with approaches that involve many ablations.
• Score: 75.61481181755744
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Given a sequence of tokens generated by a language model, we may want to identify the preceding tokens that influence the model to generate this sequence. Performing such token attribution is expensive; a common approach is to ablate preceding tokens and directly measure their effects. To reduce the cost of token attribution, we revisit attention weights as a heuristic for how a language model uses previous tokens. Naive approaches to attribute model behavior with attention (e.g., averaging attention weights across attention heads to estimate a token's influence) have been found to be unreliable. To attain faithful attributions, we propose treating the attention weights of different attention heads as features. This way, we can learn how to effectively leverage attention weights for attribution (using signal from ablations). Our resulting method, Attribution with Attention (AT2), reliably performs on par with approaches that involve many ablations, while being significantly more efficient. To showcase the utility of AT2, we use it to prune less important parts of a provided context in a question answering setting, improving answer quality. We provide code for AT2 at https://github.com/MadryLab/AT2 .

Related papers

• Stop Looking for Important Tokens in Multimodal Language Models: Duplication Matters More [18.928285521147057]
  We show that importance is not an ideal indicator to decide whether a token should be pruned.<n>We propose DART (Duplication-Aware Reduction of Tokens), which prunes tokens based on its duplication with other tokens.<n>Experiments demonstrate that DART can prune 88.9% vision tokens while maintaining comparable performance.
  arXiv  Detail & Related papers  (2025-02-17T06:56:28Z)
• AttentionPredictor: Temporal Pattern Matters for Efficient LLM Inference [51.1972443343829]
  We propose AttentionPredictor, which is the first learning-based critical token identification approach.<n> AttentionPredictor accurately predicts the attention score while consuming negligible memory.<n>We also propose a cross-token critical cache prefetching framework that hides the token time overhead to accelerate the decoding stage.
  arXiv  Detail & Related papers  (2025-02-06T13:41:46Z)
• More Expressive Attention with Negative Weights [36.40344438470477]
  We propose a novel attention mechanism, named Cog Attention, that enables attention weights to be negative for enhanced expressiveness.<n>Our approach suggests a promising research direction for rethinking and breaking the entrenched constraints of traditional softmax attention.
  arXiv  Detail & Related papers  (2024-11-11T17:56:28Z)
• When Attention Sink Emerges in Language Models: An Empirical View [39.36282162213973]
  Language Models (LMs) assign significant attention to the first token, even if it is not semantically important. This phenomenon has been widely adopted in applications such as streaming/long context generation, KV cache optimization, inference acceleration, model quantization, and others. We first demonstrate that attention sinks exist universally in LMs with various inputs, even in small models.
  arXiv  Detail & Related papers  (2024-10-14T17:50:28Z)
• Elliptical Attention [1.7597562616011944]
  Pairwise dot-product self-attention is key to the success of transformers that achieve state-of-the-art performance across a variety of applications in language and vision. We propose using a Mahalanobis distance metric for computing the attention weights to stretch the underlying feature space in directions of high contextual relevance.
  arXiv  Detail & Related papers  (2024-06-19T18:38:11Z)
• ToSA: Token Selective Attention for Efficient Vision Transformers [50.13756218204456]
  ToSA is a token selective attention approach that can identify tokens that need to be attended as well as those that can skip a transformer layer. We show that ToSA can significantly reduce computation costs while maintaining accuracy on the ImageNet classification benchmark.
  arXiv  Detail & Related papers  (2024-06-13T05:17:21Z)
• Simple linear attention language models balance the recall-throughput tradeoff [60.06020449520365]
  We propose BASED, a simple architecture combining linear and sliding window attention. We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points.
  arXiv  Detail & Related papers  (2024-02-28T19:28:27Z)
• How can objects help action recognition? [74.29564964727813]
  We investigate how we can use knowledge of objects to design better video models. First, we propose an object-guided token sampling strategy that enables us to retain a small fraction of the input tokens. Second, we propose an object-aware attention module that enriches our feature representation with object information.
  arXiv  Detail & Related papers  (2023-06-20T17:56:16Z)
• Robustifying Token Attention for Vision Transformers [72.07710236246285]
  Vision transformers (ViTs) still suffer from significant drops in accuracy in the presence of common corruptions. We propose two techniques to make attention more stable through two general techniques. First, our Token-aware Average Pooling (TAP) module encourages the local neighborhood of each token to take part in the attention mechanism. Second, we force the output tokens to aggregate information from a diverse set of input tokens rather than focusing on just a few.
  arXiv  Detail & Related papers  (2023-03-20T14:04:40Z)

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.