MHLA：通过 Token 级多头机制恢复线性注意力的表达能力 MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head
提出多头线性注意力MHLA，通过token维度分块解决全局上下文坍塌问题
前置知识
自注意力机制（Self-Attention）
Transformer的核心模块，通过计算查询Q、键K、值V之间的相似度来聚合信息。标准自注意力使用softmax函数计算注意力权重，需要计算所有token对之间的相似度，时间复杂度为O(N²d)，其中N是序列长度，d是特征维度。这种机制能够为每个查询生成独特的注意力分布，但二次复杂度限制了其在长序列任务中的应用。
理解softmax注意力的优势（表达能力强）和劣势（复杂度高）是理解本文动机的基础
线性注意力（Linear Attention）
用正特征映射φ(·)替代softmax核，使得Sim(Q_i, K_j) ≈ φ(Q_i)φ(K_j)ᵀ。通过将分子分母预计算为全局KV摘要 G = Σ_j φ(K_j)ᵀV_j 和归一化因子 z = Σ_m φ(K_m)，复杂度从O(N²)降至O(Nd_φ)。然而这种设计牺牲了查询特异性——所有查询共享同一个全局摘要。
这是本文要改进的核心对象，理解其工作原理和局限性是理解MHLA的前提
注意力矩阵的秩（Rank of Attention Matrix）
注意力矩阵 A ∈ R^{N×N} 的秩反映了模型能够表示的不同注意力模式的数量。秩越高，模型能够捕获的token间交互模式越丰富多样。对于线性注意力，A_lin = Q̃K̃ᵀ，其秩受限于 min{rank(Q̃), rank(K̃)} ≤ d，这意味着无论序列多长，表达能力都被严格限制在d维。
论文用秩作为量化'全局上下文坍塌'的核心指标，理解秩的含义才能理解为什么线性注意力性能下降
全局上下文坍塌（Global Context Collapse）
论文提出的新概念，指线性注意力中所有token被压缩到单一全局KV摘要中，导致两个问题：1）秩受限——注意力矩阵秩被限制在d以下；2）稀疏性丧失——随着序列增长，每个token的贡献变得微不足道，注意力分布趋向均匀。这使得模型无法选择性地关注相关token，性能随序列长度增加而下降。
这是论文诊断出的核心病根，MHLA的设计直接针对这一问题
研究动机
标准自注意力的O(N²)复杂度严重限制了其在高分辨率图像生成、长视频合成等大规模任务中的应用。线性注意力通过将softmax替换为核方法，将复杂度降至O(N)，但直接应用会导致显著的性能下降。现有补救方案（如深度可分离卷积、门控模块）虽然能部分恢复性能，但引入了额外的计算开销，违背了线性注意力追求高效的初衷。论文深入分析了线性注意力的失败模式：全局上下文坍塌。具体而言，所有查询共享一个固定的d×d全局KV摘要，当序列长度N远大于d时，这个固定容量的矩阵无法承载所有信息。实验证据表明，线性注意力的注意力矩阵秩被严格限制在头维度d_h ≤ 72，而softmax注意力可以达到N。同时，线性注意力的注意力熵显著高于softmax注意力（5.09 vs 4.35），表明其注意力分布过于均匀，无法有效聚焦。
本文的目标是论文的目标是设计一种新的注意力机制，能够在保持O(N)线性复杂度的前提下，恢复查询条件化的token级多样性。具体而言，不同查询应该能够检索到不同的上下文信息，而不是被迫使用同一个全局摘要。同时，该机制不应依赖额外的深度卷积或自注意力模块，以保持计算效率。
与已有工作不同的是，论文的独特切入角度是从token维度而非特征维度进行多头划分。传统多头注意力在特征维度（d）上划分头，每头独立计算注意力；而MHLA在token维度（N）上划分头，将序列分成M个非重叠块。每个块计算局部KV摘要，然后通过可学习的混合系数矩阵 M_c ∈ R^{M×M} 将这些局部摘要组合成查询特定的混合摘要。这种设计使得注意力矩阵的秩从d提升到 Σ_{b=1}^M min(n_b, d)，从而恢复了大部分softmax注意力的表达能力，同时仅引入O(M²d²)的额外开销，当M² ≤ N时仍保持线性复杂度。
核心方法
MHLA的核心思想可以用一个直觉来理解：与其让所有查询共享一个'大而全'的全局摘要，不如将序列分成多个'小组'，每个小组维护自己的局部摘要，然后让每个查询块学习如何混合这些局部摘要。这类似于一个图书馆系统：传统线性注意力是把所有书压缩成一份总目录，所有读者都用同一份目录找书；而MHLA是每个分馆有自己的目录，读者根据自己的需求组合使用不同分馆的目录。技术路线上，MHLA首先将输入序列 X ∈ R^{N×d} 投影为Q、K、V，然后沿token维度将其划分为M个非重叠块。对每个块b计算局部KV摘要 S_b = Σ_{j∈b} K̃_j V_jᵀ ∈ R^{d×d}。核心创新是多头混合模块：引入可学习系数矩阵 M_c ∈ R^{M×M}，元素 m_{i,j} 表示查询块i对键值块j的亲和度，使得每个查询块能获得独特的混合摘要 S̃_i = Σ_{b=1}^M m_{i,b} S_b。
MHLA与已有方法的本质区别在于：它不通过增加辅助模块（如卷积、门控）来'修补'线性注意力，而是从根本上改变了信息聚合的结构。传统线性注意力将所有token压缩到一个全局摘要，查询特异性完全丧失；而MHLA通过在token维度上分块并学习块间混合，恢复了查询条件化的选择性。从数学上看，对于查询向量 q̃ = φ(q)（来自块i），其输出的分子为 q̃ᵀS̃_i = Σ_{t=1}^N m_{i,b(t)} q̃ᵀK̃_t V_tᵀ，这意味着每个查询块通过m_i对整个块进行重缩放，同时在每个块内通过核内积 q̃ᵀK̃_t 区分token贡献。这种'块选择×块内重加权'的两阶段机制是MHLA恢复token级多样性的关键。
方法步骤详情
MHLA的完整流程如下：第一步，输入序列 X ∈ R^{N×d} 通过线性投影得到查询 Q = XW_Q、键 K = XW_K、值 V = XW_V，并应用特征映射 Q̃ = φ(Q)、K̃ = φ(K)。第二步，将序列沿token维度划分为M个非重叠块，在视觉模型中按照空间（2D）或时空（3D）网格定义块，而非简单展平为1D。第三步，对每个块b计算局部KV摘要 S_b = Σ_{j∈b} K̃_j V_jᵀ ∈ R^{d×d} 和归一化因子 z_b = Σ_{j∈b} K̃_j ∈ R^d。第四步，通过可学习系数矩阵 M_c 进行多头混合，为每个查询块i生成混合摘要 S̃_i = Σ_{b=1}^M m_{i,b} S_b 和混合归一化因子 z̃_i = Σ_{b=1}^M m_{i,b} z_b。第五步，计算输出 Y_i = (Q̃_i S̃_i) / (Q̃_i z̃_i)。系数矩阵 M_c 使用局部偏好初始化：m_{i,j}^(0) ∝ 1 - dist(i,j)/max_k(dist(i,k))，并归一化使得 Σ_j m_{i,j}^(0) = 1，训练过程中可学习并裁剪到(0,1)区间。
技术新颖性
MHLA的技术新颖性体现在多个层面。首先，它是首个在token维度上进行多头划分的线性注意力机制，与传统多头注意力在特征维度划分形成互补。其次，可学习混合系数矩阵 M_c 的设计使得模型能够端到端地学习块间依赖关系，而无需引入卷积或门控等辅助模块。第三，局部偏好初始化策略利用空间邻近性作为先验，既保证了训练稳定性，又不限制模型的适应能力。第四，MHLA自然适配分块并行训练范式，每个头可以直接映射到一个chunk，只需维护一个局部摘要 S_b，训练时使用混合前缀摘要 S̃_i = Σ_{b≤i} m_{i,b} S_b，推理时可支持流式/有状态执行。从理论分析看，MHLA将注意力矩阵的秩界从d提升到 Σ_{b=1}^M min(n_b, d)，这在N远大于d时是显著的提升。
实验结果
论文在图像分类、图像生成、自然语言处理和视频生成四个领域进行了全面验证。在图像分类任务中，MHLA在DeiT-T模型上达到75.8%的Top-1准确率，比标准自注意力（72.2%）高出3.6个百分点，同时仅引入最少的额外参数（5.7M vs 5.7M），FLOPs保持不变（1.1G）。在图像生成任务中，MHLA在DiT-S/2模型上将FID从68.40降至59.80（256分辨率），提升12.6%；在DiT-XL/2上，MHLA的FID为19.17，甚至略优于自注意力的19.47。在SANA模型的快速适应实验中，MHLA在2k步内即匹配预训练checkpoint，并最终将FID从6.10降至5.90。在视频生成任务中，MHLA在Wan2.1-1.3B上展现出巨大优势：vanilla线性注意力因全局上下文坍塌导致严重退化（Quality从85.23降至69.96），而MHLA恢复到84.26，同时实现2.1倍推理加速。混合模型（2/3层替换为MHLA）进一步将Quality提升至84.87，达到1.6倍加速。在NLP任务中，MHLA（340M参数）在MMLU上达到23.7，超过所有基线；在LongBench上的平均分达到7.41，领先于GDN（6.86）和Mamba2（6.62）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 图像分类（ImageNet-1K，DeiT-T） | Top-1 Accuracy (%) | 75.8 | 72.2（Self-Attention）、69.8（Linear Attention） | +3.6pp（vs Self-Attention） |
| 图像分类（ImageNet-1K，DeiT-S） | Top-1 Accuracy (%) | 81.0 | 79.8（Self-Attention）、77.6（Linear Attention） | +1.2pp（vs Self-Attention） |
| 图像生成（DiT-S/2，256分辨率） | FID ↓ | 59.80 | 68.40（Self-Attention）、89.72（Linear Attention） | 12.6%（vs Self-Attention） |
| 图像生成（DiT-XL/2，256分辨率） | FID ↓ | 19.17（w/ CPE+Gating） | 19.47（Self-Attention） | 与Self-Attention持平 |
| 视频生成（Wan2.1-1.3B） | Quality Score | 84.26（MHLA）、84.87（MHLA-Hybrid） | 85.23（Flash-Attention）、69.96（Linear Attention） | 2.1倍推理加速（vs Flash-Attention） |
| NLP（340M模型，MMLU） | MMLU Accuracy | 23.7 | 22.9（Transformer++、GLA） | +0.8pp（vs Transformer++） |
局限与改进
论文存在若干局限性。首先，MHLA的理论秩界 Σ_{b=1}^M min(n_b, d) 虽然远高于d，但仍低于softmax注意力的N，在超长序列场景下可能仍有差距。其次，论文的NLP实验规模较小（340M参数、10B tokens），尚未验证MHLA在更大规模语言模型（如7B、13B）上的表现，Wiki困惑度（71.64）显著高于Transformer++（60.46）表明在语言建模任务上仍有改进空间。第三，系数矩阵 M_c ∈ R^{M×M} 引入了O(M²)的额外复杂度，当M较大时可能成为瓶颈。第四，论文主要在视觉和NLP任务上验证，未涉及语音、多模态等其他领域。第五，消融实验表明，DW-Conv（CPE）等辅助模块在小模型上有增益但在大模型上效果减弱甚至退化，说明MHLA的优势在大规模场景下更为明显，但在小规模场景下可能需要辅助设计。
独立分析的弱点
独立分析发现以下弱点：第一，初始化策略依赖于欧氏距离计算，在非网格结构（如文本序列）中，这种空间先验可能不适用，需要探索其他初始化方案。第二，系数矩阵的裁剪操作（限制在(0,1)区间）虽然保证了稳定性，但可能限制了模型的表达能力，特别是在需要学习稀疏或尖锐注意力分布的场景。第三，论文未讨论M的选择策略，仅在消融实验中测试了4、16、64三个值，缺乏自适应选择机制。第四，在视频生成实验中，序列长度为31,500 tokens时M=105，这意味着平均每个块约300 tokens，块内仍然使用标准线性注意力，可能仍有局部信息丢失。第五，论文声称MHLA兼容流式执行，但未给出具体的流式推理算法和复杂度分析。
未来方向
基于论文成果，未来研究可向以下方向拓展：第一，探索自适应分块策略，根据输入内容动态调整块的大小和数量，而非固定M。第二，将MHLA与稀疏注意力结合，在块间混合的基础上引入块内的稀疏选择，进一步提升效率。第三，在更大规模的语言模型（7B+）上验证MHLA，特别是在长上下文理解和生成任务上。第四，研究MHLA在多模态模型中的应用，特别是视觉-语言模型中处理高分辨率图像和长文本的场景。第五，探索MHLA与硬件优化的协同设计，利用块结构的特点进行算子融合和内存优化。第六，将MHLA的分块混合思想扩展到其他高效Transformer变体（如状态空间模型），探索其通用性。
复现评估
论文在复现性方面提供了较好的支持。代码已开源在GitHub（https://github.com/DAGroup-PKU/MHLA）和HuggingFace（https://huggingface.co/DAGroup-PKU/MHLA），并提供了项目主页。论文使用标准的ImageNet-1K数据集进行图像分类和生成实验，使用FineWeb-Edu进行NLP实验，这些数据集均可公开获取。训练配置（300 epochs、batch size 1024、learning rate 1e-3等）描述清晰。然而，视频生成实验使用的是Wan2.1-1.3B预训练模型，该模型的获取可能需要额外步骤。计算资源方面，论文提到吞吐量测试在NVIDIA H100上进行，大规模实验可能需要多卡训练环境。总体而言，中等计算能力的实验室可以复现主要结果，但视频生成等大规模实验可能需要较多资源。
论文图表
