mlsanityData to 14 Oct 2026

QK-Normed MLA: QK normalization without full key caching

arXiv 2606.16310 ↗First posted 15 Jun 2026cs.LG cs.CLcited by 0

Not on the map. Why?

Abstract (arXiv)

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.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 12 of its 13 references that are in the map. First posted on arXiv 15 Jun 2026.

Builds on

12 of 13 references.

References · 12

  1. 2017Attention Is All You Needcited by 60,361
  2. 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,944
  3. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
  4. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  5. 2024Gemma 2: Improving Open Language Models at a Practical Sizecited by 1,998
  6. 2019Root Mean Square Layer Normalizationcited by 1,480
  7. 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,346
  8. 2024DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelcited by 1,135
  9. 2019Fast Transformer Decoding: One Write-Head is All You Needcited by 717
  10. 2023Scaling Vision Transformers to 22 Billion Parameterscited by 694
  11. 2025Kimi K2: Open Agentic Intelligencecited by 567
  12. 2020Query-Key Normalization for Transformerscited by 280

Built on by

Not cited yet.