MHLA

用 token 级多头，让线性注意力恢复全量表达力

MHLA 是一个面向图像生成与语言建模的线性注意力机制改进项目，旨在解决线性注意力在表达力上的不足。传统线性注意力通过核函数近似 softmax 注意力，虽降低了计算复杂度，但牺牲了模型对复杂 token 关系的建模能力。MHLA 提出在 token 级别引入多头机制，将每个 token 的注意力计算拆分为多个子空间，从而在不显著增加计算成本的前提下恢复线性注意力的表达力。项目基于 PyTorch 实现，支持图像分类（ImageNet）、图像生成（如 Sana 架构）、视频生成和语言建模等任务，提供了与 DiT 等主流架构的集成示例。核心能力包括：高效的线性注意力实现、token 级多头分解策略、以及跨模态任务的验证。该项目为 ICLR 2026 投稿，代码已开源，适合研究高效注意力机制和长序列建模的开发者。