多头低秩注意力：支持张量并行的高效注意力机制 Multi-Head Low-Rank Attention
将MLA单潜在头分解为多分支，实现4路张量并行解码，速度提升2.8倍
前置知识
KV缓存（Key-Value Cache）
在自回归语言模型解码过程中，每个token生成时都需要计算与之前所有token的注意力分数。为避免重复计算，模型会将之前token的Key和Value向量缓存起来，这就是KV缓存。对于长上下文推理（如128K-2M token），KV缓存会占用大量显存，并且每一步解码都需要从高带宽内存（HBM）加载到片上SRAM，数据搬运成为主要瓶颈。
本文的核心目标就是减少KV缓存的加载量，理解KV缓存的工作原理和瓶颈是理解论文动机的基础。
张量并行（Tensor Parallelism, TP）
一种模型并行策略，将单个层的权重矩阵分割到多个GPU设备上。在注意力机制中，TP通常将注意力头分配到不同设备上，每个设备只处理一部分头，然后通过通信聚合结果。TP可以减少单设备的计算和内存需求，但要求注意力机制能够被有效分割。
MLA的核心缺陷就是不支持TP——其单一潜在头无法被分割，导致每个设备都需要冗余加载完整KV缓存。MLRA的设计目标就是支持4路TP。
多头潜在注意力（Multi-Head Latent Attention, MLA）
DeepSeek-V2提出的注意力机制，通过将KV状态压缩到一个低维潜在空间来减少KV缓存大小。MLA使用一个共享的潜在头，通过上投影矩阵生成所有注意力头的Key和Value。缓存大小仅为每个token $4.5d_h$（对比MHA的 $2hd_h$）。但这个单一潜在头无法被分割到多个设备上。
MLRA是在MLA基础上的改进，理解MLA的工作原理是理解MLRA创新点的前提。
低秩分解（Low-Rank Decomposition）
将一个大矩阵分解为两个或多个小矩阵的乘积，利用矩阵的低秩特性减少参数和计算量。在本文中，MLRA将MLA的单一上投影矩阵分解为多个块，每个块可以独立计算注意力输出，最后求和得到最终结果。这种分解使得计算可以并行化，同时保持表达能力。
MLRA的核心技术手段就是对MLA的上投影矩阵进行块分解，将求和从KV计算移到注意力输出。
RoPE（Rotary Position Embedding）
旋转位置编码，通过在Query和Key向量上应用旋转矩阵来注入位置信息。RoPE的关键性质是：两个token的注意力分数只取决于它们的相对位置，而非绝对位置。MLA中采用部分RoPE，只对Key的一部分应用RoPE，其余部分保持位置无关。
MLRA需要保持RoPE的平移等变性，同时处理NoPE部分的块分解，这是技术实现的关键约束。
研究动机
在大语言模型的长上下文推理中，解码阶段的瓶颈已经从计算转移到了数据搬运。以128K-2M token的上下文为例，每一步解码都需要将KV缓存从片外HBM加载到片上SRAM，而每步的计算量却很小，导致GPU利用率极低。现有的多头潜在注意力（MLA）虽然通过低秩压缩将KV缓存从 $2hd_h$ 减少到 $4.5d_h$，但在分布式解码场景下存在严重缺陷：MLA的单一潜在头无法被分割到多个设备上进行张量并行。具体来说，当使用4路TP时，MLA的每个设备仍然需要加载完整的 $4.5d_h$ 的KV缓存，无法利用多设备的优势减少单设备负载。这使得MLA在TP大于1时的KV缓存加载量保持恒定，完全丧失了TP带来的带宽收益。
本文的目标是本文的目标是设计一种新的注意力机制，在保持MLA的KV缓存压缩优势的同时，原生支持4路张量并行解码。具体而言，MLRA希望实现三个目标：第一，在4路TP下将每设备KV缓存加载量从MLA的 $4.5d_h$ 降低到 $1.5d_h$；第二，保持或超越MLA的模型质量，包括困惑度和下游任务准确率；第三，实现比MLA更快的解码速度，目标是2倍以上的加速。
与已有工作不同的是，已有工作如GLA-2尝试将MLA的单一潜在头分成两个组来支持2路TP，但这只能将每设备加载量降到 $2.5d_h$，且在TP大于2时加载量不再下降。本文的独特视角是：MLA的上投影矩阵可以被分解为4个块，每个块的投影结果可以独立计算注意力输出，然后求和。这个数学等价性此前被忽视了。MLA的NoPE Key计算公式 $K_{NoPE}^{(i)} = \sum_{b=0}^{3} C_{KV}^{(b)} W_{UK}^{(b),(i)}$ 可以重写为将求和移到注意力输出的形式，即将 $O = \text{Softmax}(Q \cdot K^T) \cdot V$ 中的K和V替换为分块求和的形式。这个视角转换是MLRA的核心创新。
核心方法
MLRA的核心直觉可以用「分而治之」来类比：如果一个大任务无法被分割，那就把它拆成多个小任务，每个小任务独立完成后再合并结果。具体来说，MLA的单一潜在头就像一个不可分割的大任务，而MLRA将其拆成4个可独立计算的小任务。技术路线是：首先证明MLA的注意力输出可以数学等价地表示为4个独立分支的注意力输出之和（这是关键的数学洞察）；然后将这个等价形式显式地实现在模型架构中，每个分支有自己独立的潜在头、上投影矩阵和注意力计算；最后通过方差校正确保多分支求和后的数值稳定性。这个设计天然支持4路TP：每个设备只需要处理1个分支的潜在头，KV缓存加载量降为原来的四分之一。
MLRA与MLA最本质的区别在于求和的位置。在MLA中，计算流程是：先将4个块的KV投影求和得到完整的Key和Value，然后计算注意力输出。而在MLRA中，计算流程是：每个块独立计算自己的Key和Value，独立计算注意力输出，最后将4个注意力输出求和。数学上，MLA的输出是 $O = \text{Softmax}(Q \cdot (\sum_{b} K_b)^T) \cdot (\sum_{b} V_b)$，而MLRA的输出是 $O = \sum_{b} \text{Softmax}(Q \cdot K_b^T) \cdot V_b$。这两个表达式在数学上不等价（因为Softmax是非线性的），但MLRA通过实验证明这种近似不仅可接受，甚至能带来更好的性能。这个改变的核心意义是：每个分支的计算完全独立，只需要在最后做一次简单的加法，天然支持分布式并行。
方法步骤详情
MLRA的计算流程分为三个主要步骤。第一步是查询侧权重吸收：将上投影矩阵 $W_{UK}$ 和 $W_{UV}$ 按头和块重组为张量，然后通过爱因斯坦求和将上投影权重吸收到查询中。第二步是MQA风格解码：将KV缓存 $C_{KV}$ 和 $K_{RoPE}$ 拼接重塑为共享的Key和Value张量，然后计算注意力输出。第三步是输出上投影：通过另一个爱因斯坦求和将中间结果映射回原始维度。对于MLRA-4变体，整个流程被分解为4个独立分支并行执行，每个分支处理 $C_{KV}$ 的四分之一块，最后将4个分支的输出求和。此外，MLRA还引入了方差校正机制，通过缩放因子 $\alpha_q = \sqrt{d/d_c'}$ 和 $\alpha_{kv} = \sqrt{d/d_c}$ 确保NoPE组件与RoPE组件的方差匹配，以及注意力输出的缩放因子校正多分支求和带来的方差变化。
技术新颖性
MLRA的技术新颖性体现在三个层面。第一，数学洞察：首次发现MLA的上投影矩阵可以被分解为4个块，且每个块的投影结果可以独立计算注意力，这是一个非平凡的数学发现，因为Softmax的非线性意味着先求和再Softmax和先Softmax再求和在数学上不等价。第二，架构设计：将这个数学洞察转化为实际可执行的架构，每个分支有独立的潜在头和上投影矩阵，但共享RoPE Key，这在保持参数效率的同时实现了计算并行。第三，方差校正：系统性地分析了RoPE Key、NoPE Key、Value和Query的方差关系，提出了精确的缩放因子，这不仅适用于MLRA，也适用于MLA和GLA等现有方法。与GLA-2相比，MLRA-4将分支数从2增加到4，将每设备KV加载量从 $2.5d_h$ 降到 $1.5d_h$。
实验结果
MLRA在2.9B参数规模的实验中展现了全面的性能优势。在困惑度评估方面，MLRA-4在7个数据集上的平均困惑度为13.672，优于MLA的13.727（降低0.055）和GQA的14.139（降低0.467）。在最具代表性的FineWeb-Edu数据集上，MLRA-4的困惑度为9.193，同样创下最佳。MLRA-2的平均困惑度为13.804，介于MLA和GLA-2之间，说明增加分支数确实能带来性能提升。在零样本常识推理任务上，MLRA-4在7个基准测试上的平均准确率为58.84%，超过MLA的58.75%和GQA的57.89%。特别是在HellaSwag（62.16%对比61.39%）和Winogrande（61.48%对比60.06%）上提升明显。消融实验表明：零初始化优于高斯初始化；方差校正缩放对MLA、GLA-2和MLRA-2都有帮助；双倍注意力头数并不能提升性能，说明MLRA的优势来自分支设计而非头数增加；门控注意力能进一步提升所有模型的性能，MLRA-4加上门控达到最佳平均困惑度13.621。解码效率方面，MLRA-4在H100 GPU上的长上下文解码延迟测试中，在所有序列长度（128K-2M token）上都优于所有基线，比MLA快约2.8倍，比GQA快1.05-1.26倍。在8 GPU的吞吐量测试中，MLRA-4在TP=4、DP=2配置下在所有序列长度上都实现了最高吞吐量。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 平均困惑度（7数据集） | Perplexity | 13.672 | MLA 13.727 | 降低0.055 |
| FineWeb-Edu困惑度 | Perplexity | 9.193 | MLA 9.284 | 降低0.091 |
| 零样本常识推理（7任务平均） | Accuracy (%) | 58.84 | MLA 58.75 | 提升0.09个百分点 |
| HellaSwag | Accuracy (%) | 62.16 | MLA 61.39 | 提升0.77个百分点 |
| Winogrande | Accuracy (%) | 61.48 | MLA 60.06 | 提升1.42个百分点 |
| 长上下文解码延迟（对比MLA） | Speedup | 2.8倍 | MLA | 2.8倍加速 |
| 每设备KV缓存加载（4路TP） | Loading per device | 1.5dh | MLA 4.5dh | 减少67% |
| 解码算术强度 | Arithmetic Intensity | 2h (MLRA-4) | MLA约2h | 与MLA相当 |
局限与改进
本文存在几个明显的局限性。首先，实验规模仅限于2.9B参数，这是相对较小的模型规模。虽然结果一致且有说服力，但尚不清楚这些优势能否在更大的模型上保持。作者在论文中没有讨论更大规模的实验或计划。其次，MLRA-4的4路TP设计是固定的，这意味着在使用2路或8路TP时可能不是最优的。论文中MLRA-2变体虽然支持2路TP，但性能不如MLRA-4。第三，MLRA引入了额外的参数（4个独立的上投影矩阵对比MLA的1个），虽然通过调整FFN维度保持了总参数量不变，但这可能影响了FFN的表达能力。第四，方差校正机制基于权重矩阵独立同分布且与输入无关的假设，这个假设在实际训练中可能不严格成立，作者也承认这一点并建议通过消融实验验证。第五，论文没有讨论MLRA在微调或RLHF场景下的表现，这些场景对注意力机制的行为有不同要求。
独立分析的弱点
从独立分析的角度，MLRA有几个值得关注的弱点。第一，分支独立性假设的脆弱性：MLRA假设4个分支的注意力输出是相互独立的，但在实际模型中，不同块的潜在状态可能高度相关，特别是在训练后期。这可能导致多分支求和后的次优结果。改进方向是引入分支间的交互机制，例如在求和前对各分支输出进行加权，或者在训练中加入正交性约束。第二，固定的4路并行度：MLRA-4被硬编码为4个分支，这限制了其在不同TP配置下的灵活性。更好的设计是支持可配置的分支数，根据实际硬件配置动态调整。第三，对FlashAttention的依赖：MLRA的高效实现依赖于FlashAttention-3和FlashMLA内核，这些内核针对特定的硬件和数据类型进行了优化。在不支持这些内核的硬件上，MLRA的性能优势可能无法体现。第四，缺乏与稀疏注意力方法的对比：论文只比较了全注意力的变体，没有与MInference、H2O等稀疏注意力方法进行对比。
未来方向
基于MLRA的成果，有几个有前景的研究方向。第一，大规模验证：在7B、70B甚至400B参数规模上验证MLRA的有效性，特别是探索分支数是否需要随模型规模调整。第二，动态分支分配：设计一种机制，让模型在推理时根据输入内容动态选择激活哪些分支，类似于混合专家的思想。第三，与其他压缩技术的结合：将MLRA与KV缓存量化、稀疏注意力、滑动窗口等技术结合，探索协同效应。第四，训练效率优化：当前MLRA的训练成本与MLA相当，但多分支设计可能允许更高效的分布式训练策略。第五，扩展到其他模态：MLRA的块分解思想可以应用到视觉Transformer或跨模态注意力中，探索其通用性。
复现评估
MLRA的复现条件非常友好。作者提供了完整的代码库和预训练权重，包括训练和评估数据。代码基于nanoGPT实现，这是一个广泛使用的简洁Transformer代码库，易于理解和修改。训练配置详细记录在论文附录中，包括所有超参数。训练硬件为8块H100 80GB GPU，这对于学术实验室是可获得的。评估使用标准的lm-evaluation-harness工具，复现门槛低。主要的复现难点在于：需要实现自定义的MLRA注意力内核，虽然作者提供了参考实现，但在不同硬件上可能需要适配；需要正确实现方差校正的缩放因子，否则可能导致训练不稳定。总体而言，复现难度为中等偏易。
论文图表
