Attention Sink Forges Native MoE in Attention Layers: Sink-Aware Training to Address Head Collapse

TL;DR

The paper addresses attention sink and head collapse in large language models by revealing that sink-based attention acts as an implicit head gating factor, effectively forming a native Mixture-of-Experts structure within attention layers. It introduces a sink-aware auxiliary load-balancing loss to promote balanced head utilization, applicable to Vanilla, Sink, and Gated Attention, and demonstrates that this approach improves expressiveness and long-context retrieval, with larger models benefiting most. The key contributions include a formal MoE interpretation of attention variants, a gating-based analysis linking Sink and Gated Attention, and a practical training strategy compatible with modern kernels like Flash Attention. The findings provide a unified framework for understanding attention mechanisms and offer a scalable, effective method to mitigate head collapse and boost performance in both pre-training and fine-tuning scenarios.

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.

Figures (11)

• Figure 1: Vanilla Attention used in most open-source models. Sink Attention used in GPT-OSS with a learnable bias sink added to the softmax denominator. Gated Attention used in Qwen3-Next with a head-wise gating factor computed via sigmoid activation.
• Figure 2: $\ell_{2}$-norm of token values across layers. The first token's value vector approaches zero in models with Vanilla Attention.
• Figure 3: Visualization of attention patterns and query-key geometry. Each panel shows the attention map (left) and PCA projection of query and key vectors (right), where red stars indicate $\mathbf{k}_0$. (a)(b)(c) Vanilla Attention models exhibit attention sink and constrained query-key geometry. (d)(e)(f) Sink Attention and Gated Attention models show no sink phenomenon and more flexible vector distributions.
• Figure 4: Head importance scores over all tokens and samples for LLaMA-3.1-8B (Vanilla Attention), GPT-OSS-20B (Sink Attention), and Qwen3-Next-80B-A3B (Gated Attention). Qwen3-Next-80B-A3B uses different numbers of heads in different layers.
• Figure 5: Head load imbalance during training across model scales and attention mechanisms. Each subplot shows the coefficient of variation (Equation \ref{['eq:head_imbalance']}) over training steps for models trained with and without the auxiliary load balancing loss. Raw data and exponential moving average (EMA) smoothed trends are displayed.