Skip to content
Bud's Wiki

Training stability at scale

TL;DR

Section titled “TL;DR”

A late-2025/early-2026 cluster of papers locates large-model training stability in architectural choices (normalization variants, residual-mixing constraints, gated rescaling) rather than in parametrization recipes like µP. The shared empirical claim is that modern normalization (QK-Norm, RMSNorm, GatedNorm) and constrained residual geometry (Birkhoff-polytope residual mixing) subsume what classical parametrization tricks were originally designed to fix — and in some cases the architectural fix is strictly better because it also removes pathologies (outlier activations, identity-mapping violations across depth) that the parametrization fix never addressed.

Key claims

Section titled “Key claims”
  • Under modern architectures, the attention-logit blow-up that µP was designed to suppress simply does not occur under standard parametrization; ablating QK-Norm restores the instability µP reports, so the stability is doing the work, not the parametrization (How to Set the Learning Rate for Large-Scale Pre-training? §6.3, Figs. 8–9).
    • Conflict: Controlled LLM Training on Spectral Sphere argues the opposite resolution — µP’s motivation still stands, but current optimizers (including Muon) are only “half-aligned” (they control update spectra but allow weight spectra to drift), and the right fix is to constrain both weights and updates to a spectral sphere rather than rely on QK-Norm to mask the symptom (Controlled LLM Training on Spectral Sphere README §1, Abstract). The two papers prescribe directly opposing actions: drop µP and lean on normalization (2601.05049) vs. enforce µP harder via the optimizer (2601.08393).
    • Finetuning-side replication of the parametrization position: Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning shows that for LoRA finetuning, µP-style asymptotic analysis (extended to the joint n,r→∞n, r \to \infty limit as µA) still yields tight, predictive scaling rules for the optimal LR — including a configuration (Init[B], α=1\alpha=1) where LoRA’s optimal LR transfers directly to FFT (Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning §4.2, Fig. 5). The µP framework is not dead in the finetuning regime even if its motivating instability has been suppressed in from-scratch pretraining.
    • Vision-side replication: ViT-5: Vision Transformers for The Mid-2020s reports that adding QK-Norm to ViTs gives a modest accuracy bump and a large reduction in sharp loss spikes during optimization — the same spike-suppression role QK-Norm plays in LLMs (ViT-5: Vision Transformers for The Mid-2020s §3.6, Fig. 5).
  • Module-specific learning rates (a µP prescription) yield essentially no measurable improvement over a single global LR at 4B-MoE scale under modern normalization, because per-module AdamW update magnitudes are already uniform across layers (≈0.2) consistent with rotational equilibrium (How to Set the Learning Rate for Large-Scale Pre-training? §6.1–6.2, Figs. 5–7).
  • Attention sinks and residual sinks both arise from the same functional pressure — outliers paired with softmax/RMSNorm act as implicit rescale factors for non-outlier components, which is load-bearing for training stability (removing normalization or clipping outliers both degrade training) (A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training §3, Abstract conclusions (1)–(2)).
  • The implicit rescaling can be replaced with explicit learned gating (Gated Attention, GatedNorm), giving ≈+2-point average pretraining gain and ≈1.2-point less W4A4 quantization degradation — i.e. architecture choices that maintain stability and produce low-precision-friendly activation distributions (A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training Abstract conclusion (3), Table 1).
  • Unconstrained Hyper-Connections break the identity-mapping property of residual stacks at scale, producing forward/backward signal gains up to ~3000× and visible 27B-scale loss spikes; constraining the residual mixing matrix to doubly stochastic matrices (Birkhoff polytope) via a differentiable Sinkhorn projection bounds composite gain to ~1.6× across arbitrary depth (mHC: Manifold-Constrained Hyper-Connections §3.1, §4.1, §5.4, Figs. 2/3/7).
  • Architectural stability fixes hold across compute and data scaling — the loss-improvement gap of mHC over baseline is maintained from 3B → 9B → 27B, and across a 1T-token run at 3B (mHC: Manifold-Constrained Hyper-Connections §5.3, Fig. 6).
  • The slope of LM pre-training scaling laws — not just the intercept — appears to be set by a structural primitive (softmax + cross-entropy on peaked target distributions), yielding a universal 1/3 time-scaling exponent that interventions like optimizer choice, data mixes, and most architecture tweaks do not change (Universal One-third Time Scaling in Learning Peaked Distributions Abstract, Figs. 1–3). This is the “tyranny of the slope” framing: stability/architecture interventions move the intercept; only loss/output-head changes can move the slope.
  • Constraining both weight matrices and updates to a per-module spectral sphere of radius R=Θ(dout/din)R = \Theta(\sqrt{d_\text{out}/d_\text{in}}) — the steepest-descent direction under the spectral norm, with hard or dynamic retraction — outperforms AdamW and Muon on Dense 1.7B, MoE 8B-A1B, and a 200-layer DeepNet while suppressing outliers, bounding activations, and improving MoE router load-balancing (Controlled LLM Training on Spectral Sphere README §1, §2, §4). This is the optimizer-side analogue of the architecture-side mHC residual-mixing constraint: same manifold-constraint recipe, applied to a different stability locus.
  • Architectural stability components do not transfer trivially across modalities: porting SwiGLU MLPs (a standard LLM choice) into ViTs alongside LayerScale produces over-gating — both act as channel-wise filters, the combination drives intermediate representations to excessive sparsity, and accuracy drops (84.16% → 83.70% at ViT-B, FID 14.57 → 16.22) (ViT-5: Vision Transformers for The Mid-2020s §3.3, Table 2). The QK-Norm, RMSNorm, LayerScale, and bias-free-QKV pieces do transfer; SwiGLU does not.
  • MoE router stability has a parallel “remove the tuning axis” story: Quantile Balancing (QB) reformulates per-step expert assignment as the linear program “find the bias vector that would have produced balanced load last step” and solves it via alternating quantile iteration — eliminating the aux-loss coefficient, the fixed-bias-update magnitude, and the momentum coefficient that prior recipes (GShard, DeepSeek loss-free, Arcee SMEBU) all required. Validated at 32B-A5B / 1e22 FLOPs / 326B tokens on Marin with zero loss spikes and no leading dense layers (Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method).

