Multi-Head Latent Attention: Mathematical Foundations, Low-Rank KV Compression, and Decoupled RoPE in Transformer Architectures

Multi-Head Attention (MHA) has served as the core sequence-mixing primitive in autoregressive Transformer architectures since the introduction of the Transformer by Vaswani et al. (2017). In production serving environments, autoregressive generation requires caching key and value projections for all previous tokens in high-bandwidth GPU memory (HBM). As sequence lengths expand toward 128k tokens and beyond, this Key-Value (KV) cache grows linearly with context length, batch size, and layer count, rapidly becoming the primary memory and throughput bottleneck during inference.

Architectural variants such as Multi-Query Attention (MQA) proposed by Shazeer (2019) and Grouped-Query Attention (GQA) introduced by Ainslie et al. (2023) mitigate KV cache growth by forcing multiple query heads to share a reduced number of key-value head pairs. While GQA reduces KV cache memory by factors of 4x to 8x, it prunes the representational parameter capacity of keys and values, introducing an empirical trade-off between memory footprint and model expressivity.

Multi-Head Latent Attention (MLA), introduced by DeepSeek-AI (2024) in the DeepSeek-V2 architecture and carried forward in DeepSeek-V3, addresses this dilemma through low-rank joint key-value compression combined with decoupled positional encodings. Instead of reducing the number of attention heads, MLA factorizes key and value spaces into a compact latent vector while preserving full multi-head expressive capacity. At inference time, linear algebraic properties allow projection matrices to be absorbed into query transformations and output projections, eliminating the need to materialize uncompressed keys and values in memory.

The KV Cache Memory Bottleneck in Standard Attention

In standard Multi-Head Attention with hidden dimension $d$, $n_h$ attention heads, and per-head dimension $d_h$, the model projects an incoming token hidden state $h_t \in \mathbb{R}^d$ into query, key, and value vectors for each head $i \in \{1, \dots, n_h\}$:

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

During autoregressive generation at sequence step $t$, the key and value vectors for all preceding positions $s \le t$ must remain resident in GPU memory to avoid recalculating past projections. The per-token memory consumption in the KV cache for a single transformer layer under 16-bit precision (FP16 or BF16, 2 bytes per element) is:

Memory_MHA = 2 * 2 * n_h * d_h bytes/token/layer

For a model with $n_h = 128$ and $d_h = 128$, storing uncompressed keys and values requires $2 \times 2 \times 128 \times 128 = 65,536$ bytes (64 KB) per token per layer. Across a 60-layer foundation model serving a context of 128,000 tokens for a single user request, the KV cache alone demands approximately 491.5 GB of GPU memory, requiring multi-node tensor or pipeline parallelism merely to store activation state.

Grouped-Query Attention (GQA) reduces this memory burden by grouping query heads into $g$ groups and allocating only one key head and one value head per group ($n_{kv} = g$). For $g = 8$, the cache footprint scales down to $2 \times 2 \times 8 \times 128 = 4,096$ bytes per token per layer. However, sharing keys and values across multiple query heads reduces the model capacity to learn distinct subspace interactions between queries and tokens.

Low-Rank Key-Value Joint Compression

MLA replaces independent high-dimensional key and value projections with a low-rank joint compression mechanism. The input hidden state $h_t \in \mathbb{R}^d$ is mapped to a shared low-dimensional latent vector $c_t^{KV} \in \mathbb{R}^{d_c}$ via a down-projection matrix $W^{DKV} \in \mathbb{R}^{d_c \times d}$:

c_t^{KV} = RMSNorm(W^{DKV} h_t)

The compression dimension $d_c$ is chosen such that $d_c \ll n_h d_h$. For example, DeepSeek-V2 and DeepSeek-V3 configure $d_c = 512$ while supporting $n_h = 128$ heads with $d_h = 128$ (where $n_h d_h = 16,384$).

From this shared latent representation $c_t^{KV}$, head-specific content key vectors $k_{t, i}^C \in \mathbb{R}^{d_h}$ and value vectors $v_{t, i}^C \in \mathbb{R}^{d_v}$ are generated through head-specific up-projection matrices $W_i^{UK} \in \mathbb{R}^{d_h \times d_c}$ and $W_i^{UV} \in \mathbb{R}^{d_v \times d_c}$:

