让高效解码的MLA模型也能用QK归一化，无需缓存完整键值。
QK-Normed MLA: QK normalization without full key caching
- 将QK归一化分解为静态权重与动态统计量，前者融入查询投影，后者简化为每令牌一个标量。
- 400M参数模型训练1000亿词，损失更低，下游任务准确率优于传统截断法。
- 在256k上下文下解码延迟增加不足2%，适合大上下文高效生成场景。
查询-键（QK）归一化通过控制点积前查询和键的尺度来稳定注意力机制，但与多头潜在注意力（MLA）存在兼容性问题。MLA通过缓存低维潜在状态而非完整键值实现高效解码，而传统的投影后QK RMSNorm似乎需要每个缓存令牌的完整投影键。我们发现这种不兼容是实现方式的产物，而非架构限制。RMSNorm可分解为静态仿射权重与动态标量均方根统计量：静态键侧权重可被吸收至MLA查询侧投影；动态统计量则简化为每令牌和键值组一个逆均方根标量。由此得到的公式在精确算术下等价于显式后投影QK RMSNorm，同时保持MLA的潜在解码路径。在400M参数模型上，训练长达1000亿词，QK-Normed MLA相比QK截断达到更低训练损失和更优下游准确率；H800解码测试显示，在256k上下文下延迟增加小于2%。该结果使QK归一化成为无需全键缓存的MLA模型实用稳定选项。
原文摘要 · Abstract (English)
Query-key (QK) normalization stabilizes attention by controlling the scale of queries and keys before the dot product, but is not immediately compatible with Multi-head Latent Attention (MLA). MLA achieves efficient decoding by caching low-dimensional latent states instead of full keys, whereas post-projection QK RMSNorm appears to require the fully projected key for every cached token. We show this apparent incompatibility is an implementation artifact, not an architectural constraint. RMSNorm decomposes into a static affine weight and a dynamic scalar RMS statistic. The static key-side weight can be absorbed into the MLA query-side projection; the dynamic key statistic reduces to one inverse-RMS scalar per token and KV group. The resulting formulation is exactly equivalent to explicit post-projection QK RMSNorm in exact arithmetic and preserves MLA's latent decode path. In our 400M runs trained for up to 100B tokens, QK-Normed MLA achieves lower training loss and better downstream accuracy than QK clipping, while H800 decode benchmarks show less than 2% latency overhead up to 256k context. These results make QK normalization a practical stabilization option for MLA models without requiring full-key caching.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。