Multi-head Latent Attention (MLA)

Summary: MLA replaces standard MHA's per-head K,V projections with low-rank latent projections shared across heads. Each head attends to a compressed latent KV, reducing KV cache by 8× (DeepSeek-V3) with <1% perplexity degradation. Forms the attention backbone for DeepSeek-V3/V4.

Core Idea

Standard MHA (Reference)

For each head h:
  Q_h = x W_Q^h      [d_model × d_head]
  K_h = x W_K^h      [d_model × d_head]
  V_h = x W_V^h      [d_model × d_head]
  Attention(Q_h, K_h, V_h)
  
KV Cache per token: 2 × n_heads × d_head = 2 × d_model  (per layer)

MLA (Latent Attention)

Shared low-rank projections:
  K_latent = x W_K^latent  [d_model × d_latent]  (d_latent << d_model)
  V_latent = x W_V^latent  [d_model × d_latent]

Per head:
  Q_h = x W_Q^h           [d_model × d_head]
  K_h = K_latent W_K^h    [d_latent × d_head]
  V_h = V_latent W_V^h    [d_latent × d_head]
  Attention(Q_h, K_h, V_h)

KV Cache per token: 2 × d_latent  (per layer)  ← d_latent = d_model / 8

Mathematical Formulation

Latent Projection

K_{latent} = X W_K^{latent} ∈ R^{n × d_{latent}}

V_{latent} = X W_V^{latent} ∈ R^{n × d_{latent}}

Head-Specific Expansion

K_h = K_{latent} W_K^h ∈ R^{n × d_{head}}

V_h = V_{latent} W_V^h ∈ R^{n × d_{head}}

Q_h = X W_Q^h ∈ R^{n × d_{head}}

Attention per Head

Attention_h = softmax((Q_h K_h^T)/(sqrt(d_{head))}) V_h

Compression Ratio

Dimension Standard MHA MLA Ratio
d_{model} 4096 4096 1×
n_{heads} 32 32 1×
d_{head} 128 128 1×
d_{latent} N/A 512 —
KV/token/layer 8KB 1KB 8× smaller

d_{latent} = 512 chosen as sweet spot (DeepSeek-V3: 128 heads × 128 head_dim = 16K → 512 latent = 8×)

Key Properties

Property MLA Standard MHA
KV Cache O(d_{latent}) O(d_{model})
Training FLOPs Similar (extra matmul) Baseline
Inference Latency Lower (less memory BW) Baseline
Quality ~0.5-1% perplexity ↑ Baseline
RoPE Compatibility Requires modification Native

RoPE with MLA (Critical Detail)

Standard RoPE applies rotation to Q, K per head. With shared latent K:

  1. Option A: Apply RoPE to latent K_{latent} → all heads share same rotation
  2. Option B (DeepSeek): Apply RoPE to expanded K_h per head
    • Pre-rotate W_K^h with position-dependent complex phases
    • Equivalent: K_h = K_{latent} (W_K^h ⊗ R_theta)
# Simplified MLA with RoPE (DeepSeek-V3 style)
def mla_attention(x, pos):
    # Shared latent projections
    k_latent = x @ W_K_latent  # [n, d_latent]
    v_latent = x @ W_V_latent  # [n, d_latent]
    
    # Per-head Q
    q = x @ W_Q  # [n, n_heads, d_head]
    
    # Expand latent to heads with RoPE
    k = k_latent @ W_K_expand  # [n, n_heads, d_head]
    v = v_latent @ W_V_expand  # [n, n_heads, d_head]
    
    # Apply RoPE to Q and K per head
    q = apply_rope(q, pos)
    k = apply_rope(k, pos)
    
    return flash_attention(q, k, v)

DeepSeek-V3/V4 Integration

V3: MLA as Base Attention

  • All 61 layers use MLA
  • d_{latent} = 512, n_{heads} = 128, d_{head} = 128
  • Enables 128K context on 8×H800 (KV cache ~15GB vs 120GB)

V4: MLA + CSA + HCA Hybrid

Layer 1-20:   MLA (full attention)
Layer 21-40:  CSA (compressed sparse) + MLA
Layer 41-61:  HCA (heavily compressed) + MLA
  • MLA provides dense attention anchor
  • CSA/HCA provide sparse long-range
  • Combined: 1M context feasible

Comparison with Related

Method Compression Mechanism Quality Context
MLA 8× Low-rank latent KV High 128K
GQA 4-8× Grouped K,V heads Medium 32K
MQA 32× Single K,V head Low 32K
Linear Attn ∞ (O(1)) Kernel approx Low-Med ∞
Sparse (DSA) 10-100× Learned sparsity High 1M+

Implementation Considerations

Memory Savings (DeepSeek-V3, 67B params)

Component Standard MHA MLA Savings
KV Cache (128K ctx) 120 GB 15 GB 105 GB
Activation Memory Baseline +5% (latent) Minor ↑

Training Overhead

  • Extra matmul: X W_K^{latent} + K_{latent} W_K^h per layer
  • ~3-5% additional FLOPs
  • Offset by longer context / larger batch

Inference Optimization

  • Fuse latent projection + head expansion
  • Cache K_{latent}, V_{latent} instead of per-head
  • Compatible with FlashAttention-2/3

Limitations

  1. RoPE coupling: Can't use standard RoPE; requires custom implementation
  2. Head independence reduced: Latent sharing couples heads slightly
  3. Not ideal for short context: Overhead not worth it < 4K tokens
  4. Quantization sensitivity: Low-rank projections sensitive to INT4/NF4

Related Concepts

Sources

  • arxiv-2512.02556: DeepSeek-V3 paper, §2.1 — MLA introduction
  • arxiv-2606.19348: DeepSeek-V4 paper, §2.1 — MLA in hybrid
  • deepseek-dsa-paper: DSA paper — MLA as efficient base