k_{t, i}^C = W_i^{UK} c_t^{KV}
v_{t, i}^C = W_i^{UV} c_t^{KV}

To reduce training activation memory and parameter overhead, MLA also applies a low-rank decomposition to query projections:

c_t^Q = RMSNorm(W^{DQ} h_t)
q_{t, i}^C = W_i^{UQ} c_t^Q

where $W^{DQ} \in \mathbb{R}^{d_c' \times d}$ down-projects the query into latent dimension $d_c'$, and $W_i^{UQ} \in \mathbb{R}^{d_h \times d_c'}$ up-projects to individual query head subspaces. Because queries are computed only for active tokens during the current decoding step and are not retained across time steps, $c_t^Q$ does not occupy persistent KV cache memory during inference.

The RoPE Incompatibility Problem and Decoupled RoPE

Modern autoregressive language models rely on Rotary Position Embeddings (RoPE), formulated by Su et al. (2021), to inject relative positional information via complex rotation matrices applied directly to queries and keys:

q_{t, i}^{RoPE} = R_t q_{t, i}
k_{s, i}^{RoPE} = R_s k_{s, i}

where $R_t \in \mathbb{R}^{d_h \times d_h}$ is a block-diagonal orthogonal rotation matrix for sequence index $t$. The inner product between rotated vectors naturally encodes the relative distance $(s - t)$ because $R_t^\top R_s = R_{s - t}$.

A fundamental architectural tension arises when combining RoPE with low-rank key compression. If RoPE were applied directly to the up-projected keys $k_{s, i}^C = W_i^{UK} c_s^{KV}$, the attention score between query token $t$ and cached token $s$ would be:

Score_{t, s, i} = (R_t q_{t, i}^C)^\top (R_s W_i^{UK} c_s^{KV})

Because the rotation matrix $R_s$ depends on position $s$ and does not commute with the up-projection matrix $W_i^{UK}$ ($R_s W_i^{UK} \neq W_i^{UK} R_s$), the up-projection cannot be factored out or pre-computed. If keys were rotated in this manner, the inference runtime would be forced to compute and store the full uncompressed vectors $k_{s, i}^{RoPE} \in \mathbb{R}^{n_h \times d_h}$ across all $n_h$ heads for every cached token, entirely negating the memory benefits of the low-rank latent representation.

To resolve this incompatibility, MLA decouples positional encodings from content representations. The query and key vectors are divided into content sub-vectors and dedicated positional sub-vectors:

q_{t, i} = [q_{t, i}^C ; q_{t, i}^R]
k_{t, i} = [k_{t, i}^C ; k_t^R]

The positional key vector $k_t^R \in \mathbb{R}^{d_R}$ is projected directly from the hidden state $h_t$ via a shared projection matrix $W^{KR} \in \mathbb{R}^{d_R \times d}$ and rotated with standard RoPE:

k_t^R = RoPE(W^{KR} h_t)

Crucially, $k_t^R$ is shared across all attention heads rather than being replicated per head. The positional query vector $q_{t, i}^R \in \mathbb{R}^{d_R}$ is projected per head and rotated accordingly:

q_{t, i}^R = RoPE(W_i^{QR} c_t^Q)

The attention score between query $t$ and cached token $s$ in head $i$ decomposes into the sum of a content inner product and a positional inner product:

q_{t, i}^\top k_{s, i} = (q_{t, i}^C)^\top k_{s, i}^C + (q_{t, i}^R)^\top k_s^R
                       = (q_{t, i}^C)^\top W_i^{UK} c_s^{KV} + (q_{t, i}^R)^\top k_s^R

Because content keys $k_{s, i}^C$ are free of position-dependent rotations, their projection matrices remain static across all token positions.

Weight Absorption: Zero Decompression at Inference

The linearity of the un-rotated content transformations enables weight absorption during inference serving, eliminating the requirement to materialize or uncompress $k_{s, i}^C$ or $v_{s, i}^C$ in memory.

Key Matrix Absorption into Queries

During inference, rather than multiplying the cached vector $c_s^{KV}$ by $W_i^{UK}$ for every past token $s$, the serving engine transposes and absorbs $W_i^{UK}$ directly into the current query vector $q_{t, i}^C$:

q'_{t, i} = (q_{t, i}^C)^\top W_i^{UK} \in \mathbb{R}^{1 \times d_c}

The content attention score is then computed via a single vector-matrix multiplication directly against the cached latent states:

Score_{t, s, i}^C = q'_{t, i} c_s^{KV}

This reduces the dot-product computation from head dimension $d_h$ across $n_h$ individual cached vectors to a single dot product in the latent space $d_c$.

Value Matrix Absorption into Output Projections

A complementary absorption occurs for value vectors during attention aggregation. The attention output for head $i$ is defined as the weighted sum over value vectors:

o_{t, i} = \sum_s A_{t, s, i} v_{s, i}^C = \sum_s A_{t, s, i} (W_i^{UV} c_s^{KV})

Because matrix multiplication distributes over addition, the static matrix $W_i^{UV}$ factors outside the summation:

o_{t, i} = W_i^{UV} \left( \sum_s A_{t, s, i} c_s^{KV} \right) = W_i^{UV} u_{t, i}

where $u_{t, i} = \sum_s A_{t, s, i} c_s^{KV} \in \mathbb{R}^{d_c}$ is the attention-weighted sum of raw cached latent vectors.

When the multi-head outputs are concatenated and multiplied by the standard attention output projection matrix $W^O \in \mathbb{R}^{d \times (n_h d_v)}$ (partitioned into per-head blocks $W_i^O \in \mathbb{R}^{d \times d_v}$), the final layer output is:

O_t = \sum_{i=1}^{n_h} W_i^O o_{t, i} = \sum_{i=1}^{n_h} W_i^O (W_i^{UV} u_{t, i}) = \sum_{i=1}^{n_h} (W_i^O W_i^{UV}) u_{t, i}

By pre-computing the combined output projection matrix:

W'_i^O = W_i^O W_i^{UV} \in \mathbb{R}^{d \times d_c}

the inference engine aggregates attention directly over the compressed latent KV vectors and projects directly to the model hidden state $d$, completely bypassing value decompression.

KV Cache Footprint Comparison

Through weight absorption, the inference KV cache needs to retain only two tensors per token per layer: the compressed content latent vector $c_t^{KV} \in \mathbb{R}^{d_c}$ and the decoupled shared positional key $k_t^R \in \mathbb{R}^{d_R}$.

Consider the concrete architecture dimensions of DeepSeek-V2 and DeepSeek-V3 ($n_h = 128$, $d_h = 128$, $d_c = 512$, $d_R = 64$):

• Standard Multi-Head Attention (MHA):
• Caches $2 \times n_h \times d_h = 2 \times 128 \times 128 = 32,768$ elements per token per layer.
• At 16-bit precision: 65,536 bytes/token/layer.
• Grouped-Query Attention (GQA, 8 groups):
• Caches $2 \times n_{kv} \times d_h = 2 \times 8 \times 128 = 2,048$ elements per token per layer.
• At 16-bit precision: 4,096 bytes/token/layer.
• Multi-Head Latent Attention (MLA):
• Caches $d_c + d_R = 512 + 64 = 576$ elements per token per layer.
• At 16-bit precision: 1,152 bytes/token/layer.

MLA achieves a 98.24% reduction in per-token KV cache memory relative to full MHA and a 71.88% reduction relative to 8-group GQA. Unlike GQA, which restricts the model to 8 effective key-value head representations, MLA maintains 128 distinct query heads attending to rank-compressed representations with separate up-projection pathways.

Empirical evaluations reported by DeepSeek-AI (2024) and verified in independent architectural studies such as MHA2MLA by Wang et al. (2025) demonstrate that MLA matches or exceeds the modeling perplexity and downstream benchmark accuracy of standard Multi-Head Attention while operating at a fraction of the memory footprint.

Sources

• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (DeepSeek-AI, 2024)
• DeepSeek-V3 Technical Report (DeepSeek-AI, 2024)
• Attention Is All You Need (Vaswani et al., 2017)
• Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
• RoFormer: Enhanced Transformer with Rotary Position Embedding (Su et al., 2021)
• Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs (Wang et al., 2025)