Multi-Head Latent Attention: Low-Rank KV Compression, Decoupled RoPE, and Matrix Absorption

Multi-Head Latent Attention (MLA), introduced in DeepSeek-V2, addresses the KV cache bottleneck that has constrained long-context LLM inference. Rather than reducing the number of heads as in Grouped-Query Attention (GQA) or Multi-Query Attention (MQA), MLA compresses keys and values into a shared low-rank latent representation, caches that compressed form, and reconstructs full K and V matrices at use time. A decoupled rotary position embedding (RoPE) preserves positional information without preventing compression, and an "absorb" transformation folds the up-projection and output-projection matrices into the query weights so inference never materializes the decompressed K and V. This article walks through the mathematics of standard MHA, the MLA modifications, and the resulting inference economics.

Standard Multi-Head Attention and the KV Cache Problem

In standard Multi-Head Attention (MHA), a hidden state $h_t \in \mathbb{R}^d$ at position $t$ is projected to queries, keys, and values for each of $H$ heads:

$$q_{t,i} = W_i^Q h_t, \quad k_{t,i} = W_i^K h_t, \quad v_{t,i} = W_i^V h_t$$

where $W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d_h \times d}$ and $d_h = d / H$ is the head dimension. The attention score for head $i$ at step $t$ against all prior positions $s \le t$ is

$$\alpha_{t,s,i} = \frac{q_{t,i}^\top k_{s,i}}{\sqrt{d_h}}$$

and the head output is $o_{t,i} = \sum_{s \le t} \text{softmax}_s(\alpha_{t,s,i}) v_{s,i}$. The concatenated output $o_t = [o_{t,1}; \dots; o_{t,H}]$ is then projected by $W^O \in \mathbb{R}^{d \times d}$.

During autoregressive generation, each layer must store $k_{s,i}$ and $v_{s,i}$ for all $s \le t$ and all heads $i$. For a model with $L$ layers, $H$ heads, head dimension $d_h$, and context length $T$, the KV cache size is $2 L H d_h T$ floats — $2 L d T$ in total. DeepSeek-V2 ($L=60, H=128, d_h=128, T=128\text{K}$) would require roughly 488 GB of KV cache under MHA, making long-context serving impractical.

MLA: Low-Rank Key-Value Joint Compression

MLA replaces the separate $W_i^K$ and $W_i^V$ projections with a single down-projection to a shared latent vector $c_t^{KV} \in \mathbb{R}^{d_c}$ (where $d_c \ll H d_h$), followed by up-projections per head:

$$c_t^{KV} = W^{DKV} h_t, \quad W^{DKV} \in \mathbb{R}^{d_c \times d}$$

$$k_{t,i} = W_i^{UK} c_t^{KV}, \quad v_{t,i} = W_i^{UV} c_t^{KV}, \quad W_i^{UK}, W_i^{UV} \in \mathbb{R}^{d_h \times d_c}$$

Only $c_t^{KV}$ is stored in the cache — a single vector of size $d_c$ per token per layer. For DeepSeek-V2, $d_c = 512$ versus $H d_h = 16,384$ for full K/V, yielding a 32× reduction in cached elements per token (from 32,768 to 1,024 floats per layer per token). The authors report a 93.3% KV cache reduction compared to their 67B dense MHA baseline.

Queries are handled symmetrically with a separate down-projection $W^{DQ} \in \mathbb{R}^{d_c' \times d}$ and per-head up-projections $W_i^{UQ} \in \mathbb{R}^{d_h \times d_c'}$, though the query latent $c_t^Q$ is not cached since queries are only needed at the current step.

Decoupled Rotary Position Embedding

Standard RoPE applies a position-dependent rotation to queries and keys:

$$q_{t,i} \leftarrow \text{RoPE}(q_{t,i}, t), \quad k_{t,i} \leftarrow \text{RoPE}(k_{t,i}, t)$$

In MLA, if RoPE were applied after the up-projection, the rotated $k_{t,i}$ would depend on the full $W_i^{UK} c_t^{KV}$, breaking the ability to cache only $c_t^{KV}$. MLA therefore applies RoPE to a dedicated decoupled set of query and key projections that bypass the compression:

$$q_{t,i}^R = W_i^{QR} h_t, \quad k_t^{KR} = W^{KV_R} h_t$$

where $W_i^{QR} \in \mathbb{R}^{d_R \times d}$ and $W^{KV_R} \in \mathbb{R}^{d_R \times d}$ with $d_R$ typically much smaller than $d_h$ (DeepSeek-V2 uses $d_R=64$). The RoPE rotation is applied to these low-dimensional vectors, and the positional attention score component is computed separately:

$$\beta_{t,s,i} = \frac{(q_{t,i}^R)^\top \text{RoPE}(k_s^{KR}, s-t)}{\sqrt{d_R}}$$

