MHLA：トークンレベルマルチヘッドによる線形注意機構の表現力回復

要旨

Transformerアーキテクチャは多くの分野で支配的である一方、その二次的な自己注意機構の計算複雑性が大規模応用における使用を妨げている。線形注意機構は効率的な代替手段を提供するが、直接適用すると性能劣化が生じることが多く、既存の修正手法では通常、深度分離可能畳み込みなどの追加モジュールを通じて計算オーバーヘッドを再導入するため、本来の目的が損なわれる。本研究では、これらの手法における主要な失敗モードとして、モデルが表現の多様性を失う「大域的文脈の崩壊」を特定する。この問題に対処するため、トークン次元に沿って分割されたヘッド内で注意を計算することで多様性を保持するMulti-Head Linear Attention（MHLA）を提案する。MHLAが線形複雑性を維持しながらsoftmax注意の表現力の大部分を回復することを理論的に証明し、複数領域でその有効性を検証する。同一時間計算量条件下で、ImageNet分類では3.6%、自然言語処理では6.3%、画像生成では12.6%、動画生成では41%の性能向上を達成した。

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.

MHLA：トークンレベルマルチヘッドによる線形注意機構の表現力回復

MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head

要旨

Support