Paper 2601.07832
MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head
- Published
- Jan 2026
- Research lab
- Independent
- Citations
- 5
- GitHub
- 152 stars
01 In brief
Summary
The paper introduces Multi-Head Linear Attention (MHLA), a linear attention mechanism that restores expressivity by partitioning tokens into blocks (heads) along the token dimension, computing local key-value summaries, and mixing them with query-specific learned coefficients.
This addresses 'global context collapse' in standard linear attention, which suffers from rank deficiency and high entropy in attention maps.
MHLA maintains O(N) complexity and requires no extra modules.
Experiments show improvements over baselines: 3.6% accuracy gain on ImageNet classification, 6.3% on NLP, 12.6% on image generation (DiT), and 41% on video generation (Wan2.1-1.3B), with throughput comparable to linear attention and up to 2.1x speedup over self-attention.
The method is validated across DeiT, VLT, DiT, SANA, and language models, and includes ablations on initialization and head number.
02 From the paper
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.