Adaptive optimizers¶

Imagine hiking down a steep canyon whose floor snakes diagonally across the landscape, yet your compass only lets you step north-south or east-west. Every move is a zig or a zag, and half your energy is wasted correcting your direction rather than closing distance. That is what happens when AdamW or its kin chase gradients in a coordinate system that has no relationship to the thin, twisted subspace the data actually lives in. Adaptive optimizers used to mean per-coordinate scaling—rescale each gradient element by its RMS history and hope the world straightens out. Today the bottleneck is not the per-coordinate rate but the coordinate system itself. Imagine instead being able to rotate your compass on the fly so that its axes naturally align with the ravine and you can sprint along the easiest descent path. By the end of this page you can explain how spectral orthogonalizers such as Muon inject that rotation, why their adaptivity is more stable because it respects the low stable rank of activations, and how to implement a practical Muon-inspired optimizer to feel the difference against AdamW.

The territory¶

Optimization in deep learning is still fundamentally gradient descent, but its landscape has become much harder. Deep nets stretch across billions of parameters and very narrow manifolds of activations, so the curvature varies wildly in different directions. AdamW solved the immediate problem of coordinate-wise scale by dividing each gradient element by a running RMS estimate, yet it still assumes that the most difficult directions align with axes. In practice, the real directions of high curvature are rotated combinations of weights shaped by the activations feeding them, a phenomenon especially severe in transformers and diffusion models (Author et al. 2026) arxiv:2603.18168. The modern pretraining stack therefore leans toward geometry-aware optimizers: instead of rescaling along messy axes, they learn a low-rank basis for the gradient, rotate into that space, perform a stable update, and then rotate back. This new family borrows from spectral methods, model reduction, and orthogonalization routines that were once confined to second-order solvers, but now they execute cheaply in deep networks. As we trace this transition, consider the key tension: how do we keep the flavor of Adam’s element-wise adaptivity while adding a coordinated, low-rank rotation that respects stable rank? The mechanism is best understood by starting from the gradient covariance of the layer activations and building an adaptive orthonormal basis.

How it works¶

Modern spectral adaptive optimizers replace the coordinate-wise RMS scaling with a two-phase update: (1) estimate the dominant directions of the layer’s activation-gradient alignment, and (2) apply a scaled, orthogonal step in that subspace before returning to the original coordinates. The first insight comes from measuring stable rank. Given the layer activations \(A \in \mathbb{R}^{B \times d}\) where \(B\) is the batch size and \(d\) the feature size, the stable rank is defined as \(\text{srank}(A) = \frac{\|A\|_F^2}{\|A\|_2^2}\), where \(\|\cdot\|_F\) is the Frobenius norm and \(\|\cdot\|_2\) is the spectral norm. Because activations in deep nets are often low stable rank, most of the meaningful gradient energy lies in a concentric set of singular vectors. When the gradient \(\nabla_\theta \mathcal{L}\) aligns with those vectors, a standard per-coordinate optimizer wastes steps along the null directions. A spectral optimizer estimates a basis \(U\) of the top \(r\) singular vectors of the activation–gradient covariance, where \(r \ll d\), and applies an update that respects those directions.

The second insight is how to combine this rotated update with element-wise adaptivity. The Muon-inspired optimizers keep track of both a spectral direction and an RMS-like scale so the step's magnitude mirrors what Adam would have taken. Let \(g_t = \nabla_\theta \mathcal{L}(\theta_t)\) be the gradient at iteration \(t\). Adam would produce the update \(\Delta\theta^{\text{Adam}}_t = -\eta_t \frac{m_t}{\sqrt{v_t} + \epsilon}\) where \(m_t\) and \(v_t\) are the biased first and second moment estimates and \(\eta_t\) the learning rate. In spectral Muon, we rewrite the update as [ \Delta\theta_t = -\eta_t \left(U_t \Lambda_t U_t^\top\right) \tilde{m}_t, ] where \(\tilde{m}_t\) is a sign-stabilized version of the gradient (to avoid direction flips when RMS estimates are noisy), \(U_t \in \mathbb{R}^{d \times r}\) is the orthonormal basis capturing the stable directions, and \(\Lambda_t \in \mathbb{R}^{r \times r}\) is a diagonal matrix that scales each component by an RMS-aligned factor. The term \(U_t \Lambda_t U_t^\top\) is a low-rank projector built from the current activation sketch, and it replaces the identity matrix implicit in Adam. Each element of \(\Lambda_t\) is computed so that the magnitude of the spectral update matches the per-coordinate RMS of Adam, ensuring stability: if \(v_t\) is the RMS estimate, we set \(\lambda_{t,i} = \frac{1}{\sqrt{v_t + \epsilon}}\), except that Muon also stretches or shrinks \(\lambda_{t,i}\) to reflect the curvature captured by the sketch (Zhao et al. 2025) [https://arxiv.org/abs/2507.11005].

