MHLA : Restauration de l'expressivité de l'attention linéaire via une approche multi-têtes au niveau des tokens

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

January 12, 2026

papers.authors: Kewei Zhang, Ye Huang, Yufan Deng, Jincheng Yu, Junsong Chen, Huan Ling, Enze Xie, Daquan Zhou

cs.AI

papers.abstract

Bien que l'architecture Transformer domine de nombreux domaines, sa complexité attentionnelle quadratique entrave son utilisation dans les applications à grande échelle. L'attention linéaire offre une alternative efficace, mais son application directe dégrade souvent les performances, les correctifs existants réintroduisant généralement une surcharge computationnelle via des modules supplémentaires (par exemple, la convolution séparable en profondeur) qui contredisent l'objectif initial. Dans ce travail, nous identifions un mode d'échec clé de ces méthodes : l'effondrement du contexte global, où le modèle perd sa diversité représentationnelle. Pour y remédier, nous proposons l'Attention Linéaire Multi-Têtes (MHLA), qui préserve cette diversité en calculant l'attention au sein de têtes divisées le long de la dimension des tokens. Nous démontrons que MHLA maintient une complexité linéaire tout en retrouvant une grande partie de la puissance expressive de l'attention softmax, et vérifions son efficacité dans plusieurs domaines, obtenant une amélioration de 3,6 % sur la classification ImageNet, un gain de 6,3 % en TAL, une amélioration de 12,6 % sur la génération d'images et une augmentation de 41 % sur la génération vidéo à complexité temporelle égale.

English

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.