DeepSeek-V2: MLA Explained

MLA is a novel attention mechanism designed for Transformer-based Large Language Models, aiming to boost inference efficiency by massively compressing how Key-Value (KV) cache is stored and accessed during generation. Conventional Multi-Head Attention (MHA) keeps a heavy cache of keys and values per token, which limits batch sizes and sequence lengths for LLMs. MLA replaces this with a joint, low-rank compression of keys and values (and optionally queries), dramatically reducing memory footprint while maintaining or improving performance.

MLA also introduces decoupled Rotary Position Embedding (RoPE) to address the incompatibility between standard RoPE and low-rank compression.

---

MLA Full Formulas#

Below are the complete step-by-step formulas of MLA, as defined in the DeepSeek-V2 paper (Appendix C):

1. Query Compression

• Compress the attention input for queries:

  $$c_{Q_t} = W_{DQ} h_t$$

  where $c_{Q_t}$ is the compressed latent vector for queries, $h_t$ is the input at time $t$, and $W_{DQ}$ is the down-projection matrix.

• Up-project to produce query candidates:

  $$[q^C_{t,1}; \ldots; q^C_{t,n_h}] = q^C_t = W_{UQ} c_{Q_t}$$

• Apply Rotary Position Embedding (RoPE) to decoupled queries:

  $$[q^R_{t,1}; \ldots; q^R_{t,n_h}] = q^R_t = \text{RoPE}(W_{QR} c_{Q_t})$$

• Concatenate compressed and RoPE queries for each head:

  $$q_{t,i} = [q^C_{t,i}; q^R_{t,i}]$$

2. Key-Value Joint Compression

• Compress the attention input for keys and values:

  $$c_{KV_t} = W_{DKV} h_t$$

• Up-project keys:

  $$[k^C_{t,1}; \ldots; k^C_{t,n_h}] = k^C_t = W_{UK} c_{KV_t}$$

• Apply RoPE to the decoupled key:

  $$k^R_t = \text{RoPE}(W_{KR} h_t)$$

• Concatenate compressed and RoPE keys for each head:

  $$k_{t,i} = [k^C_{t,i}; k^R_t]$$

• Up-project values:

  $$[v^C_{t,1}; \ldots; v^C_{t,n_h}] = v^C_t = W_{UV} c_{KV_t}$$

3. Attention Calculation

• For each head, compute attention output:

  $$o_{t,i} = \sum_{j=1}^{t} \text{Softmax}_j\left( \frac{q_{t,i}^T k_{j,i}}{\sqrt{d_h + d^R_h}} \right) v^C_{j,i}$$

  where $d_h$ is the per-head dimension, and $d^R_h$ is the decoupled per-head dimension for RoPE.

• Concatenate all heads and project:

  $$u_t = W_O [o_{t,1}; o_{t,2}; \ldots; o_{t,n_h}]$$

  where $W_O$ is the output projection matrix.

Caching Strategy

• During inference, only compressed latent vectors ($c_{KV_t}$, $k^R_t$) need to be cached, greatly reducing memory requirements compared to MHA.

Optimization Note

• Due to distributive properties in matrix multiplication, some projection matrices (e.g., $W_{UK}, W_{UV}$) can be merged with other projections during inference, avoiding redundant computation of keys and values for each query.

---

Summary Table of MLA vs Other Attention Mechanisms#

Where:

• $n_h$: # of attention heads
• $d_h$: per-head dimension
• $l$: # of layers
• $n_g$: # of groups (GQA)
• $d_c$: compression dimension (MLA)
• $d^R_h$: decoupled query/key dimension (MLA)

For DeepSeek-V2, typically, $d_c = 4 d_h$, $d^R_h = d_h / 2$, making MLA's cache much leaner.

Source: [DeepSeek-V2 Paper](https://arxiv.org/pdf/2405.04434