🤖 machine learning

QK-Normed MLA: QK normalization without full key caching

This paper demonstrates that QK normalization can be seamlessly integrated with Multi-head Latent Attention (MLA) by mathematically reformulating the operation to avoid full key caching, thereby achieving improved training stability and downstream accuracy with negligible latency overhead.

Original authors: Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun

Published 2026-06-16

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive library where a librarian (the AI) needs to find the perfect book (information) to answer a question. To do this efficiently, the librarian uses two main tricks:

The "Latent Summary" Trick (MLA): Instead of keeping a giant, heavy copy of every single book ever written in the room (which takes up too much space), the librarian keeps a tiny, compressed "summary card" for each book. When a question comes in, the librarian expands the summary card on the fly to find the answer. This saves huge amounts of space.
The "Volume Control" Trick (QK Normalization): Sometimes, the librarian gets too excited. The "volume" of the connection between a question and a book gets so loud that it drowns out everything else, causing the librarian to panic or make mistakes. To fix this, we usually put a "volume knob" on the books to keep the sound level steady.

The Problem:
For a long time, these two tricks didn't get along. The "Volume Control" trick seemed to require the librarian to have the full, heavy copy of the book in hand to adjust the volume. But the "Latent Summary" trick was specifically designed to avoid having those heavy copies. It seemed like you had to choose: either save space (Latent Summary) or keep the volume steady (Volume Control), but not both.

The Solution (QK-Normed MLA):
The authors of this paper discovered a clever math hack that lets you have your cake and eat it too. They realized that the "Volume Control" knob actually has two parts:

A static setting (a fixed dial that never changes).
A dynamic reading (a quick glance at how loud the current book is).

They showed that:

The static setting can be moved before the librarian even looks at the summary card. It gets baked into the question itself.
The dynamic reading is so simple that instead of needing the whole book, the librarian only needs to write down one tiny number (a single scalar) for each summary card.

The Result:
Now, the librarian can keep using the tiny "summary cards" (saving massive amounts of space) while still having perfect volume control. They don't need to carry the heavy books anymore.

What the Experiments Showed:
The team tested this new method on a 400-million-parameter model (a medium-sized AI) trained on a massive amount of text (100 billion words).

Better Learning: The model with this new "Volume Control" learned faster and made fewer mistakes than models that used a different, cruder method called "clipping" (which is like just shouting "STOP!" when things get too loud).
Better Answers: When tested on various tasks, the new method gave more accurate answers.
Speed: The only downside was checking that one tiny number, which added less than 2% to the time it took to read a long story. This is so small it's barely noticeable.

In short: The paper proves you don't have to sacrifice memory efficiency to keep your AI stable. You can use a tiny, clever math trick to control the "volume" without ever needing to store the heavy, full-sized data you were trying to avoid in the first place.

Technical Summary: QK-Normed MLA

Problem Statement
Modern large language models face two competing pressures: the need for memory-efficient decoding in long-context scenarios and the need for training stability.

Memory Efficiency: Standard Multi-Head Attention (MHA) caches full key and value vectors for every token, creating high memory and bandwidth costs. Multi-Head Latent Attention (MLA) addresses this by caching low-dimensional latent states ( $c_t$ ) rather than full keys, recovering the full key effect via an up-projection ( $W^K_g$ ) during decoding.
Training Stability: Uncontrolled scales in queries and keys can lead to attention saturation and training instabilities. Query-Key (QK) normalization (specifically RMSNorm) stabilizes training by normalizing vectors before the dot product. However, standard post-projection QK RMSNorm appears incompatible with MLA because it seemingly requires the fully projected key ( $k_{t,g} = c_t W^K_g$ ) for every cached token to compute the normalization statistic. A naive implementation would require materializing full keys, negating MLA's memory advantages.

Methodology
The paper proposes QK-Normed MLA, a formulation that integrates post-up-projection QK RMSNorm into the MLA latent decode path without materializing full keys. The core insight is the algebraic decomposition of RMSNorm:
$\text{RMSNorm}(z) = z \odot \frac{\gamma}{S(z)}$
where $\gamma$ is a static affine weight and $S(z)$ is a dynamic scalar RMS statistic.

The authors derive an "absorbed" formulation that separates these components:

Static Weight Absorption: The static key-side weight $\gamma_k$ is absorbed into the query-side projection matrix. Instead of projecting the query $q$ and then normalizing the key, the query is pre-multiplied by $\gamma_k$ (and the static weight $\gamma_q$ ) before the latent dot product.
Dynamic Scalar Caching: The dynamic statistic $S(z)$ reduces to a single inverse-RMS scalar ( $\alpha_{t,g} = 1/S_{t,g}$ ) per token and per KV group.
Cache Layout: The decode cache retains the original low-dimensional latent state $c_t$ and adds a small scalar cache $A^k_{cache}$ containing the inverse-RMS values. It does not store the full normalized keys.
Blockwise Normalization: To maintain efficiency, the method normalizes the "content" and "RoPE" components separately. The content path uses the absorbed latent formulation, while the RoPE path uses standard materialized vectors. This avoids synchronization points that would complicate fused kernel implementations.

During decoding, the attention logits are computed as:
$\ell_{i,t,h} = (\tilde{q}_{i,h} \cdot c_t) \cdot \alpha^q_{i,h} \cdot \alpha^k_{t,g(h)}$
where $\tilde{q}$ includes the absorbed weights, and the scalar factors are applied as a row-column scaling step before the softmax.

Key Contributions

Exact Formulation: The paper derives an algebraically exact formulation of post-up-projection QK RMSNorm that is compatible with MLA's latent decode path.
Efficient Cache Design: It proposes a cache layout that preserves the latent KV cache as the dominant state, adding only one scalar key RMS statistic per token and KV group (e.g., increasing cache size by ~1.56% in tested configurations).
Empirical Validation: The authors demonstrate that this approach achieves lower training loss and better downstream accuracy compared to QK clipping, while introducing negligible latency overhead.

Results
Experiments were conducted on 400M-parameter MLA models trained for up to 100B tokens on H800 GPUs:

Training Stability: QK-Normed MLA achieved consistently lower training loss than QK-Clip throughout the 100B-token run. In high-learning-rate stress tests ( $2 \times 10^{-2}$ ), QK-Clip diverged with gradient NaNs, while QK-Normed MLA remained stable.
Downstream Performance: On a suite of downstream tasks, QK-Normed MLA improved accuracy on 7 out of 8 tasks, raising the average score from 44.75 to 46.33, and reduced LAMBADA perplexity from 16.28 to 14.18.
Decode Efficiency: Benchmarks on contexts ranging from 4k to 256k tokens showed that the method adds less than 2% latency overhead (average 1.40%). The scalar cache increased the total cache size by only ~2.8% per layer.

Significance and Claims
The paper claims that the apparent incompatibility between QK normalization and MLA is an implementation artifact rather than an architectural constraint. By showing that the static weights can be absorbed and the dynamic statistics reduced to scalars, the authors enable QK normalization as a practical stabilization mechanism for MLA models without sacrificing the memory efficiency of latent decoding.

The authors remain modest regarding the scope of their claims:

The results are based on 400M-parameter models; they do not claim to establish frontier-scale behavior for billion-parameter models, though the algebraic equivalence holds regardless of scale.
The logit-scale analysis is presented as a mechanism-level diagnostic rather than a complete optimization theory.
The practical speed depends on kernel integration (fusing scalar factors into the attention path) to avoid unnecessary memory traffic.

Technical Summary: QK-Normed MLA

More like this