MHLA: Ripristinare l'Espressività dell'Attenzione Lineare tramite Multi-Head a Livello di Token
MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head
January 12, 2026
Autori: Kewei Zhang, Ye Huang, Yufan Deng, Jincheng Yu, Junsong Chen, Huan Ling, Enze Xie, Daquan Zhou
cs.AI
Abstract
Sebbene l'architettura Transformer domini molti campi, la sua complessità quadratica dell'auto-attenzione ne ostacola l'uso in applicazioni su larga scala. L'attenzione lineare offre un'alternativa efficiente, ma la sua applicazione diretta spesso degrada le prestazioni, e le soluzioni esistenti tipicamente re-introducono un sovraccarico computazionale attraverso moduli aggiuntivi (ad esempio, convoluzione separabile in profondità) che vanificano lo scopo originale. In questo lavoro, identifichiamo una causa principale di fallimento in questi metodi: il collasso del contesto globale, dove il modello perde diversità rappresentativa. Per affrontare questo problema, proponiamo la Multi-Head Linear Attention (MHLA), che preserva questa diversità calcolando l'attenzione all'interno di teste separate lungo la dimensione dei token. Dimostriamo che MHLA mantiene una complessità lineare recuperando gran parte del potere espressivo dell'attenzione softmax e ne verifichiamo l'efficacia in molteplici domini, ottenendo un miglioramento del 3,6% sulla classificazione ImageNet, un guadagno del 6,3% sull'NLP, un miglioramento del 12,6% sulla generazione di immagini e un incremento del 41% sulla generazione video a parità di complessità temporale.
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.