Attention Sink Forges Native MoE in Attention Layers: Sink-Aware Training to Address Head Collapse
Abstract
Large Language Models (LLMs) often assign disproportionate attention to the first token, a phenomenon known as the attention sink. Several recent approaches aim to address this issue, including Sink Attention in GPT-OSS and Gated Attention in Qwen3-Next. However, a comprehensive analysis of the relationship among these attention mechanisms is lacking. In this work, we provide both theoretical and empirical evidence demonstrating that the sink in Vanilla Attention and Sink Attention naturally construct a Mixture-of-Experts (MoE) mechanism within attention layers. This insight explains the head collapse phenomenon observed in prior work, where only a fixed subset of attention heads contributes to generation. To mitigate head collapse, we propose a sink-aware training algorithm with an auxiliary load balancing loss designed for attention layers. Extensive experiments show that our method achieves effective head load balancing and improves model performance across Vanilla Attention, Sink Attention, and Gated Attention. We hope this study offers a new perspective on attention mechanisms and encourages further exploration of the inherent MoE structure within attention layers.
Lay Summary
Large language models use many attention heads to choose which earlier words should guide the next prediction. Yet these models often place unusually high attention on a special sink, such as the first token, even when it carries little meaning. This paper shows that the sink is more than a side effect: it acts like a hidden switch that controls how much each attention head contributes. In this view, attention heads behave like a team of specialists, but current models often rely on only a few of them, a problem the paper calls head collapse. To fix this, we introduce sink-aware training, which adds a small balancing loss that encourages different heads to share the work more evenly. The method works with Vanilla attention, Sink Attention, and Gated Attention, and it can be used during both pre-training and fine-tuning. Experiments on newly trained models and larger pretrained models show better head utilization and consistent gains on language, reasoning, coding, and long-context tasks. Overall, the work offers a clearer way to understand attention and a practical path to make language models more capable without changing their overall architecture.