MHLA: Herstel van de Expressiviteit van Lineaire Aandacht via Token-Level Multi-Head
MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head
January 12, 2026
Auteurs: Kewei Zhang, Ye Huang, Yufan Deng, Jincheng Yu, Junsong Chen, Huan Ling, Enze Xie, Daquan Zhou
cs.AI
Samenvatting
Hoewel de Transformer-architectuur veel domeinen domineert, vormt de kwadratische complexiteit van de zelf-attentie een belemmering voor grootschalige toepassingen. Lineaire aandacht biedt een efficiënt alternatief, maar de directe toepassing ervan leidt vaak tot prestatieverlies. Bestaande oplossingen voeren doorgaans opnieuw rekenkosten in via extra modules (zoals depthwise separable convolutie), wat het oorspronkelijke doel tenietdoet. In dit werk identificeren we een cruciaal faalmechanisme in deze methoden: het ineenstorten van de globale context, waarbij het model representatiediversiteit verliest. Om dit aan te pakken, stellen we Multi-Head Lineaire Aandacht (MHLA) voor, die deze diversiteit behoudt door aandacht te berekenen binnen gescheiden 'heads' langs de token-dimensie. Wij bewijzen dat MHLA lineaire complexiteit behoudt en tegelijkertijd een groot deel van het expressieve vermogen van softmax-attentie herstelt. We verifiëren de effectiviteit ervan in meerdere domeinen, met een verbetering van 3,6% op ImageNet-classificatie, een winst van 6,3% op NLP, een verbetering van 12,6% op beeldgeneratie en een verbetering van 41% op videogeneratie bij gelijke tijdcomplexiteit.
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.