Recent contributions

Section titled “Recent contributions”
  • [2026-07-16] Kimi K3 — Open Frontier Intelligence (2.8T MoE with KDA + AttnRes): Kimi K3 pairs Per-Head Muon (Muon with per-attention-head decoupling) with QAT-from-SFT in MXFP4 weights + MXFP8 activations and reports stable training at 2.8T parameters with 16/896 expert sparsity.
  • [2026-07-15] Introducing Inkling — Thinking Machines' open-weights 975B/41B MoE multimodal reasoning model: Inkling’s public training recipe: hybrid Muon (matrix weights) + Adam (everything else) with weight decay coupled to LR² — the LR²-coupling is claimed to keep weight-matrix norms stable across training horizons, letting the same recipe transfer across LR schedules. First frontier-scale from-scratch pretraining run to name modular-manifolds-style spectral control + LR²-coupled decay as its stability primitive; 45T-token, 975B/41B MoE run on GB300 NVL72.
  • [2026-07-13] Modular Duality in Deep Learning: Bernstein & Newhouse’s foundational duality-map paper — the theoretical origin of the rectangular Newton-Schulz iteration that Muon uses to orthogonalize its momentum. Derives per-layer duality maps from operator norms (RMS→RMS for Linear, ∞→RMS for Embed, max-over-kernel for Conv2D) and composes them via the modular norm’s max-weighted composition, giving straightforward all-layer updates. Shows µP and Shampoo both fall out as partial approximations to the single √(d_out/d_in)·UVᵀ duality map for Linear — reframing “should we tune per-module LRs?” as “which pieces of this duality map does the optimizer actually implement?” Grounds the SSO claim that Muon is “half-aligned with µP” in a precise object.
  • [2026-07-10] The 4-bitter Lesson: Balancing Stability and Performance in NVFP4 RL: Dequantized backward as a new stability-under-quantization lever: replacing y = x·w_bf16 in the BF16 backward with y = x·DQ(Q(w_bf16)) matches the chain rule to the actual quantized forward, eliminating the grad-norm spikes that appear under the “quantized forward, higher-precision backward” baseline. Under Adam’s momentum + adaptive second moment the extra NVFP4 gradient variance is absorbed, yielding far fewer large spikes than BF16 backward. Combined with 4/6 and selective precision, NVFP4 RL matches BF16 grad-norm across five runs where any two-out-of-three combination still spikes.
  • [2026-07-09] Scaling with Collapse: Efficient and Predictable Training of LLM Families: Bergsma et al. deploy scaling-collapse residuals as a production stability signal on a Cerebras 300M–3.9B (Celerity) family: at 1.8B, collapse residuals surfaced a numerical instability at ~60% of training, whereas the raw loss trend only revealed it at ~90%. Ablations against the collapse reference isolated the bug to a microbatch-size-triggered loss-kernel numerical issue and enabled a targeted restart. Concrete demonstration that the Qiu et al. supercollapse diagnostic operationalizes at LLM scale as an earlier-and-sharper warning than trend-watching. Also documents CompleteP outperforming µP for joint width–depth scaling on this family, adding a data point to the µP-vs-CompleteP debate elsewhere on the page.
  • [2026-07-09] Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks: Qiu et al. establish scaling collapse and supercollapse as a live diagnostic: mis-parameterization (dropping µP for a constant LR across widths) breaks collapse even at small scales where final losses are indistinguishable. Adds a new class of stability signal to the page — collapse residuals give an earlier, more sensitive warning of scaling-parametrization pathologies than final-loss comparison or trend-watching alone. Complements Training stability at scale‘s architecture-vs-optimizer-vs-parametrization framing with a fourth “does dynamics stay consistent across scales?” axis.
  • [2026-07-08] LingBot-Video: Scaling Mixture-of-Experts Video Pretraining for Embodied Intelligence: LingBot-Video’s stability recipe for scaling a MoE video DiT to 120B/11B-active: QK-Norm (per-head RMSNorm on Q/K before attention), adaLN-single modulation (one shared timestep MLP + per-layer trainable modulation tables, zero-init on shared projection), auxiliary-loss-free bias for MoE load balancing plus sequence-wise Lseq to prevent per-video routing collapse, and √(d/d_a)-style variance-preserving expert init.
  • [2026-07-01] To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters: Dragutinović et al. attack Muon from a new angle — not stability, not µP-alignment, but inductive bias. Under an exact-SVD, no-momentum abstraction of Muon (Spectral GD) on 2-layer deep linear networks, they prove GD’s classical saddle-to-saddle dynamics sequentially increase solution rank (largest singular value first) while Spectral GD learns every singular component in parallel at equal rate — losing GD’s implicit rank-based regularization. Two constructed head-to-heads have SGD beat Muon: (a) a multi-source “routing” task where SGD discovers a shared 4-D representation and generalizes to unseen source pairs while stochastic Spectral GD memorizes each pair despite matching train loss, and (b) MNIST + per-class spurious pixel where SGD’s dominant-first learning gives higher clean-digit accuracy up to spurious intensity ≈2. Also observes that after the theoretical low-rank modes are learned, Spectral GD’s discretization noise makes the update full-rank and it starts fitting that noise in a highly oscillatory trajectory that momentum does not damp — a mechanism-level hint that Muon may be less robust at high learning rates in richer settings than the current stability literature captures.
  • [2026-06-30] What do near-optimal learning rate schedules look like?: Naganuma et al. find that on the small WikiText-103 Transformer workload, the most aggressive near-optimal LR shapes produce some training-instability blowups (large-perplexity runs) that inflate confidence intervals, echoing the broader story that the optimal shape sits near the stability boundary — and on linear regression the analytical optimum is exactly at the edge-of-stability. AdamW weight decay (a known stability knob) is the strongest covariate of optimal shape, supporting the architecture/optimizer-side stability framing on this page.
  • [2026-06-24] Krea 2 Technical Report: Krea 2’s architecture-stability moves on a production T2I MMDiT: GQA + sigmoid-gated attention (small wall-clock cost, no quality gain but visibly smoother loss and gradient-norm curves), zero-centered RMSNorm with weight decay on its learnable params, QKNorm, per-block tunable bias replacing AdaLN MLPs. Reports Muon (Dion implementation + Moonlight RMS-matched HP transfer) is stable enough to beat AdamW at low and high resolution only after excluding the first and last linear layers and adding Nesterov momentum — otherwise it triggers loss/gradient-norm spikes.
  • [2026-05-26] Laguna M.1/XS.2 Technical Report: Laguna documents the M.1→XS.2 stability redo: keep the first Transformer layer dense, FP32 RMSNorm/RoPE, configurable gradient-reduction dtypes per collective (FP32 default for sensitive paths including the LM-head input-gradient all-reduce), and a WSD schedule with the peak chosen from an N,D scaling law rather than a one-shot cosine tune.
  • [2026-05-24] Addition of Muon optimizer to torch.optim — PyTorch team welcomes PR (Issue #148819): PyTorch maintainers’ July-2025 comment welcoming a PR to add Muon to torch.optim, explicitly citing Moonshot’s MuonClip / Kimi K2 deployment as the social proof. Marks the point at which Muon transitioned from “research-lab optimizer” to a framework default. The infrastructure caveat — distributed Muon (FSDP/HSDP) is still not supported in core, and Kimi K2 disables FP8 GEMM for the Newton–Schulz step — is the open piece the SSO line (Controlled LLM Training on Spectral Sphere) and Aurora (Aurora: A Leverage-Aware Optimizer for Rectangular Matrices) are addressing.
  • [2026-05-23] Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method: hyperparameter-free MoE load balancer (QB) from Jianlin Su, formulated as a per-step LP over bias offsets and solved via alternating quantile iteration. Sits in the same lineage as MuonClip / qk-clip — “eliminate the tuning knob entirely” — but applied to the router rather than to attention logits. Validated by the Marin team at 32B-A5B / 1e22 FLOPs (326B tokens) with zero loss spikes; 1e23-FLOP run underway. Suggests MoE routing is another locus where the “constrained optimization replaces hand-tuned heuristic” pattern works.
  • [2026-05-23] Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning: introduces µA (Maximal-Update Adaptation), the LoRA analogue of µP, characterizing optimal-LR scaling in the joint n,r→∞n, r \to \infty limit. Identifies Init[B] + α=1\alpha=1 as the unique LoRA configuration whose optimal LR coincides with full-finetuning’s, enabling direct LR transfer. Provides a finetuning-side counterpoint to 2601.05049: µP-style asymptotic analysis is still load-bearing here, even though from-scratch LLM pretraining no longer needs it under QK-Norm. Verified across SFT (5 models), RLVR, and Stable Diffusion.
  • [2026-05-23] ViT-5: Vision Transformers for The Mid-2020s: vision-side replication of the architectural-stability story — QK-Norm in ViTs suppresses sharp loss spikes during optimization, mirroring the LLM-side finding. Adds a cross-modality negative result: SwiGLU + LayerScale together over-gate ViTs (both are channel-wise filters; combined sparsity hurts), so the LLM→ViT port has to drop SwiGLU even though every other component transfers. Hits ImageNet-L/384² SoTA (86.0%) and SiT-XL FID 1.84.
  • [2026-05-23] Controlled LLM Training on Spectral Sphere: introduces SSO, an optimizer that constrains both weights and updates to a per-module spectral sphere via steepest descent under the spectral norm + Newton–Schulz msign + retraction. Argues Muon is “half-aligned” with µP (controls updates, not weights) and that closing this loop gives strict Θ(1)\Theta(1) activation control end-to-end. Outperforms AdamW and Muon on Dense 1.7B, MoE 8B-A1B, and 200-layer DeepNet; ships as a Megatron-LM fork with Atomic Module Sharding and cached singular vectors. Directly conflicts with 2601.05049 on whether µP’s motivation still stands.
  • [2026-05-23] Universal One-third Time Scaling in Learning Peaked Distributions: derives a universal 1/3 time-scaling exponent for cross-entropy loss when the target distribution is peaked, attributing it to softmax saturation and the resulting power-law-vanishing gradient. Reframes most pretraining interventions (Muon, data mixes, normalization) as intercept-only shifts of the scaling law — its slope is structural.
  • [2026-05-23] A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training: unifies attention sinks and residual sinks as outlier-driven rescaling; replaces the implicit mechanism with explicit Gated Attention / GatedNorm, yielding pretraining gains and W4A4 quantization robustness. Locates a previously-unexamined coupling between training stability and low-precision friendliness.
  • [2026-05-22] How to Set the Learning Rate for Large-Scale Pre-training?: shows µP’s motivating attention-logit instability is suppressed by QK-Norm under SP, so module-specific LRs no longer pay off and the Fitting paradigm (scaling-law extrapolation of LR*(N,D)) beats µTransfer at 4B and 12B MoE scale.
  • [2026-05-22] mHC: Manifold-Constrained Hyper-Connections: makes residual-stream geometry the load-bearing stability primitive by constraining the residual-mixing matrix to the Birkhoff polytope, suppressing 27B-scale loss spikes and bounding composite signal gain at any depth.

Open questions

Section titled “Open questions”
  • All three architectural-stability papers locate stability in different places in the network — Q/K projections (QK-Norm), attention output + RMSNorm (GatedNorm / Gated Attention), residual mixing (mHC). Are these complementary (stack them) or do they substitute for one another at sufficient strength? No filed paper has run the stacked ablation.
  • Architecture-side vs optimizer-side stability fixes are now an explicit fork. SSO (Controlled LLM Training on Spectral Sphere) suppresses outliers and bounds activations as a side-effect of spectral retraction; Gated Attention (A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training) preserves outliers but adds explicit learned rescaling; mHC bounds composite signal gain at the residual-mixing level. Do these stack (SSO + Gated Attention + mHC) or do they each saturate the same underlying degree of freedom (bounded effective spectra of the forward/backward map)? No filed paper has done the cross-comparison.
  • Pretraining vs finetuning may be a real split. How to Set the Learning Rate for Large-Scale Pre-training? declares µP’s stability motivation stale for from-scratch pretraining under QK-Norm; Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning shows µP-style analysis remains predictive for LoRA finetuning. Whether this means “stability is solved by architecture, but rank/init/multiplier coupling needs separate parametrization analysis” or “µP’s relevance is regime-dependent in a way no filed paper has fully mapped” is open.
  • The Gated Attention story explicitly couples stability with quantization-robustness — outliers are both why training works and why PTQ fails. Does the same coupling apply to mHC’s residual-stream constraint (does bounding composite gain to ~1.6× also make activations PTQ-friendly)? SSO claims the same coupling (bounded activations → low-precision friendly). The mHC paper does not measure this; the SSO paper asserts it without quantifying PTQ deltas.
  • µP was designed for LLM decoder-only stacks. Do these architecture-first stability fixes transfer to DiT / MMDiT pretraining, where AdaLN modulations are an additional outlier source and the residual streams carry conditioning rather than just token state? Untested across all five filed papers — though ViT-5: Vision Transformers for The Mid-2020s supplies a partial answer for plain SiT backbones (QK-Norm transfers cleanly to vision; SwiGLU does not) and Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning does run Stable Diffusion finetuning (Init[B] + α=1\alpha=1 optimum holds; FFT is sharply LR-sensitive on UNets). The SSO Megatron fork is the lowest-friction way to test the optimizer-side hypothesis on a DiT codebase.
  • Cross-modality component transfer is not uniform: ViT-5 finds QK-Norm, RMSNorm, LayerScale, RoPE, and bias-free-QKV all transfer from LLMs to ViTs, but SwiGLU does not (over-gates with LayerScale). Is there a similar transfer asymmetry going the other direction — components that work in ViTs but not LLMs? The fact that LayerScale ≈ post-RMSNorm in ViTs (ViT-5: Vision Transformers for The Mid-2020s §3.1, Table 1) but LayerScale is rare in LLMs raises this question directly.
  • The Fitting paradigm in How to Set the Learning Rate for Large-Scale Pre-training? fits LR*(N, D) under WSD. If the architectural stability story is right, the LR* surface should become flatter under stronger normalization (less penalty for slight LR overshoot). Could be a clean prediction to falsify. SSO’s claim that µP transfer still holds under spectral-sphere constraints is a related, testable prediction.
  • All papers run under LLM-style pretraining (decoder-only, large-N MoE). Do the same instabilities (sinks, identity-violation, attention-logit blowup, 1/3-exponent slope) appear at the smaller-N high-batch regime that diffusion-from-scratch pretraining actually lives in? The 1/3 paper makes a specific prediction here: MSE / flow-matching objectives should not exhibit the softmax-CE bottleneck and so the DiT slope should be in a different universality class.
  • If the 1/3 exponent is structural to softmax-CE on peaked targets, does it survive when the target distribution is broadened (label smoothing, mixture-of-token-distributions distillation, temperature-scaled teachers)? A clean falsifier for whether the slope can be moved.
  • Is “Muon is half-aligned with µP” the right framing, or does Muon’s implicit weight-spectrum control via update orthogonalization already deliver most of SSO’s gains? SSO’s MuonSphere variant ablates exactly this — claim is the gap matters, but the precise delta between Muon and MuonSphere is the load-bearing experiment for this whole line.
  • Does QB’s LP framing generalize to other discrete-routing decisions beyond MoE expert assignment — e.g. mixture-of-depths gating, conditional-compute heads, layer-skip routing? QB’s claim is that the routing target (uniform load) is what makes the LP solvable in closed-quantile form; whether other targets (e.g. accuracy-weighted load, capacity-aware load) admit similar one-shot solutions is open.
  • How does QB interact with the SSO finding that spectral-sphere training “improves MoE router load-balancing” as a side-effect? If the optimizer already bounds router activations and reduces outlier-driven imbalance, do you still benefit from QB on top, or do they substitute? Both touch MoE balancing but from different loci (optimizer vs. routing-time correction).
  • What does framework-default Muon (no distributed support) imply for adoption? Per Addition of Muon optimizer to torch.optim — PyTorch team welcomes PR (Issue #148819), torch.optim.Muon shipped without FSDP/HSDP support, and Kimi K2’s deployment explicitly disables FP8 GEMM for the Newton–Schulz step. The follow-on optimizer papers (SSO, Aurora, Newton-Muon) increasingly target the implementation gap — distribution, FP8, leverage-awareness — rather than the algorithm itself. Whether one of these becomes the canonical “distributed Muon” recipe in torchtitan is the live infra question.

Papers

Section titled “Papers”

2026-05

Section titled “2026-05”

2026-04

Section titled “2026-04”

2026-03

Section titled “2026-03”

2026-02

Section titled “2026-02”

2026-01

Section titled “2026-01”

2025-12

Section titled “2025-12”

2025-11

Section titled “2025-11”

2025-10

Section titled “2025-10”

2025-09

Section titled “2025-09”

2025-08

Section titled “2025-08”

2025-07

Section titled “2025-07”

2025-06

Section titled “2025-06”

2025-05

Section titled “2025-05”

2025-04

Section titled “2025-04”

2025-03

Section titled “2025-03”

2025-02

Section titled “2025-02”

2025-01

Section titled “2025-01”

2024-01

Section titled “2024-01”

2020-12

Section titled “2020-12”

Unknown date

Section titled “Unknown date”