Attention Sink Forges Native MoE in Attention Layers: Sink-Aware Training to Address Head Collapse
This work addresses head collapse in attention mechanisms for LLMs, offering a new perspective on their inherent MoE structure, though it is incremental as it builds on existing methods like Sink Attention and Gated Attention.
The paper tackles the attention sink problem in Large Language Models, where disproportionate attention to the first token causes head collapse, and proposes a sink-aware training algorithm that improves model performance by achieving effective head load balancing.
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.