🤖 machine learning

QK-Normed MLA: QK normalization without full key caching

本文证明了通过将操作进行数学重构以避免全量键缓存（key caching），QK 归一化可以无缝集成到多头潜在注意力（Multi-head Latent Attention, MLA）中，从而在实现训练稳定性提升和下游准确度提高的同时，仅产生微乎其微的延迟开销。

原作者： Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun

发布于 2026-06-16

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下你正在经营一家规模宏大的图书馆，一位图书管理员（AI）需要找到完美的书籍（信息）来回答一个问题。为了高效地完成这项工作，图书管理员使用了两个主要窍门：

“潜变量摘要”窍门 (MLA)： 为了不让房间里堆满每一本已写就的巨型、沉重的书籍（这太占空间了），图书管理员为每本书保留了一张微小的、压缩过的“摘要卡”。当问题进来时，图书管理员会即时展开这张摘要卡以寻找答案。这节省了大量的空间。
“音量控制”窍门 (QK Normalization)： 有时候，图书管理员会变得过于兴奋。问题与书籍之间的“音量”变得太大，盖过了其他一切，导致图书管理员陷入恐慌或犯错。为了解决这个问题，我们通常会在书上安装一个“音量旋钮”，以保持声音水平稳定。

问题所在：
长期以来，这两个窍门并不兼容。“音量控制”窍门似乎需要图书管理员在手中握有书籍的完整、沉重的副本才能调节音量。但“潜变量摘要”窍门的设计初衷恰恰是为了避免持有这些沉重的副本。这似乎让你必须做出选择：要么节省空间（潜变量摘要），要么保持音量稳定（音量控制），但两者不可兼得。

解决方案 (QK-Normed MLA)：
论文的作者们发现了一个聪明的数学黑科技，让你能鱼与熊掌兼得。他们意识到，这个“音量控制”旋钮实际上有两个部分：

一个静态设置（一个永不改变的固定刻度）。
一次动态读取（快速瞥一眼当前这本书的声音有多大）。

他们证明了：

静态设置可以被移动到图书管理员查看摘要卡之前。它被直接植入到了问题本身之中。
动态读取非常简单，以至于图书管理员不需要整本书，只需要为每张摘要卡记录下一个极小的数字（一个标量）即可。

结果：
现在，图书管理员可以继续使用微小的“摘要卡”（节省大量空间），同时仍拥有完美的音量控制。他们不再需要搬运那些沉重的书籍了。

实验表明：
团队在一个拥有 4 亿参数的模型（中型 AI）上测试了这种新方法，该模型是在海量文本（1000 亿词）上训练而成的。

更好的学习能力： 使用这种新“音量控制”的模型学习得更快，且比使用另一种更粗糙的方法——“裁剪”（类似于在声音过大时大喊“停！”）的模型犯错更少。
更好的答案： 在各项任务测试中，新方法给出了更准确的答案。
速度： 唯一的缺点是检查那个微小的数字会增加不到 2% 的阅读长篇故事的时间。这一点微乎其微，几乎察觉不到。

简而言之： 论文证明了你不需要为了保持 AI 的稳定性而牺牲内存效率。你可以使用一个微小且聪明的数学技巧来控制“音量”，而无需在最初试图避免存储的那些沉重、全尺寸的数据中进行操作。

技术摘要：QK-Normed MLA

问题陈述
现代大语言模型面临着两种相互竞争的压力：在长文本场景下对内存高效解码的需求，以及对训练稳定性的需求。

内存效率： 标准的多头注意力（MHA）机制会为每个 token 缓存完整的键（key）和值（value）向量，从而产生高昂的内存和带宽成本。多头潜在注意力（MLA）通过缓存低维潜在状态（ $c_t$ ）而非完整的键来解决这一问题，并在解码过程中通过上投影（ $W^K_g$ ）恢复完整的键效应。
训练稳定性： 查询（query）和键（key）中不受控制的尺度可能导致注意力饱和及训练不稳定。查询-键（QK）归一化（特别是 RMSNorm）通过在点积前对向量进行归一化来稳定训练。然而，标准的后投影 QK RMSNorm 在 MLA 中似乎是不兼容的，因为它看似需要为每个缓存的 token 计算完整的投影键（ $k_{t,g} = c_t W^K_g$ ）以获取归一化统计量。一种天真的实现方式将需要实例化完整的键，从而抵消了 MLA 的内存优势。

