Saved in:
| Main Authors: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2601.07832 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866914254405238784 |
|---|---|
| author | Zhang, Kewei Huang, Ye Deng, Yufan Yu, Jincheng Chen, Junsong Ling, Huan Xie, Enze Zhou, Daquan |
| author_facet | Zhang, Kewei Huang, Ye Deng, Yufan Yu, Jincheng Chen, Junsong Ling, Huan Xie, Enze Zhou, Daquan |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_07832 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head Zhang, Kewei Huang, Ye Deng, Yufan Yu, Jincheng Chen, Junsong Ling, Huan Xie, Enze Zhou, Daquan Computer Vision and Pattern Recognition Artificial Intelligence 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. |
| title | MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head |
| topic | Computer Vision and Pattern Recognition Artificial Intelligence |
| url | https://arxiv.org/abs/2601.07832 |