用分头机制让线性注意力恢复表达力，速度不变性能大升
MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head
- 将输入分头在令牌维度做线性注意力，避免全局信息坍塌
- 图像分类提升3.6%，视频生成提速41%且保持同复杂度
- 无需额外模块，适合对效率和表现兼得的场景
尽管Transformer架构主导多个领域，但其二次方自注意力复杂度限制了大规模应用。线性注意力提供了高效替代方案，但直接使用常导致性能下降；现有修复方法多引入额外模块（如深度可分离卷积），增加计算开销，违背初衷。本文识别出关键缺陷：全局上下文坍塌，导致表征多样性丧失。为此提出多头线性注意力（MHLA），通过在分头中沿令牌维度计算注意力，保留表征多样性。理论证明MHLA维持线性复杂度的同时，恢复了接近softmax注意力的表达能力。在多领域验证有效，在相同时间复杂度下，图像分类提升3.6%，自然语言处理增益6.3%，图像生成改进12.6%，视频生成性能提高41%。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。