方法论
本文提出了 QK-Normed MLA，这是一种将后上投影 QK RMSNorm 集成到 MLA 潜在解码路径中的公式化方法。其核心洞察在于 RMSNorm 的代数分解：
$\text{RMSNorm}(z) = z \odot \frac{\gamma}{S(z)}$
其中 $\gamma$ 是静态仿射权重，而 $S(z)$ 是动态标量 RMS 统计量。

作者推导出了一个“吸收”后的公式，将这些组件分离：

静态权重吸收： 静态键侧权重 $\gamma_k$ 被吸收进查询侧的投影矩阵中。与其先投影查询 $q$ 再归一化键，不如在潜在点积之前将查询预乘 $\gamma_k$ （以及静态权重 $\gamma_q$ ）。
动态标量缓存： 动态统计量 $S(z)$ 简化为每个 token 和每个 KV 组一个单一的反 RMS 标量（ $\alpha_{t,g} = 1/S_{t,g}$ ）。
缓存布局： 解码缓存保留原始的低维潜在状态 $c_t$ ，并增加一个包含反 RMS 值的微小标量缓存 $A^k_{cache}$ 。它并不存储完整的归一化键。
分块归一化： 为了保持效率，该方法分别对“内容（content）”和“RoPE”组件进行归一化。内容路径使用吸收后的潜在公式，而 RoPE 路径使用标准的实例化向量。这避免了会使融合算子实现变得复杂的同步点。

在解码期间，注意力逻辑值（logits）计算如下：
$\ell_{i,t,h} = (\tilde{q}_{i,h} \cdot c_t) \cdot \alpha^q_{i,h} \cdot \alpha^k_{t,g(h)}$
其中 $\tilde{q}$ 包含了吸收后的权重，标量因子作为 Softmax 前的行-列缩放步骤进行应用。

核心贡献

精确公式： 本文推导出了一个与 MLA 潜在解码路径兼容的、关于后上投影 QK RMSNorm 的代数精确公式。
高效缓存设计： 提出了一种缓存布局，在保留潜在 KV 缓存作为主导状态的同时，仅增加每个 token 和每个 KV 组一个标量键 RMS 统计量（例如，在测试配置中使缓存大小仅增加了约 1.56%）。
实证验证： 作者证明了该方法与 QK 裁剪（QK Clipping）相比，能实现更低的训练损失和更好的下游任务准确度，且引入的延迟开销微乎其微。

结果
实验在 H800 GPU 上进行的、训练时长达 100B token 的 400M 参数量 MLA 模型上完成：

训练稳定性： 在整个 100B token 的运行过程中，QK-Normed MLA 的训练损失始终低于 QK-Clip。在（ $2 \times 10^{-2}$ ）的高学习率压力测试中，QK-Clip 因梯度出现 NaN 而发散，而 QK-Normed MLA 保持稳定。
下游性能： 在一系列下游任务中，QK-Normed MLA 在 8 个任务中的 7 个上提升了准确度，将平均得分从 44.75 提高到 46.33，并将 LAMBADA 困惑度从 16.28 降低至 14.18。
解码效率： 在 4k 到 256k token 上下文的基准测试显示，该方法增加的延迟不到 2%（平均 1.40%）。标量缓存使每层的总缓存大小仅增加了约 2.8%。

意义与主张
论文声称，QK 归一化与 MLA 之间的不兼容性是一个实现层面的伪命题，而非架构层面的约束。通过展示静态权重可以被吸收且动态统计量可以简化为标量，作者使得 QK 归一化能够作为一种实用的 MLA 模型稳定机制，同时不牺牲潜在解码的内存效率。

作者对其主张的范围保持谦逊：

结果是基于 400M 参数模型的；他们并未声称确立了十亿级参数模型的规模化行为，尽管代数等价性在任何规模下均成立。
对 Logit 尺度的分析被呈现为一种机制层面的诊断，而非完整的优化理论。
实际速度取决于算子集成（将标量因子融合进注意力路径），以避免不必要的内存传输。

技术摘要：QK-Normed MLA

类似论文