MHLA: Restauración de la Expresividad de la Atención Lineal mediante Múltiples Cabezales a Nivel 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
Resumen
Si bien la arquitectura Transformer domina numerosos campos, su complejidad cuadrática de autoatención dificulta su uso en aplicaciones a gran escala. La atención lineal ofrece una alternativa eficiente, pero su aplicación directa a menudo degrada el rendimiento, y las soluciones existentes suelen reintroducir sobrecarga computacional mediante módulos adicionales (por ejemplo, convolución separable en profundidad) que frustran el propósito original. En este trabajo, identificamos un modo de fallo clave en estos métodos: el colapso del contexto global, donde el modelo pierde diversidad representacional. Para abordarlo, proponemos la Atención Lineal Multi-Cabezal (MHLA), que preserva esta diversidad calculando la atención dentro de cabezales divididos a lo largo de la dimensión de tokens. Demostramos que MHLA mantiene una complejidad lineal mientras recupera gran parte del poder expresivo de la atención softmax, y verificamos su eficacia en múltiples dominios, logrando una mejora del 3.6% en clasificación de ImageNet, una ganancia del 6.3% en PLN, una mejora del 12.6% en generación de imágenes y una mejora del 41% en generación de vídeo bajo la misma complejidad 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.