Attention Internals advanced 8 min read 5 flashcards

Multi-Head Latent Attention

DeepSeek's answer to the KV cache problem — cache one low-rank latent vector per token instead of every head's keys and values, absorb the up-projection into the query and output weights, and carve out a separate untouched channel for RoPE.

Grouped-query attention shrinks the KV cache by giving several query heads one shared key-value head, which is a blunt instrument: you choose a group count, you accept the quality loss it causes, and you are done (see multi-query and grouped-query attention). Multi-head latent attention asks a different question. Rather than reducing how many key-value heads exist, it changes what gets stored: keep full multi-head expressiveness, but cache a single compressed latent vector per token and reconstruct the heads on the fly. DeepSeek-V2 reports a 93.3% KV cache reduction against DeepSeek 67B alongside 42.5% lower training cost and 5.76x higher maximum generation throughput (DeepSeek-AI, 2024, DeepSeek-V2, arXiv:2405.04434).

Low-rank joint compression

For hidden state \(h_t\), MLA projects down to a joint latent of dimension \(d_c \ll d_h n_h\) and back up separately for keys and values:

\[ c_t^{KV} = W^{DKV} h_t, \qquad k_t^{C} = W^{UK} c_t^{KV}, \qquad v_t^{C} = W^{UV} c_t^{KV} \]

Only \(c_t^{KV}\) is cached. Everything else is recomputed from it. This is a rank constraint on the key-value subspace, and it is a bet that the information genuinely needed from past positions lives in a much smaller subspace than \(n_h \times d_h\) dimensions suggest.

The step that makes it fast rather than merely small is weight absorption. Since \(q_t^\top k_s = (W^{UQ} c_t^{Q})^\top (W^{UK} c_s^{KV}) = c_t^{Q\top} (W^{UQ\top} W^{UK}) c_s^{KV}\), the up-projection \(W^{UK}\) can be folded into the query projection at load time and never materialised at inference. The same trick folds \(W^{UV}\) into the output projection. At decode time you never reconstruct per-head keys and values at all; you attend directly in latent space.

The RoPE problem, and the decoupled fix

Weight absorption depends on \(W^{UK}\) being position-independent so it can be multiplied into \(W^{UQ}\) once. Rotary position embeddings break exactly this, because RoPE applies a position-dependent rotation \(R_s\) between the two projections, and \(W^{UQ\top} R_{t-s} W^{UK}\) cannot be precomputed for every relative offset (see rotary position embeddings).

MLA's answer is to split the head into two channels. A compressed channel carries no positional rotation and enjoys absorption; a small decoupled RoPE channel of extra dimensions carries the rotation and is cached separately, shared across heads. Query and key are the concatenation of the two, so the attention score is the sum of a content term computed in latent space and a positional term computed the ordinary way. The cache is therefore \(d_c\) latent dimensions plus a handful of RoPE dimensions per token, not \(2 n_h d_h\).

When it breaks

  • It is an architecture, not a retrofit. MLA changes the shape of the attention block, so an existing MHA or GQA checkpoint cannot be converted without retraining or a nontrivial distillation procedure. GQA can sometimes be approximated by mean-pooling existing KV heads; MLA cannot.
  • Kernel support is the practical bottleneck. Absorbed-weight decoding is a different memory access pattern from standard attention, so FlashAttention-style kernels do not apply unchanged, and early MLA deployments ran slower than their FLOP counts predicted until dedicated kernels landed. Verify your serving stack implements the absorbed path rather than materialising keys and values per step.
  • The latent dimension is a quality knob with a cliff. Too small a \(d_c\) and the rank constraint starts discarding information the model needed; the loss shows up unevenly across tasks, typically hitting long-context retrieval first.
  • It shifts the bottleneck rather than removing it. With the cache 93% smaller, decode stops being KV-bandwidth bound and starts being bound by something else — weight loading, or expert routing in an MoE model. The end-to-end speedup is always smaller than the cache reduction, and 5.76x throughput against a 93.3% cache cut is the honest shape of that gap.
  • The published gains are vendor-measured, against a specific baseline. They compare DeepSeek-V2 to DeepSeek 67B, which differs in more than attention. Read them as evidence the direction works, not as a portable multiplier.
Check yourself

5 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track