Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks

TL;DR

Kosson, Messmer & Jaggi extend Spherical Motion Dynamics (Spherical Motion Dynamics: Learning Dynamics of Neural Network with Normalization, Weight Decay, and SGD) from SGDM to modern optimizers (AdamW, Adam+L2, Lion) and to non-scale-invariant weights. They derive closed-form steady-state weight norm and angular update per neuron for each optimizer, show empirically that the predictions hold across ResNet-50/ImageNet, DeiT/ImageNet, and GPT-2 on WikiText/OpenWebText, and construct “Rotational Variants” (RVs) — optimizers that explicitly enforce the equilibrium rotation without weight decay. Their central mechanistic claim: balanced rotation across layers and neurons is the underlying benefit of AdamW over Adam+L2 and of Weight Standardization over vanilla LayerNorm — because Adam+L2’s equilibrium rotation depends on gradient magnitude (which varies across layers) while AdamW’s does not.

Key claims

For a scale-invariant weight vector under any of {SGDM, AdamW, Adam+L2, Lion}, weight-decay-driven radial shrinkage and gradient-driven orthogonal growth converge to an equilibrium with closed-form expected norm and closed-form expected angular update [Table 1, §3].
AdamW’s equilibrium angular update E[Δθ] ≈ √(2ηλ/(1-β₁)) depends only on hyperparameters; Adam+L2’s depends additionally on the per-tensor gradient RMS E[|g|], so identical hyperparameters produce different equilibrium rotation across layers whenever gradient scales differ [Table 1, §3.3].
The “Rotational Variant” (RV) of an optimizer replaces weight decay with an explicit projection to the equilibrium ‖w‖ after each step and enforces the equilibrium Δθ; RV-AdamW matches baseline AdamW accuracy zero-shot on CIFAR-10/ResNet-20, GPT-2/OpenWebText, GPT-2/WikiText, and with one weight-decay retune on DeiT/ImageNet and Transformer-S/IWSLT2014 [Table 2, §5.1].
Balanced rotation is the load-bearing property: enforcing per-layer equilibrium rotation via RV closes the ≈0.5% AdamW-vs-Adam+L2 gap on CIFAR-100/ResNet-18/LayerNorm; the actual angular-update trajectory shows Adam+L2 produces ~30× variation in per-layer rotation (corresponding to ~1000× effective LR variation) while AdamW keeps all layers within a narrow band [§5.3, Fig. 6].
Weight Standardization = balanced rotation per neuron: LayerNorm alone makes only the whole layer scale-invariant, so individual neurons can still exhibit radial gradient components producing imbalanced per-neuron rotation; WS makes each neuron scale-invariant and eliminates the imbalance. On CIFAR-100/ResNet-18/LayerNorm, both WS and RV significantly outperform baseline training and produce visibly balanced per-neuron rotation curves [§5.3, Fig. 7].
LR warmup partially substitutes for enforcing rotational equilibrium: on GPT-2-124M/OpenWebText and ResNet-50/ImageNet-8k-batch, baseline AdamW’s benefit from a 5% warmup is large; the same runs under RV-AdamW show almost no warmup benefit — the transient rotation instability that warmup was needed to suppress is instead eliminated by construction [§5.2, Fig. 5].

Method

The paper’s core object is the Total Update Contribution (TUC) of a single step’s gradient — the sum of all future weight updates it produces via momentum accumulation — as opposed to the immediate update. Under a random-walk assumption (batch gradient dominated by noise, common in stochastic-differential-equation analyses of optimizers), the gradient TUC is orthogonal to the weight in expectation via symmetry, and the weight-decay TUC is anti-parallel. Applying Pythagoras to the “TUC triangle” gives a fixed-point equation for E[‖w‖²] that admits closed-form solutions per optimizer.

For AdamW specifically, the RMS update size becomes η·(1-β₁)/√(1-β₂) · d in expectation under the random-walk assumption; combined with the equilibrium norm this yields the tabulated E[Δθ] = √(2ηλ/(1-β₁)). The Rotational Variants (Algorithm 1) enforce this in two moves: after each step (a) rescale w ← w · (‖w‖_target / ‖w‖) and (b) optionally learn a per-tensor gain to compensate. Since the RV is always in equilibrium, the specified LR schedule directly controls Δθ — unlike in standard optimizers where the transient phase decouples Δθ from η(t).

Results

AdamW vs Adam+L2 (Fig. 6): Reproduces Loshchilov & Hutter’s ≈0.5% CIFAR-10 gap. AdamW’s per-layer angular update stays within a narrow band throughout training; Adam+L2’s varies by ~30° between layers, corresponding to ~1000× effective-LR variation.
Weight Standardization vs LayerNorm (Fig. 7): On CIFAR-100/ResNet-18, WS and RV both outperform vanilla LayerNorm by ~1 point and produce visibly balanced per-neuron rotation; the baseline is imbalanced across neurons within a single LN layer.
Warmup ablation (Fig. 5): On GPT-2-124M/OpenWebText, no-warmup runs open a validation-loss gap that never closes under baseline AdamW; under RV-AdamW the gap is negligible across LRs. Similar on ResNet-50/ImageNet-8k-batch.
Zero-shot RV matches baseline on 5/6 configurations in Table 2; the IWSLT2014 Transformer-S needs a light retune (baseline was trained with effectively zero weight decay, so equilibrium was never reached — the RV surfaces this).

Why it’s interesting

Turns the SMD framing from a descriptive tool into a prescriptive one: if balanced equilibrium rotation is what modern optimizers are secretly optimizing for, you can build it in directly (RV) and dispense with weight decay tuning. This is the direct theoretical ancestor of a whole cluster of 2025-2026 weight-geometry recipes on the wiki:

Modular Manifolds generalizes RV’s “project to equilibrium norm” to arbitrary per-module manifolds
Controlled LLM Training on Spectral Sphere (SSO) constrains the spectral norm (not Frobenius) to a per-module sphere and demonstrates µP transfer at 1.7B/8B
Delphi: Scaling Laws That Extrapolate 300× Past the Fit‘s AdamH is the Frobenius-sphere version deployed in a public 300× scaling-law extrapolation
Introducing Inkling — Thinking Machines' open-weights 975B/41B MoE multimodal reasoning model‘s WD ∝ LR² is a modular-manifolds-style coupling shipped at 975B/41B MoE

The paper also directly grounds How to Set the Learning Rate for Large-Scale Pre-training?‘s empirical observation that per-module AdamW update magnitudes are ≈0.2 across the network — the Fitting paradigm’s argument for a single global LR under modern normalization is the rotational-equilibrium prediction. Where Kyle’s skepticism lands is worth naming: the paper’s own zero-shot RV matches baselines on 5/6 configs — impressive as a validation of the theory but not a wall-clock or benchmark win. Whether “study angular updates in isolation” captures enough of the geometry to matter dynamically is unresolved.