Gated Multi-head Latent Attention (Gated MLA)

Summary: Gated MLA extends Multi-head Latent Attention (MLA) by adding per-head gating mechanisms that improve attention selectivity — controlling which heads participate in which expert routing decisions. Combined with SiTU activation, it enables stable attention-expert coordination at 2.8T scale.

Overview

Standard MLA (Multi-head Latent Attention, from DeepSeek-V2/V3) compresses KV cache by projecting keys/values to a latent dimension. Gated MLA adds a gating layer per head to modulate head participation based on routing context, improving alignment between sparse MoE routing and attention computation.

Key Components

Base MLA Recap

  • Compress K, V to latent dimension d_c << d_model
  • K_latent = K @ W_K, V_latent = V @ W_V where W_K, W_V ∈ ℝ^{d_model × d_c}
  • Attention: Attn(Q, K_latent, V_latent) with Q projected to d_c
  • Benefits: reduced KV cache, faster decoding

Gating Mechanism

  • Per-head gate: Each head h has gate g_h = σ(W_g @ x + b_g) where x is input/context
  • Selective participation: Head output scaled by gate: out_h = g_h * Attention_h(Q, K_latent, V_latent)
  • Routing alignment: Gates can be conditioned on expert routing decisions — heads relevant to active experts get higher gates

Integration with Stable LatentMoE

  • 16 active experts per token, each with potentially different attention needs
  • Gated MLA allows attention heads to specialize for different expert subsets
  • Reduces interference between expert-specific attention patterns

Mathematical Formulation

In plain English: Standard MLA computes attention for all heads equally. Gated MLA asks "which heads are useful for this token's active experts?" and scales head outputs accordingly — like a soft attention head selection.

Standard MLA:

# Q: [B, L, H, d_head], K, V: [B, L, d_c]
Q_c = Q @ W_Qc  # [B, L, H, d_c]
Attn = softmax(Q_c @ K_c.transpose(-2,-1) / sqrt(d_c)) @ V_c  # [B, L, H, d_c]
Out = Attn @ W_O  # [B, L, H, d_head] -> merge heads

Gated MLA:

# Additional gating
gate_logits = x @ W_g + b_g  # [B, L, H] - x can include routing info
gates = sigmoid(gate_logits)  # [B, L, H]

Q_c = Q @ W_Qc
Attn_h = softmax(Q_c @ K_c.transpose(-2,-1) / sqrt(d_c)) @ V_c  # per head
Attn_gated = gates[..., None] * Attn_h  # [B, L, H, d_c]
Out = merge_heads(Attn_gated) @ W_O

Variants / Extensions

  • Hard Gating: Top-k heads per token (discrete, requires straight-through estimator)
  • Expert-Conditioned Gates: Gate network takes expert routing weights as input
  • Quantized Gates: Binary/ternary gates for inference efficiency

Applications

  • Stable LatentMoE (Kimi K3: 896 experts, 16 active, hybrid KDA attention)
  • Sparse MoE with heterogeneous attention requirements
  • Long-context models where different heads handle different context regions
  • Mixture-of-Depths / conditional attention computation

Historical Context

Year Work Contribution
2017 Multi-head Attention Parallel attention heads
2023 DeepSeek-V2 MLA Latent KV compression
2024 DeepSeek-V3 MLA + fine-grained MoE
2024 Various Head pruning, head selection
2026 Gated MLA (Kimi K3) Per-head gating for MoE-attention alignment at 3T scale

Related Concepts

Sources

  • kimi-k3-doc: Introduction — Gated MLA improves attention selectivity
  • kimi-k3-tech-blog: Architecture and Infrastructure — SiTU and Gated MLA improve activation control and attention selectivity; enable stable 2.8T training