Constructing \(U_t\) cheaply without computing a full SVD is critical. Muon and similar optimizers use activation sketches: draw a random projection matrix \(S \in \mathbb{R}^{d \times k}\) and compute \(Y = A^\top (g \odot \text{sign}(g)) S\), where \(\odot\) denotes element-wise multiplication. The columns of \(Y\) capture the action of the gradient across the activation space, and an orthogonalization such as Newton–Schulz iterations on \(Y^\top Y\) yields \(U_t\). Newton–Schulz solves \(X^2 = Y^\top Y\) by iterating \(X_{i+1} = \frac{1}{2} X_i (3I - X_i^2)\), and because \(Y^\top Y\) is only \(k \times k\), the cost is negligible if \(k\) is small (often in the tens). This is where the geometry meets adaptivity: we only measure the curvature in the subspace that the current gradients occupy. The orthogonal component \(U_t\) rotates the gradient into this subspace, \(\Lambda_t\) rescales it by RMS and curvature, and the preimage \(U_t \Lambda_t U_t^\top\) brings us back to the full parameter space.

There is also a stability trick borrowed from Moons and SOAP-type methods that handles the low stable rank of activations. Because the sketch \(Y\) can collapse when the activation covariance is nearly low rank, Muon applies a sign stabilization: \(\tilde{g}_t = \text{sign}(g_t) \odot \text{max}(|g_t|, \sigma)\), where \(\sigma\) is a small floor that prevents zeroing out directions, and \(\text{sign}(g_t)\) ensures the update does not flip when the RMS estimate lags. This matches the intuition from the spectral gradient analysis: the ratio between the nuclear norm and the Frobenius norm of the gradient determines whether a rotation helps (Author et al. 2026) arxiv:2603.18168. When the ratio is high, the gradient has significant mass along fewer singular directions; rotation reduces variance. When it is low, the gradient is already isotropic, so the optimizer gracefully falls back to Adam-like scaling.

Implementing these ideas in code means updating the optimizer state twice per step. First, gather the activation sketch by multiplying the activations with the stabilized gradient. That gives the covariance that defines \(U_t\). Second, compute \(\Lambda_t\) by dividing the RMS moments by the spectral norms derived from the sketch. The per-element RMS acts as a thermostat—it keeps each spectral direction from growing too fast. In the original Muon paper, rescaling each spectral component to match Adam's RMS reduces the need for tedious learning rate tuning, and Zhao et al. (2025) report training speedups of over 40% across vision transformers, diffusion pipelines, and language models on the same hardware.

Because the updates are low-rank, the spectral optimizer is inexpensive and can share a single shared sketch across layers of similar dimension. Furthermore, the sketch can be updated infrequently: a layer whose activations have high stable rank might need re-orthogonalizing every few steps, while a low stable rank layer can reuse the same basis for longer, measured by tracking the Frobenius norm of the difference between successive sketches. This is the knob that future work aims to automate. For now, we add a scheduler that triggers orthogonalization only when the normalized trace of the sketch drops below a threshold.