The full attention score combines content and positional components:

$$\alpha_{t,s,i} = \frac{q_{t,i}^\top k_{s,i}}{\sqrt{d_h}} + \frac{(q_{t,i}^R)^\top \text{RoPE}(k_s^{KR}, s-t)}{\sqrt{d_R}}$$

Because $k_s^{KR}$ is projected directly from $h_s$ (not from the compressed $c_s^{KV}$), it can be RoPE-rotated without requiring decompression. Only $c_s^{KV}$ and $k_s^{KR}$ are cached, adding $d_c + d_R = 576$ floats per token per layer — still a 57× reduction versus full MHA.

Matrix Absorption: Eliminating Decompression at Inference

The content attention score can be rewritten by substituting the up-projections:

$$q_{t,i}^\top k_{s,i} = (W_i^{UQ} c_t^Q)^\top (W_i^{UK} c_s^{KV}) = (c_t^Q)^\top (W_i^{UQ \top} W_i^{UK}) c_s^{KV}$$

Define $\hat{W}_i^{UK} = W_i^{UQ \top} W_i^{UK} \in \mathbb{R}^{d_c \times d_c}$. Then

$$\alpha_{t,s,i}^{\text{content}} = \frac{(c_t^Q)^\top \hat{W}_i^{UK} c_s^{KV}}{\sqrt{d_h}}$$

During inference, $c_t^Q$ is computed once per step. The matrix $\hat{W}_i^{UK}$ can be pre-computed and fused into the query projection:

$$\tilde{W}_i^Q = \hat{W}_i^{UK \top} W_i^{UQ \top} \in \mathbb{R}^{d_c \times d_c'}$$

so that the score becomes $(c_t^Q)^\top \tilde{W}_i^Q c_s^{KV} / \sqrt{d_h}$ — a single matrix-vector product against the cached $c_s^{KV}$. The full K matrix is never materialized.

Similarly, the output aggregation can absorb the value up-projection and output projection. The head output is

$$o_{t,i} = \sum_s \text{softmax}_s(\alpha_{t,s,i}) v_{s,i} = \sum_s \text{softmax}_s(\alpha_{t,s,i}) W_i^{UV} c_s^{KV}$$

The concatenated output $o_t = [o_{t,1}; \dots; o_{t,H}]$ is multiplied by $W^O$. Defining $W_i^{OUV} = W^O_{:, i} W_i^{UV} \in \mathbb{R}^{d \times d_c}$ (where $W^O_{:, i}$ selects the columns corresponding to head $i$), the final output is

$$y_t = \sum_i \sum_s \text{softmax}_s(\alpha_{t,s,i}) W_i^{OUV} c_s^{KV}$$

All up-projections $W_i^{UK}, W_i^{UV}$ and the output projection $W^O$ are absorbed into modified query and output matrices. At inference time, the model only reads the compressed latents $c_s^{KV}$ and $k_s^{KR}$ from cache — no decompression step occurs.

Inference Economics

| Configuration | KV Cache per Token per Layer | 128K Context (60 layers) | |---------------|------------------------------|---------------------------| | MHA (DeepSeek 67B) | $2 H d_h = 32,768$ floats | ~488 GB | | GQA (8 groups) | $2 \times 8 \times 128 = 2,048$ floats | ~30 GB | | MLA (DeepSeek-V2) | $d_c + d_R = 576$ floats | ~8.6 GB |

MLA achieves a 57× reduction versus MHA and a 3.5× reduction versus 8-group GQA. The compute cost of the absorb operations is negligible: the pre-computed $\tilde{W}_i^Q$ and $W_i^{OUV}$ add one small matrix multiply per head per layer per token, which is dominated by the attention softmax and the linear projections that must run regardless.

Training Considerations

During training, MLA uses the same forward pass but caches the full $k_{t,i}$ and $v_{t,i}$ for gradient computation, since the absorb transformation is not applied until inference. The decoupled RoPE projections add a small parameter overhead (~0.5% of attention params). DeepSeek-V2 reports that MLA matches or exceeds MHA quality while enabling 128K context at a fraction of the inference memory.

Sources

• DeepSeek-AI, "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model," arXiv:2405.04434, §2.1.2, Appendix D.2.
• Lior Sinai, "DeepSeek's Multi-Head Latent Attention," https://liorsinai.github.io/machine-learning/2025/02/22/mla.html
• Chris McCormick, "The Inner Workings of Multihead Latent Attention (MLA)," https://mccormickml.com/2025/04/26/inner-workings-of-mla
• Sebastian Raschka, "Multi-Head Latent Attention (MLA)," https://sebastianraschka.com/llms-from-scratch/ch04/05_mla
• Vizuara, "Decoding Multi-Head Latent Attention (Part 1): The KV Cache Compression," https://vizuara.substack.com/p/decoding-multi-head-latent-attention