Muon Optimizer

Summary: Muon (Momentum Orthogonalization Update) is an optimizer that combines momentum with Newton-Schulz iteration to orthogonalize update matrices. Used for most parameters in DeepSeek-V4 training, it achieves faster convergence and improved stability vs AdamW. AdamW retained for embeddings, prediction heads, RMSNorm weights, and mHC static biases/gating factors.

Motivation

In plain English: Standard optimizers (Adam) treat gradients as element-wise. For matrix parameters (like attention projections), the structure of the update matters. Muon uses the gradient's singular vectors (via Newton-Schulz) to make updates that preserve orthogonality — leading to better conditioning and faster convergence, especially for large models with attention.

Algorithm (Algorithm 1, arxiv-2606.19348)

Input: Learning rate η, momentum μ, weight decay λ, rescaling γ

For each training step t:
  For each logically independent weight matrix W ∈ ℝ^{n×m}:
    G_t = ∇_W L_t(W_{t-1})              # Gradient
    M_t = μ M_{t-1} + G_t              # Momentum buffer
    
    # Nesterov trick + Hybrid Newton-Schulz
    O'_t = HybridNewtonSchulz(μ M_t + G_t)
    
    # Rescale update RMS
    O_t = O'_t · √max(n, m) · γ
    
    # Weight decay + update
    W_t = W_{t-1} · (1 - ηλ) - η O_t

Hybrid Newton-Schulz Orthogonalization

For matrix M, target orthogonalized form: UV^T where M = USigma V^T (SVD).

Standard Newton-Schulz iteration:

M_k = a M_{k-1} + b (M_{k-1} M_{k-1}^T) M_{k-1} + c (M_{k-1} M_{k-1}^T)^2 M_{k-1}

Muon's Hybrid (10 steps):

Steps Coefficients (a, b, c) Purpose
1-8 (3.4445, -4.7750, 2.0315) Rapid convergence, singular values → 1
9-10 (2, -1.5, 0.5) Stabilize exactly at 1

Pre-normalization: M_0 = M / ||M||_F to keep max singular value ≤ 1.

Configuration Details (§2.4)

Component Optimizer Reason
Embeddings AdamW Standard
Prediction head AdamW Standard
RMSNorm weights AdamW Standard
mHC static biases (S^pre, S^res, S^post) AdamW Small, sensitive
mHC gating factors (alpha^pre, alpha^res, alpha^post) AdamW Small, sensitive
All other params (attn, MoE, MTP, etc.) Muon Matrix params benefit from orthogonalization

Differences from muon_kimi:

  • ✅ Weight decay on Muon params
  • ✅ Nesterov momentum
  • ✅ RMS rescaling (reuses AdamW hyperparams)
  • ✅ Hybrid Newton-Schulz (vs standard)
  • ✅ No QK-Clip — RMSNorm on queries/KV prevents exploding logits

Why It Works

  1. Orthogonal updates: Newton-Schulz extracts UV^T — updates move along singular vectors, preserving spectrum
  2. Faster convergence: Empirically fewer steps to same loss
  3. Stability: Bounded updates (rescaling + orthogonalization) prevent gradient explosion
  4. Compatibility: Works alongside AdamW for non-matrix params

Usage in DeepSeek-V4

From §2.4 and §3 (post-training §37-38):

  • Pre-training: 32T+ tokens on V4-Flash, 33T on V4-Pro
  • Post-training: FP4 quantization-aware training for MoE experts + indexer QK path
  • Framework: Custom fused kernels, ZeRO hybrid strategy for Muon, tensor-level checkpointing

Key Claims with Sources

Claim Source Locator Confidence
Muon uses momentum + hybrid Newton-Schulz for orthogonal updates arxiv-2606.19348 §2.4, Alg 1 0.95
Hybrid N-S: 8 steps rapid convergence + 2 steps stabilization arxiv-2606.19348 §2.4 0.9
Coefficients: (3.4445, -4.7750, 2.0315) then (2, -1.5, 0.5) arxiv-2606.19348 §2.4 0.95
AdamW kept for embeddings, head, norms, mHC biases/gating arxiv-2606.19348 §2.4 0.95
Weight decay + Nesterov + RMS rescaling applied to Muon arxiv-2606.19348 §2.4 0.9
QK-Clip not needed due to RMSNorm on queries/KV arxiv-2606.19348 §2.4 0.9
Faster convergence & stability vs AdamW for matrix params arxiv-2606.19348 §2.4 intro 0.9

Related Concepts

Sources

  • arxiv-2606.19348: DeepSeek-V4 paper (§2.4, Algorithm 1)