Training Transformers with Enforced Lipschitz Constants

TL;DR

Trains transformers with a Lipschitz certificate enforced throughout training (not just at initialization), by maintaining norm-constrained weight matrices via novel computationally-efficient projections. The headline finding is that switching from AdamW to Muon strictly improves the Lipschitz-vs-performance tradeoff of standard recipes (weight decay, spectral normalization), and a co-designed weight constraint (spectral soft cap, spectral hammer) tuned to Muon’s fixed-spectral-norm update pushes the frontier further. A <2-Lipschitz transformer reaches 60% validation accuracy on Shakespeare; a <10-Lipschitz 145M model reaches 21% on internet text; matching the NanoGPT 39.4% baseline requires letting the bound blow up to 10^264 — but the Lipschitz models train without layer norm, QK-Norm, or logit tanh softcapping.

Key claims

Optimizer choice changes the Pareto frontier of Lipschitz-bound vs. loss: weight decay and spectral normalization both reach better validation loss at a lower Lipschitz bound under Muon than under AdamW [Fig. 2].
A new family of weight-constraint methods — spectral soft cap (a smooth approximation of spectral projection containing weight decay as a special case) and spectral hammer — is co-designed with Muon’s spectral-norm-preserving update and improves over weight decay / spectral norm regularization on MLPs and 2M-param transformers [Fig. 2, App. A].
A <2-Lipschitz transformer on Shakespeare achieves 60% validation accuracy [Abstract].
A <10-Lipschitz 145M-parameter transformer reaches 21% accuracy on internet text [Abstract].
Matching the NanoGPT validation-accuracy baseline of 39.4% requires the Lipschitz upper bound to grow to ~10^264 — i.e. a tight global Lipschitz bound at frontier-LM accuracy remains an open gap [Abstract].
Lipschitz-constrained transformers train stably without the usual stabilization stack: no LayerNorm, no QK-Norm, no logit tanh softcapping [Abstract].
Framing: “don’t just regulate the gradient like Muon does, regulate the weights too” — Muon orthogonalizes updates (fixed spectral norm), so the natural companion is to keep the weight spectra bounded as well [project repo, §1].

Method

The paper develops a toolkit of weight-projection operators that can be applied at each optimizer step to keep ‖W‖₂ ≤ τ for chosen module-wise τ. Beyond standard weight decay (which targets the average singular value) and spectral norm regularization (which targets only the top singular value), the authors derive a one-parameter family they call spectral soft cap — a smooth approximation of hard spectral projection that contains weight decay as a limit, computed via Newton–Schulz-style iterations on the singular spectrum — and spectral hammer, which targets sufficient singular values to enforce the bound directly. Combined with Muon (whose update step itself has a fixed spectral norm via its orthogonalization), this gives matched control over both the weight matrices and the updates that perturb them. The composite Lipschitz constant of the network is then the product of per-layer spectral bounds, attention re-scaling factors, and activation Lipschitz constants — tracked analytically through training. Critically, no LayerNorm / RMSNorm / QK-Norm / logit-cap stabilization layers are inserted; stability comes from the weight constraint plus Muon dynamics alone.

Results

Sweep of 1600 MLPs on CIFAR-10 and 800 transformers on Shakespeare across optimizers × weight-constraint methods. Weight decay and spectral normalization both yield lower validation loss at equal Lipschitz bound when run under Muon vs. AdamW. Spectral soft cap and spectral hammer extend the Pareto frontier further [Fig. 2, App. E].
<2-Lipschitz Shakespeare transformer: 60% validation accuracy (no normalization layers) [Abstract].
<10-Lipschitz 145M-param transformer on internet text: 21% validation accuracy [Abstract].
Closing the gap to the NanoGPT baseline (39.4%) currently requires the Lipschitz bound to grow to ~10^264 — i.e. tight bounds are not yet free at frontier accuracy [Abstract].

Why it’s interesting

The paper is the inverse of the dominant “architectural-stability” line on the wiki: rather than adding more normalization layers (QK-Norm, GatedNorm in A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training; Birkhoff-constrained residual mixing in mHC: Manifold-Constrained Hyper-Connections; LayerScale-style stabilization in ViT-5: Vision Transformers for The Mid-2020s) to manage outliers, it constrains the weight spectra so that the outliers never get a chance to grow — and the stabilization layers become unnecessary. This is a third position alongside the architecture-side fix and the optimizer-side fix proposed by Controlled LLM Training on Spectral Sphere (SSO); in fact the philosophical alignment with SSO is striking (“Muon is half-aligned with µP because it controls updates but not weights” — SSO and this paper give the same diagnosis and very similar prescriptions of constraining weights to a spectral sphere). For Luma, two practical implications: (1) if the no-normalization Lipschitz recipe holds at DiT scale, it removes the AdaLN/QK-Norm modulation stack that currently fights with conditioning; (2) the spectral-soft-cap operator is a candidate drop-in for Muon implementations that already do Newton–Schulz on updates (Addition of Muon optimizer to torch.optim — PyTorch team welcomes PR (Issue #148819), Aurora: A Leverage-Aware Optimizer for Rectangular Matrices).