MHLA: 通过令牌级多头恢复线性注意力的表达性
计算机视觉与模式识别
2026-01-15 v2 人工智能
摘要
虽然 Transformer 架构主导着许多领域，但其二次自注意力复杂度阻碍了其在大规模应用中的使用。线性注意力提供了一种高效替代方案，但其直接应用通常会导致性能下降，现有修复方法通常通过额外模块（如深度可分离卷积）重新引入计算开销，违背了原始目的。在本工作中，我们识别到了这些方法的一个关键失效模式：全局上下文坍塌，即模型丢失表示多样性。为此，我们提出了多头线性注意力 (MHLA)，通过在令牌维度上对注意力进行划分来计算，从而保持这种多样性。我们证明了 MHLA 在保持线性复杂度的同时，恢复了大量 softmax 注意力的表达能力，并在多个领域验证了其有效性，在相同的时间复杂度下，ImageNet 分类上提升 3.6\,\%，NLP 上提升 6.3\,\%，图像生成上提升 12.6\,\%，视频生成上提升 41\,\%。
关键词
引用
@article{arxiv.2601.07832,
title = {MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head},
author = {Kewei Zhang and Ye Huang and Yufan Deng and Jincheng Yu and Junsong Chen and Huan Ling and Enze Xie and Daquan Zhou},
journal= {arXiv preprint arXiv:2601.07832},
year = {2026}
}
备注
Code: https://github.com/DAGroup-PKU/MHLA/ Project website: https://dagroup-pku.github.io/MHLA/