π€ AI Summary
This work addresses the issues of representation collapse and attention sink arising from suboptimal allocation in standard attention mechanisms. It presents a unified explanation for both phenomena and introduces Lazy Attention, a novel mechanism that integrates positional discrimination, Elastic-Softmax normalization, and multi-head multidimensional optimization to produce sparse yet semantically focused attention distributions. Evaluated on FineWeb-Edu, the proposed model achieves an attention sparsity of 59.58% while matching the performance of leading architectures across nine benchmark tasks, effectively mitigating the attention sink problem.
π 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.