Attention Needs to Focus: A Unified Perspective on Attention Allocation

TL;DR

The paper addresses fundamental inefficiencies in Transformer self-attention, namely representational collapse from Attention Overload and spurious focus from Attention Underload (attention sink). It introduces Lazy Attention, which combines Positional Discrimination (RoPE-based dimension-wise rotation plus learnable head-wise distance biases) with Elastic-Softmax (a learnable head-specific offset and ReLU-based filtering) to produce targeted, sparse attention. Across nine benchmarks on FineWeb-Edu scales, Lazy Attention reduces attention sink, achieves up to 59.58% sparsity, and maintains competitive language modeling performance, while improving long-context robustness via length extrapolation. These findings suggest a practical path toward more efficient, focused attention in large-scale language models and open avenues for hardware- and computation-efficient LLM deployment.

Abstract

The Transformer architecture, a cornerstone of modern Large Language Models (LLMs), has achieved extraordinary success in sequence modeling, primarily due to its attention mechanism. However, despite its power, the standard attention mechanism is plagued by well-documented issues: representational collapse and attention sink. Although prior work has proposed approaches for these issues, they are often studied in isolation, obscuring their deeper connection. In this paper, we present a unified perspective, arguing that both can be traced to a common root -- improper attention allocation. We identify two failure modes: 1) Attention Overload, where tokens receive comparable high weights, blurring semantic features that lead to representational collapse; 2) Attention Underload, where no token is semantically relevant, yet attention is still forced to distribute, resulting in spurious focus such as attention sink. Building on this insight, we introduce Lazy Attention, a novel mechanism designed for a more focused attention distribution. To mitigate overload, it employs positional discrimination across both heads and dimensions to sharpen token distinctions. To counteract underload, it incorporates Elastic-Softmax, a modified normalization function that relaxes the standard softmax constraint to suppress attention on irrelevant tokens. Experiments on the FineWeb-Edu corpus, evaluated across nine diverse benchmarks, demonstrate that Lazy Attention successfully mitigates attention sink and achieves competitive performance compared to both standard attention and modern architectures, while reaching up to 59.58% attention sparsity.

Figures (10)

• Figure 1: Overview of two attention failure modes. Attention Overload: broadly assigned high weights blur important patterns. Attention Underload: uniformly low relevance gets normalized.
• Figure 2: Analysis of the location of the attention sink and the characteristics of the sink token.
• Figure 3: Comparison of attention weight patterns in five LLMs employing different position encodings. The models are evaluated with two types of input: (a) natural text and (b) repeated tokens. Under repeated tokens, the sink pattern of attention weights changes significantly.
• Figure 4: Learned position-dependent attention biases across different layers and attention heads.
• Figure 5: Distribution of learned Elastic-Softmax offsets across layers, with violin plots showing variation across heads and white dots denoting per-layer means.