Skip to content
Bud's Wiki

Online KL Shampoo (OKLS) — Tilde Research's KL-optimal full-matrix AdaGrad approximation optimizer with 1.45× parameter efficiency over Muon

TL;DR

Section titled “TL;DR”

Online KL Shampoo (OKLS) is a Tilde Research optimizer that closes the gap between diagonal preconditioners (which ignore gradient-coordinate correlations), Muon (which orthogonalizes each step but has no history of correlations), and full-matrix AdaGrad (which captures correlation history but demands quadratic state). OKLS approximates full-matrix AdaGrad by KL-optimal Kronecker factors, whitening matrix gradients across both row and column directions while remaining scale-invariant. The main practical obstacle — computing fresh inverse-square-root preconditioners every step, since even one-step staleness destabilizes training — is solved with Scaled CANS Coupled Newton–Schulz: 10 iterations, 27 FP16 GEMMs, FP32 accumulation. Reported result: 1.45× the parameter efficiency of Muon at 98% of Muon’s throughput, fit across 200M–1B non-embedding-parameter models trained for 20B tokens each.

Key claims

Section titled “Key claims”
  • Diagonal optimizers (Adam family) ignore correlations between gradient coordinates; full-matrix AdaGrad captures the geometry but requires quadratic state; Muon considers correlations at each step but not their history — OKLS is positioned as the first practical optimizer that keeps both the correlation geometry and its history [tweet §1].
  • OKLS approximates full-matrix AdaGrad by KL-optimal Kronecker factors that whiten matrix gradients along both row and column directions, while remaining naturally scale-invariant [tweet §2].
  • The main challenge is zero-staleness preconditioning: even one-step staleness in the inverse-square-root preconditioners destabilizes training, so the preconditioner must be recomputed fresh every step [tweet §3].
  • Scaled CANS Coupled Newton–Schulz makes zero-staleness preconditioning practical: 10 iterations, 27 FP16 GEMMs per step, with FP32 accumulation for numerical stability [tweet §3].
  • Headline scaling-law claim: median 1.46× parameter-efficiency advantage for OKLS over Muon, fit across models from 79M to 1.2B non-embedding parameters, each trained for 20B tokens [linked X article figure caption, cover image].
  • OKLS retains 98% of Muon’s training throughput — the KL-optimal Kronecker approximation + coupled Newton–Schulz kernel is not a wall-clock cost [tweet §4].
  • Concrete restatement of the parameter-efficiency claim: across 200M–1B models, an OKLS model matches a Muon model roughly 1.5× larger at equal training compute [tweet §4].

Method

Section titled “Method”

The setup is Tilde’s usual optimizer framing: name an object the current optimizer mis-controls, derive the correction, verify empirically. Here the mis-controlled object is the gradient’s covariance history across coordinates. Diagonal preconditioners (Adam) approximate the AdaGrad covariance ∑tgtgt⊤\sum_t g_t g_t^\top by its diagonal, throwing away all off-diagonal correlations. Muon controls per-step spectral geometry (via Newton–Schulz orthogonalization of the current gradient matrix) but does not accumulate any history of that geometry — every step starts fresh. Full-matrix AdaGrad keeps the full covariance but the state is quadratic in parameter count, which is infeasible at scale.

OKLS’s compression is Kronecker factorization of the covariance under a KL-divergence objective, giving two much smaller row and column factors instead of the full matrix. The KL-optimal choice — rather than Frobenius-optimal (Shampoo) or other norms — is presented as the key algorithmic novelty over Shampoo-family preconditioners.

The kernel-side story is Scaled CANS Coupled Newton–Schulz, an inverse-square-root routine designed to hit zero staleness (recomputed every step) at usable cost. The concrete kernel budget — 10 iterations × ~27 FP16 GEMMs per step, with FP32 accumulation — is small enough that overall training throughput stays within 2% of Muon.

The launch tweet references a longer X-article writeup (https://x.com/i/article/2082164286638768289) that hosts a scaling-law figure across 79M–1.2B non-embedding-parameter models at 20B-token isoflop. The article was behind an X login wall at filing time, so the exact fit form (Chinchilla additive vs Kaplan multiplicative), the parameter-efficiency scaling exponent, and the ablation set (partial Kronecker updates, staleness schedules, kernel-iteration count) are not verifiable from the tweet alone.

Results

Section titled “Results”
  • Parameter efficiency: median 1.46× advantage over Muon, fit across a 79M → 1.2B non-embedding-parameter scaling law, 20B tokens per point [linked X-article figure].
  • Throughput: OKLS retains 98% of Muon’s training throughput despite the every-step inverse-square-root recompute [tweet §4].
  • Restated efficiency claim: across 200M–1B models, an OKLS run matches a Muon run of ~1.5× the parameter count [tweet §4].
  • No numbers reported at filing for: absolute validation loss, downstream benchmark performance, weight-decay coupling, LR-transfer behavior across width, or interaction with QK-Norm / µP.

Why it’s interesting

Section titled “Why it’s interesting”

OKLS is Tilde Research’s third optimizer release in the ~2-month window this wiki tracks, sitting alongside Aurora: A Leverage-Aware Optimizer for Rectangular Matrices (row-leverage-aware Muon for tall MLP matrices) and Compositional Muon — partner-whitened optimizer that controls composed functional updates in attention (Compositional Muon — partner-whitened updates for QK/OV compositions). All three are pathology-first derivations, but OKLS attacks a different pathology from either: not per-matrix geometry (Aurora, Muon), not composed-object geometry (Compositional Muon), but the history of gradient correlations that Muon deliberately throws away by re-orthogonalizing every step. That makes OKLS the first Tilde optimizer that’s directly comparable to the Shampoo / SOAP / K-FAC line — microsoft/dion — distributed Muon, Dion2, Dion, and NorMuon orthonormal optimizers ships Dion and Dion2 as distributed Kronecker-factored variants, but Dion uses Frobenius-optimal factors; OKLS’s KL-optimal factors are the novel algorithmic claim, and the 1.45× parameter-efficiency headline is the first public comparison against Muon at ≥1B scale.

The zero-staleness requirement is the load-bearing empirical claim, and it complements two other 2026 threads. First, How Muon Lost Its Geometry shows that Muon variants’ per-matrix scaling factors have quietly diverged from what µP prescribes; OKLS re-imports geometry that Muon deleted (correlation history) and pays for freshness via a hand-tuned kernel, so it re-opens the “how much geometry does the optimizer need to explicitly control” fork Training stability at scale tracks. Second, the 10-iteration coupled Newton–Schulz kernel with FP16 GEMMs + FP32 accumulation sits in the same design space as the Muon-orthogonalization kernels that Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon‘s Gram-Newton–Schulz work targets — both papers are effectively saying that per-step inverse-square-root is now a practical primitive for optimizers, which was not true 12 months ago. The clean follow-up is whether OKLS composes with Aurora / Compositional Muon (each attacks a non-overlapping pathology), and whether the 1.45× parameter-efficiency headline survives the same critique Muon is Not That Special: Random or Inverted Spectra Work Just as Well leveled at Muon — that geometric preconditioners may look strong against Adam but tie randomized/inverted-spectrum baselines.

The absence of load-bearing verifiable details at filing (fit form, exponent, ablations, wall-clock numbers vs the 98% claim) is the main risk. The launch tweet is the artifact; the linked longform is login-walled. A repo / arxiv follow-up should be appended as a pointer when it lands.

See also

Section titled “See also”