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