MHLA: Restaurando a Expressividade da Atenção Linear via Múltiplas Cabeças em Nível de Token
MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head
January 12, 2026
Autores: Kewei Zhang, Ye Huang, Yufan Deng, Jincheng Yu, Junsong Chen, Huan Ling, Enze Xie, Daquan Zhou
cs.AI
Resumo
Embora a arquitetura Transformer domine muitos campos, sua complexidade quadrática de autoatenção dificulta seu uso em aplicações em larga escala. A atenção linear oferece uma alternativa eficiente, mas sua aplicação direta frequentemente degrada o desempenho, com as correções existentes normalmente reintroduzindo sobrecarga computacional por meio de módulos extras (por exemplo, convolução separável em profundidade) que frustram o objetivo original. Neste trabalho, identificamos um modo de falha fundamental nesses métodos: o colapso do contexto global, onde o modelo perde diversidade representacional. Para resolver isso, propomos a Atenção Linear Multi-Cabeça (MHLA), que preserva essa diversidade ao calcular a atenção dentro de cabeças divididas ao longo da dimensão do *token*. Demonstramos que a MHLA mantém complexidade linear enquanto recupera grande parte do poder expressivo da atenção *softmax*, e comprovamos sua eficácia em múltiplos domínios, alcançando uma melhoria de 3,6% na classificação do ImageNet, um ganho de 6,3% em PLN, uma melhoria de 12,6% na geração de imagens e um aprimoramento de 41% na geração de vídeo sob a mesma complexidade temporal.
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.