MHLA: Wiederherstellung der Ausdruckskraft linearer Aufmerksamkeit durch token-basierten Multi-Head-Ansatz
MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head
January 12, 2026
papers.authors: Kewei Zhang, Ye Huang, Yufan Deng, Jincheng Yu, Junsong Chen, Huan Ling, Enze Xie, Daquan Zhou
cs.AI
papers.abstract
Während die Transformer-Architektur in vielen Bereichen dominiert, behindert ihre quadratische Selbstaufmerksamkeitskomplexität den Einsatz in großmaßstäblichen Anwendungen. Lineare Aufmerksamkeit bietet eine effiziente Alternative, doch ihre direkte Anwendung führt oft zu Leistungseinbußen, wobei bestehende Lösungsansätze typischerweise durch zusätzliche Module (z.B. tiefenweise separierbare Faltungen) erneut Rechenaufwand verursachen, der den ursprünglichen Zweck zunichtemacht. In dieser Arbeit identifizieren wir einen zentralen Schwachpunkt dieser Methoden: den Kollaps des globalen Kontexts, bei dem das Modell seine Repräsentationsvielfalt verliert. Um dies zu beheben, schlagen wir Multi-Head Linear Attention (MHLA) vor, das diese Vielfalt durch die Berechnung von Aufmerksamkeit innerhalb aufgeteilter Köpfe entlang der Token-Dimension erhält. Wir beweisen, dass MHLA bei linearem Komplexitätsaufwand einen Großteil der Ausdruckskraft der Softmax-Aufmerksamkeit zurückgewinnen kann, und bestätigen seine Wirksamkeit in mehreren Domänen: Es erzielt eine Verbesserung von 3,6 % bei der ImageNet-Klassifikation, einen Zuwachs von 6,3 % im NLP-Bereich, eine Steigerung von 12,6 % bei der Bildgenerierung und eine Verbesserung von 41 % bei der Videogenerierung bei gleicher Zeitkomplexität.
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.