Weight Update & Optimizers: How Parameters Move
What actually happens in one update
Section titled “What actually happens in one update”Every training step, the model goes through a “checkup, then prescription” cycle:
- Forward pass: feed data, compute predictions and the loss (how wrong we are)
- Backward pass: use the chain rule to attribute the loss to every parameter (gradients)
- Optimizer update: decide, per parameter, which direction to move and how far
Step 3 is the star of this page. It sounds trivial — “step against the gradient” — but every devil lives in the details: how big is the step? Same step for every parameter? What if it oscillates? How do we smooth the noise? How do we save memory?
A brief history of optimizers: from “one size fits all” to “tailored”
Section titled “A brief history of optimizers: from “one size fits all” to “tailored””graph LR A[SGD<br/>one LR for all] --> B[Momentum<br/>smooth past gradients] B --> C[AdaGrad<br/>per-parameter LR] C --> D[Adam<br/>1st + 2nd moments] D --> E[AdamW<br/>decoupled weight decay] E --> F[Adafactor / 8-bit<br/>slim the states] F --> G[Muon<br/>orthogonalized momentum] style G fill:#f9f,stroke:#333Core intuition: gradients are noisy and biased — the direction from one batch is not the global direction. An optimizer’s entire job is to trade off “trust the current gradient” against “respect history”, while deciding how large each parameter’s step should be.
AdamW: the default answer for a decade
Section titled “AdamW: the default answer for a decade”Adam maintains two quantities:
- First moment m: an EMA of gradients (“where have we been going lately”)
- Second moment v: an EMA of squared gradients (“how noisy is this parameter”)
Update rule: parameter delta ≈ m / √v. Noisy parameters (large v) get shrunk steps; calm parameters get amplified steps. AdamW adds decoupled weight decay — multiply weights by (1 − lr·λ) each step as a regularizer.
The standard LLM configuration is remarkably settled:
| Setting | Typical value | Note |
|---|---|---|
| β₁ / β₂ | 0.9 / 0.95 | β₂ = 0.95 instead of the 0.999 default: LLM data drifts fast, the second moment must forget fast |
| Learning rate | 1e-5 ~ 3e-4 | Bigger models use smaller LRs |
| Weight decay λ | 0.1 | AdamW’s decoupled decay |
| Gradient clipping | global norm 1.0 | Guards against single-step explosions |
| Warmup | first 0.1%–2% of steps | linear ramp to peak LR |
Optimizer State Memory Calculator
Memory breakdown (GB / GPU)
AdamW: first moment m + second moment v · Assumes mixed precision (BF16 weight copy + gradients). Numbers are approximate, ignoring sharding metadata and alignment overhead.
One underestimated fact: AdamW’s optimizer states cost 8 bytes per parameter (two FP32 moments). Adding FP32 master weights, a BF16 copy, and gradients, the total is ~16–18 bytes/param — training a 70B model burns 1.1TB+ on states alone. That is the practical motivation behind Adafactor (factorized second moment), 8-bit Adam (quantized states), and Muon (a single momentum buffer).
Learning-rate schedules: WSD is replacing cosine
Section titled “Learning-rate schedules: WSD is replacing cosine”The learning rate is not a constant; it is a timetable:
- Warmup: a few hundred to thousand steps of linear ramp. Parameters are still random noise; a large LR flips the table
- Main phase: hold at peak, or decay smoothly (cosine)
- Cooldown: decay so parameters “settle” into a good local minimum
The GPT-3 era loved cosine (decay all the way down). Its engineering flaw: the decay schedule is welded to the total length — continue training and you must redesign the schedule.
WSD (Warmup–Stable–Decay) turns the middle into a constant-LR plateau and only decays in the final ~20%. Practice shows: flexible continual training, fork-anytime semantics, and no quality loss from the shorter decay. MiniCPM popularized it; Kimi K2 used WSD across 15.5T tokens of pretraining (Muon + stable-phase LR 2e-4, decaying to 2e-5).
Mixed precision: BF16 is the floor, FP8 is the frontier
Section titled “Mixed precision: BF16 is the floor, FP8 is the frontier”- BF16 training: weights and compute in BF16 (same dynamic range as FP32, no underflow), master weights and optimizer states in FP32. No loss scaling needed; the LLM default
- FP16 training: more mantissa but narrow range; requires dynamic loss scaling (multiply loss to keep gradients from underflowing). Mostly inference now
- FP8 training: DeepSeek-V3 first validated end-to-end at 671B scale. The trick is fine-grained quantization: weights quantized in 128×128 blocks, activations in 1×128 groups, with high-precision accumulation. Sensitive ops (softmax, norms) stay in BF16. Roughly doubles training throughput and is the default direction for 2025–2026 clusters
Gradient clipping and loss-spike engineering
Section titled “Gradient clipping and loss-spike engineering”Loss spikes are a fact of large-model life. The standard three-piece kit:
- QK normalization / logit soft-capping: cap attention logit magnitude structurally (Gemma-2 soft-capping, PaLM z-loss)
- Global gradient clipping: scale down any update whose global norm exceeds the threshold (typically 1.0)
- Circuit breakers and rollback: skip the spiking batch; if it persists, roll back to the last stable checkpoint
Kimi K2’s MuonClip is the newest member of this family: Muon’s update is coupled with QK-clip (anticipatorily clip attention logits once they exceed 100), and the 1T-parameter run finished with zero loss spikes and zero rollbacks — turning stability from “firefighting” into “vaccination.”
Why Muon works: the spectral-norm view
Section titled “Why Muon works: the spectral-norm view”Muon’s update has three steps: accumulate momentum → orthogonalize the momentum matrix (Newton–Schulz iterations approximating U·Vᵀ, the orthogonal factor of its SVD) → apply. Intuition: replace the raw gradient matrix with the closest orthogonal matrix, so every singular direction updates at equal magnitude.
Jeremy Bernstein and collaborators’ explanation: this is steepest descent inside the unit ball of the spectral norm. Adam corresponds to element-wise (∞-type) steepest descent, which over-weights a few coordinates; the spectral-norm constraint equalizes updates across the whole matrix — exactly what a weight matrix needs. Embeddings, MoE routers, and other non-matrix parameters stay on AdamW: Muon complements rather than replaces.
The 2025–2026 optimizer landscape
Section titled “The 2025–2026 optimizer landscape”- Scale validation: Moonlight (Moonshot) added weight decay and shape-aware scaling (×0.2·√max(m,n)), achieving ~2x AdamW’s token efficiency on a 16B MoE / 5.7T tokens; Kimi K2 pushed Muon to 1T parameters with zero spikes; DeepSeek V4 (2026) also migrated most modules to Muon
- Variant explosion: AdaMuon, AdaGO, NorMuon, NAMO, Muon2, CacheMuon — iterating on how to avoid second moments and how to reuse statistics across devices
- Open theory: Muon’s convergence theory is still incomplete; there is no operational criterion for when it clearly beats AdamW (or degrades) across data/architecture regimes
🔬 Open Research Questions
Key questions and research directions in this area:
- Can the spectral-norm steepest-descent explanation predict Muon's applicability boundaries on MoE, multimodal, or RL training?
- Is there an order-of-magnitude compression left for optimizer states (AdamW's 8 bytes/param)? What should an optimizer look like in the FP4 era?
- How short can WSD's decay phase be before final quality degrades? Could 'fork anytime during the stable phase' change the product shape of continual pretraining?
- Can MuonClip-style anticipatory stabilization (QK-clip) generalize into a principled architecture-optimizer co-design rule for explosion prevention?
References
- Adam: A Method for Stochastic Optimization
Introduces Adam, which keeps exponential moving averages of gradient first and second moments and rescales each parameter learning rate by the square root of the second moment. Fast and robust, it dominated deep learning for a decade; its decoupled variant AdamW remains the default optimizer for LLM training.
- Decoupled Weight Decay Regularization
Shows that L2 regularization inside Adam gets rescaled by the adaptive denominator, so effective weight decay drifts with gradient magnitude, and proposes AdamW with decoupled weight decay. Nearly all LLM pretraining today uses AdamW with weight decay around 0.1.
- Adafactor: Adaptive Learning Rates with Sublinear Memory Cost
Factorizes the second-moment matrix of Adam into low-rank terms, cutting optimizer state memory from O(mn) to O(m+n) with little loss in convergence. Used to train T5; a classic memory-conscious alternative to Adam and an early hint that optimizer states must slim down for large models.
- Mixed Precision Training
Systematically validates mixed precision training: FP16 storage and compute with FP32 master weights and accumulators plus dynamic loss scaling to prevent gradient underflow. Nearly doubles throughput without quality loss; BF16 later removes the need for loss scaling and becomes the LLM default, and FP8 doubles throughput again.
- Symbolic Discovery of Optimization Algorithms
Discovers optimizers via AutoML program search, yielding Lion: it compares the sign of momentum with the sign of the gradient and emits constant-magnitude updates without any second moment, saving memory and matching or beating AdamW on vision and language tasks. Showed that optimizers can be searched for, inspiring a wave of sign-based update rules.
- MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies
Systematically validates the Warmup-Stable-Decay (WSD) learning-rate schedule on small models: a long stable phase allows flexible continual training and mid-flight branching, while a short final decay recovers cosine-level quality. The recipe was adopted by the MiniCPM family and Kimi K2, becoming a mainstream schedule for long-horizon training.
- Muon: An Optimizer for Hidden Layers
Introduces Muon: orthogonalize the momentum buffer (Newton-Schulz iteration approximating UV^T) before applying it to hidden-layer matrices, interpretable as steepest descent under the spectral norm so every layer update has consistent spectral magnitude. Over 1.35x faster than a tuned AdamW on the NanoGPT speedrun at the cost of one momentum orthogonalization, making it the most watched post-Adam optimizer.
- Muon is Scalable for LLM Training
The Moonlight engineering report: adds weight decay to Muon and rescales updates by matrix shape (0.2*sqrt(max(m,n))) so learning rates transfer from AdamW recipes. Validates Muon at production scale on a 16B MoE trained on 5.7T tokens with roughly 2x the token efficiency of AdamW, proving Muon can move from speedrun tracks to real pretraining.
- Kimi K2: Open Agentic Intelligence
The Kimi K2 technical report: a 1T-total / 32B-active MoE pretrained on 15.5T tokens. On the optimization side it pushes Muon to trillion-parameter scale with MuonClip (QK-clip: anticipatorily clip attention logits once they exceed a threshold of 100) together with a WSD learning-rate schedule, achieving zero loss spikes and zero training rollbacks; weights and detailed infra notes are open-sourced.