Finally, the spectral optimizer interacts smoothly with other stability layers such as layer normalization and dropout because it does not assume any particular weight distribution; it simply rotates the gradient into a subspace defined by the activations, then rescales. When the activation stable rank shifts—for example, during warm-up when the weights change rapidly—the sketch naturally picks a new basis because the covariances move. The smooth interplay between the per-element RMS and the global geometry is what lets Muon achieve both adaptivity and fast convergence.

Where the field is now¶

The latest research paints a clear picture: when the gradient is compressible, spectral adaptivity wins. AdaMuon (Zhao et al. 2025) [https://arxiv.org/abs/2507.11005] put this into practice by combining the Muon sketch with sign-stabilized updates and RMS-aligned rescaling, reporting 40% faster steps on transformer pretraining and diffusion flows compared to AdamW baselines with the same steps per second. The benchmark for how much geometry helps comes from the theoretical work “When do spectral gradient updates help in deep learning?” (Author et al. 2026) [https://www.arxiv.org/pdf/2603.18168], which measures the nuclear-to-Frobenius norm ratio of gradients across layers and shows that when this ratio exceeds the stable rank of the activations, a spectral rotation yields a tighter bound on the expected loss decrease than any diagonal scaling can. These two works together explain why Muon splits the difference between efficient computation and Newton-like directionality.

Research is also stressing that this approach is not limited to transformers. The “Optimization Benchmark for Diffusion Models on Dynamical Systems” study (Author et al. 2026) [https://www.arxiv.org/pdf/2602.07145] re-ran diffusion model training with spectral optimizers and found that the geometry-aware updates improved sample quality on both continuous-time and discrete-time schedulers, even though diffusion gradients are notoriously noisy. This robustness led to a new class of optimizers, including SOAP-style updates and the sketching-based BASIS method. BASIS: Balanced Activation Sketching with Invariant Scalars for “Gh (Author et al. 2026) [https://arxiv.org/abs/2604.16324] extends the sketch by subtracting a bias term that tracks the moving activation mean, which suppresses direction drift in self-supervised pretraining. SOURCE’s engineering piece (Author et al. 2026) — “DDCL-INCRT: A Self-Organising Transformer with Hierarchical Prototype Structure” [https://arxiv.org/abs/2604.01880v1] — uses hierarchical prototypes whose gradients are naturally low stable rank, so the paper argues that only spectral updates make their layered self-organization tractable at scale.

On the engineering frontier, spectral adaptive optimizers are already entering production inference pipelines. Neural search companies deploy SOAP variants inside their transformers to reduce latency by eliminating zigzag gradient paths during fine-tuning, and the same sketching operations sweep through inference-time adapter training because they are lightweight compared to full second-order updates. The practical lesson is that these optimizers are the first general-purpose methods that can replace both AdamW and LAMB in large-scale distributed runs without extra hyperparameter tuning, so you can train the same NanoGPT or diffusion pipeline with fewer epochs and the same GPU count.

What's still open¶

Can we learn the orthogonalization frequency of the sketch itself? Right now, we trigger a Newton–Schulz iteration once every \(k\) steps or when a heuristic threshold of sketch drift is crossed. What if the optimizer monitored real-time activation stable rank and only re-orthogonalized when the rank suddenly rose or fell? Automating that could eliminate the remaining cost of periodic sketch recomputation, but it must avoid false positives that would kill training throughput.

Is there a provable way to interpolate between scalar RMS rescaling and full spectral projectors so we know when to trade one for the other? In regimes where the nuclear-to-Frobenius ratio matches the stable rank, the spectral update adds no benefit, yet existing Muon variants still pay for sketching. A formal criterion or a learned gating function that shuts off the projector in isotropic regions would let a single optimizer handle both extremes optimally.

Can spectral adaptivity be combined with preconditioners that already operate on weight matrices, such as Shampoo or K-FAC? Shampoo is costly because it stores \(O(d^2)\) matrices per layer, but its preconditioners and Muon’s sketches are both approximations of the curvature tensor. If one could share or re-use the same factors, we could achieve Shampoo-level conditioning at Muon-level cost.

Where to read next¶

If you want the probabilistic foundation, → Score matching explains why rotating the gradient space is equivalent to shaping the score function, and if you want to understand the conditioning side that Muon scrambles, → second order methods shows how Kronecker-factored curvature and Shampoo work with these low-rank sketches; for the engineering counterpart, → Distributed Training describes how to implement these optimizers across shards and keep the sketches synchronized without extra bandwidth.

Build it¶

Training a Muon-inspired adaptive optimizer head-to-head with AdamW on a character-level NanoGPT over Tiny Shakespeare exposes the geometry gap in under two hours on a Colab T4.

What you're building: a char-level NanoGPT trained with both AdamW and a simplified Muon optimizer (RMS-aligned spectral updates based on activation sketches), plus a comparison plot of validation loss / perplexity on Tiny Shakespeare.

Why this is valuable: it forces you to implement sketch generation, orthogonalization, and RMS rescaling so you actually touch the parts of Muon that break the compass analogy, and the resulting plot answers the practical question “does a geometry-aware rotation beat AdamW’s diagonal scaling on this dataset?”

Stack: - Model: karpathy/nanoGPT — community reference for char-level transformers - Dataset: tiny_shakespeare — 1MB text corpus, safe for Colab uploads - Framework: PyTorch 2.1 + accelerate 0.28 - Compute: single Colab T4 (16 GB VRAM), ~2 hours for 10 epochs of 512-token context

The recipe: 1. pip install torch torchvision accelerate einops matplotlib and clone karpathy’s NanoGPT transformer; use accelerate config to select the single T4 device. 2. Preprocess tiny_shakespeare: tokenize with the existing byte-pair tokenizer, pad to 512 tokens per example, and log the mean activation norms for each layer during a dry run. 3. Implement Muon: for each transformer block, maintain a sketch matrix \(Y_t = A^\top (\text{sign}(g_t) \odot \max(|g_t|, \sigma)) S\) where \(S \in \mathbb{R}^{d \times k}\) is a fixed random Gaussian matrix with \(k=16\), normalize \(Y_t\), apply one Newton–Schulz iteration to get \(U_t\), and set \(\Lambda_t\) so that \(\|\Lambda_t U_t^\top g_t\|\approx \|\frac{g_t}{\sqrt{v_t}+\epsilon}\|\). Train for 10 epochs with learning rate 3e-4 and compare to AdamW’s baseline. 4. Evaluate by computing validation cross-entropy per token every epoch, then plot both losses and the ratio of Muon update norm to AdamW update norm; target perplexity below 1.25 for Muon and a 5–7% lower loss than AdamW after 10 epochs. 5. What you now have is a paired training log and plot demonstrating how geometry-aware updates reduce zig-zagging, plus checkpointed weights from both optimizers that can be loaded into a text-sampling script.

Expected outcome: Artifact is the two checkpoints plus a pair of loss curves that show Muon’s spectral update producing lower perplexity on Tiny Shakespeare within the same training budget.

CS student: Run the same recipe on a free Colab A100 or RTX 4070 by reducing the context size to 256 tokens, halving the batch size to 16, and shrinking the sketch rank to \(k=8\); the smaller footprint lets you run the build within Colab’s session limit while still seeing the spectral update drift.
Applied engineer: Apply the Muon optimizer to meta-llama/Llama-3-7B-Instruct fine-tuning on a small custom dataset, quantize the final model to int4 with vLLM for inference, and measure p95 latency ≤ 120 ms at batch 1; the same sketching logic can live inside the LoRA adapter updates.
Applied researcher: Test the hypothesis that doubling the sketch rank from \(k=16\) to \(k=32\) improves convergence when the activation stable rank exceeds 20; falsify by plotting validation loss vs. sketch rank and declare failure if Muon’s loss does not drop by at least 1% relative to the baseline within 3 epochs.
Frontier researcher: Extend the orthogonalization frequency question from §What’s still open by implementing a trigger based on the real-time stable rank ratio; the falsifier criterion is that if the automated trigger changes the orthogonalization cadence by more than 30% compared to a fixed schedule, the validation loss should not exceed the fixed-schedule loss by more than 0.5% after 5 epochs.

If this build worked for you — a ⭐ on GitHub is the only signal we collect.