Per-Head Muon
Summary: Per-Head Muon extends the Muon optimizer (momentum + orthogonalization preconditioning) by maintaining separate optimizer state per attention head, enabling adaptive learning rates and preconditioning tailored to each head's curvature at extreme scale.
Overview
Standard Muon applies the same preconditioning (Newton-Schulz orthogonalization) to all parameters. At 2.8T scale with hundreds of attention heads, different heads learn different features at different rates. Per-Head Muon gives each head its own optimizer state.
Key Components
Base Muon Recap
- Momentum-based optimizer with orthogonalization preconditioning
- For weight matrix W: maintain momentum M, update via orthogonalized gradient
- Newton-Schulz iteration:
G_orth = G @ (3I - G.T @ G) @ (3I - G @ G.T) / 2(approx) - Benefits: better conditioning, scale invariance, works well for transformers
Per-Head Extension
- State separation: Each attention head h ∈ {0..H-1} gets independent momentum buffer M_h and preconditioner state
- Head-specific learning: Heads attending to different positions (local vs global) or features (syntax vs semantics) adapt independently
- Memory overhead: O(H * d_head²) — manageable since d_head = 64-128 typically
Integration with KDA
- Kimi Delta Attention has hybrid linear + standard attention heads
- Per-Head Muon handles different head types with different dynamics
- Critical for stable training when heads have heterogeneous gradient statistics
Mathematical Formulation
In plain English: Standard Muon computes one preconditioned update for the whole weight matrix. Per-Head Muon splits the weight matrix by head dimension, computes independent preconditioned updates for each head's slice, then stitches back together.
Let W ∈ ℝ^{d_model × d_model} be attention output projection, split into H heads:
W = [W_0, W_1, ..., W_{H-1}] where each W_h ∈ ℝ^{d_model × d_head}
Standard Muon:
M = β M + (1-β) G
W = W - lr * Orthogonalize(M)
Per-Head Muon:
For each head h:
M_h = β M_h + (1-β) G_h
W_h = W_h - lr * Orthogonalize(M_h)
(Similarly for Q, K, V projections split by head)
Variants / Extensions
- Per-Layer Per-Head: Separate state per layer AND head (more memory, more adaptivity)
- Shared Preconditioner, Per-Head LR: Shared Newton-Schulz, but per-head learning rate scaling
- Quantized Per-Head State: MXFP4 momentum buffers for memory efficiency
Applications
- Stable LatentMoE training at 2.8T (Kimi K3)
- Ultra-deep transformers (100+ layers, heterogeneous head roles)
- Mixture-of-Depths / conditional computation
- Models with mixed attention types (local + global + linear)
Historical Context
| Year | Work | Contribution |
|---|---|---|
| 2020 | AdamW | Standard for transformers |
| 2022 | Muon (Liu et al.) | Orthogonalization preconditioning for transformers |
| 2023 | Shampoo / SOAP | Full-matrix preconditioning |
| 2024 | Muon variants | Distributed, quantized, fused |
| 2026 | Per-Head Muon (Kimi K3) | Head-wise state for heterogeneous head dynamics at 3T scale |
Related Concepts
- stable-latentmoe
- quantile-balancing
- muon-optimizer
- kimi-delta-attention
- attention-residuals
- kimi-k3
- optimizer-preconditioning
Sources
- kimi-k3-doc: Introduction — Per-Head Muon extends Muon by optimizing attention heads independently for adaptive learning at scale
- kimi-k3-tech-blog: Architecture and Infrastructure — Per-Head Muon extends Muon by optimizing attention heads independently for more adaptive learning at scale