MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head

TL;DR

MHLA identifies global context collapse as a central limitation of linear attention and solves it by partitioning tokens into multiple heads along the token dimension, enabling query-dependent, token-level diversity with linear-time complexity. The method introduces local KV summaries per block and a learnable coefficient matrix to mix these summaries for each query block, boosting expressive power without resorting to extra modules. Across image classification, image generation, NLP, and video generation, MHLA achieves consistent improvements over prior linear attention methods and approaches or matches self-attention performance at comparable costs. This makes MHLA a practical, hardware-friendly alternative for long-sequence modeling in vision, language, and multimodal generation tasks.

Abstract

While the Transformer architecture dominates many fields, its quadratic self-attention complexity hinders its use in large-scale applications. Linear attention offers an efficient alternative, but its direct application often degrades performance, with existing fixes typically re-introducing computational overhead through extra modules (e.g., depthwise separable convolution) that defeat the original purpose. In this work, we identify a key failure mode in these methods: global context collapse, where the model loses representational diversity. To address this, we propose Multi-Head Linear Attention (MHLA), which preserves this diversity by computing attention within divided heads along the token dimension. We prove that MHLA maintains linear complexity while recovering much of the expressive power of softmax attention, and verify its effectiveness across multiple domains, achieving a 3.6\% improvement on ImageNet classification, a 6.3\% gain on NLP, a 12.6\% improvement on image generation, and a 41\% enhancement on video generation under the same time complexity.

Figures (7)

• Figure 1: (a) Generation results from our fine-tuned SANA model using MHLA. (b) Performance and efficiency comparison between the proposed MHLA and baselines. The throughput was tested on the NVIDIA H100 Tensor Core GPU. Following the previous method, we report the FID in the table at a resolution of $256 \times 256$. (c) Multi-domain performance of MHLA. We evaluate MHLA across diverse domains, demonstrating its strong and universal performance. (d) Throughput of DiT-S/2 at 4096 resolution across different devices. All improvements are solely due to MHLA, and can be further combined with orthogonal techniques for even greater speedups.
• Figure 2: Comparison between the proposed MHLA and other linear attentions. MHLA divides multiple heads on the token dimension. Through Multi-Head Mixing, MHLA restores query-conditioned selectivity by mixing KV summaries with query-specific weight, improving token-level diversity while keeping linear complexity.
• Figure 3: (a) Visualization of attention score and attention maps of MHLA and baselines. (b) Average rank and entropy of attention scores for DeiT-T, showing MHLA yields richer and more focused attention.
• Figure 4: (a) Overview of the proposed Multi-Head Linear Attention. (b) We visualize two rows of the initialized Learnable Coefficient Matrix corresponding to $Block~1$ and $Block~14$ separately when M is 25. We reshape the two rows and the M dimension in 2D for better understanding.
• Figure 5: Loss comparison.