1: Introduction

Modern neural-network optimizers can largely be understood as different ways of preconditioning the gradient. Given a gradient $g$, an optimizer applies a transformation that changes the scale and orientation of the resulting update.

At one extreme, stochastic gradient descent applies no adaptive preconditioning, while Adam maintains a diagonal estimate of the gradient second moment. Full-matrix AdaGrad [3] instead maintains

and updates using

This full-matrix update is attractive because it accounts for correlations among parameters rather than treating each coordinate independently. By equalizing the magnitude of updates across different directions, it achieves optimal regret bounds in adversarial online learning environments.

Unfortunately, the full AdaGrad matrix grows quadratically with the number of parameters. It is intractable for modern neural networks.

Shampoo [4] makes full-matrix AdaGrad more practical by approximating its preconditioner with a Kronecker product. For a matrix-shaped parameter $W \in \mathbb{R}^{m \times n}$ with gradient $G$, Shampoo maintains smaller left and right factors rather than a single $(mn)\times(mn)$ matrix.

The classical construction uses row and column gradient covariances such as

These factors are then combined with matrix roots to precondition the gradient from both sides. This construction is intuitive, but not optimal. In particular, the standard covariance factors do not jointly provide both gradient whitening and scale invariance. Practical Shampoo variants therefore often introduce additional mechanisms, including Adam grafting, auxiliary diagonal statistics, eigenbasis updates, and other normalization rules [5].

1.1 The KL-optimal approximation

A theoretically cleaner approach is to choose the Kronecker factors as the best approximation to the full AdaGrad matrix with respect to a suitable divergence. Let the vectorized gradient second moment be

We seek factors $S_a$ and $S_b$ such that

Taking our divergence to be KL divergence, the resulting factors define the idealized KL Shampoo preconditioner.

The corresponding matrix update then has the form

This construction has two especially useful properties:

1. It approximately whitens the matrix gradient.
2. It is naturally scale-invariant.

In particular, unlike classical Shampoo and SOAP-style methods, it does not require a separate Adam statistic to determine its update scale.

Lin et al. [1] derived this idealized update and an online rule that moves the factors toward the KL-optimal solution. However, the paper argues that the idealized algorithm was too expensive and impractical because it requires fresh inverses and $-1/2$ roots at every training step. Computing these roots using a linear solver, eigendecomposition, Schur decomposition, or related dense linear-algebra routines is prohibitively expensive. Most QR-based Shampoo implementations therefore use a single QR factorization with power iteration to approximate an eigenbasis [5] [1]. The resulting basis is not mathematically exact since orthogonality alone does not make it an eigenbasis, and Appendix E develops this distinction and its resulting mathematical imprecision.

In this post, we show that with proper hardware-aligned numerical algorithm design, idealized KL Shampoo can be executed in reasonable wall-clock time; we call this optimizer Online KL Shampoo (OKLS). We find that OKLS achieves 1.59× Muon's parameter efficiency at the 1.2B Muon-equivalent scale, with the advantage increasing toward larger models [6] [7].

2: Making Idealized KL Shampoo Practical

For each parameter matrix, idealized KL Shampoo needs the inverse and $-1/2$ root for both the left and right preconditioners:

at every optimization step.

However, the only place the inverse term is used is in the preconditioner EMA update, for example

Notice that the inverse term can be rewritten using the inverse square root:

Thus, it suffices to compute an accurate matrix inverse square root every step using operations that map efficiently to modern accelerators.

2.1 Slowness of standard numerical methods

Eigendecomposition, linear solvers, and Schur decomposition are poor matches for GPU hardware because they require irregular control flow, have less favorable memory access patterns, and often demand higher precision for numerical stability. QR-based approaches mitigate some of these issues, but they remain expensive in practice, and their cost encourages implementations to refresh the eigenbasis only once every tens or hundreds of steps, which can introduce staleness into the optimizer.

GEMM-only methods are more promising because GPUs have specialized tensor cores that are designed to execute such operations at extremely high throughput. However, such methods are useful only if they converge quickly and may be executed in low precision.

2.2 Coupled Newton–Schulz

Given a symmetric positive-definite matrix $S$ (normalized so the largest eigenvalue is smaller than 1), coupled Newton–Schulz [8] maintains two sequences $Y_k$ and $Z_k$:

At each iteration,

As the iteration converges,

The coupled iteration is exceptionally stable, as its Frechet derivative is zero at the exact solution. Thus, it suppresses local numerical perturbations and works reliably in TF32.

Unfortunately, it converges slowly for small eigenvalues. To reach our desired accuracy over the relevant eigenvalue range, naive coupled Newton–Schulz requires roughly 21 iterations, corresponding to approximately 60 GEMMs after minor algebraic simplifications. At this point, it is no longer decisively faster than QR. It also suffers from a hardware disadvantage. Tensor cores support TF32 GEMM, but their throughput is substantially lower than that of native FP16 or BF16. Muon, by contrast, is often implemented with 16-bit GEMMs.

To make online KL Shampoo affordable, we thus need fewer iterations and native FP16 execution.

2.3 CANS: optimizing for a fixed iteration budget

Classical Newton–Schulz uses the same coefficients at every iteration, because these coefficients are designed for asymptotic convergence. They are thus not necessarily optimal when the iteration budget is fixed. The Chebyshev-optimized Newton–Schulz method, or CANS [2], instead chooses a sequence of polynomials tailored to a fixed number of steps and a target error tolerance.

The generalized update takes the form

The coefficients $(a_k,b_k)$ are selected so that the composed polynomial rapidly moves small eigenvalues into the quadratic convergence regime while controlling the worst-case final error. Although CANS was originally developed in the context of computing Muon's polar factor [2], the same coefficients apply to the coupled inverse-root iteration. In both cases, the scalar polynomials push the relevant singular values or eigenvalues toward one.

The 10-step CANS achieves a level of precision similar to that of the 21-step classical Coupled Newton-Schulz, and the resulting implementation requires 27 GEMMs. Figure 1 shows that the convergence of CANS Accelerated Coupled NS (10 steps) matches that of classical Coupled NS (21 steps), while classical Coupled NS (10 steps) fails to converge for small eigenvalues.

This resolves the algorithmic convergence problem, but not yet the low-precision execution problem.

2.4 Per-step scaling for FP16 execution

TF32 and FP16 both use a 10-bit mantissa. Their primary difference is exponent range. Consequently, when FP16 GEMMs accumulate into FP32, their precision is comparable to TF32 as long as the FP16 inputs do not overflow or underflow. The obstacle is therefore dynamic range, not mantissa precision.

During the CANS iteration, the eigenvalues of the intermediate matrices $Y_k$, $Z_k$, and $P_k$ (which bound the largest elements in the matrices) evolve according to deterministic scalar polynomials that can be solved offline, as demonstrated in Figure 2. We can therefore rescale each intermediate matrix so that its numerical representation remains safely within the FP16 representable range while fully utilizing its limited dynamic range to prevent underflow.

We store the intermediate matrices in their scaled form:

where the scaling factors are chosen from the known stepwise eigenvalue bounds. Substituting these scaled variables into the CANS recurrence produces updates

And finally:

All scale corrections can be folded into the scalar $\alpha$ and $\beta$ arguments of the GEMM kernels:

The matrix inputs are cast to FP16, while accumulation and persistent outputs remain FP32. At the final step, the normalization and unscaling constants are folded into the output coefficient, directly materializing the inverse square root in FP32.

The result is Scaled CANS Coupled Newton–Schulz:

Across the evaluated matrix sizes and condition numbers, the scaled FP16 method closely matches the error of the unscaled TF32 implementation (Figure 3).

By combining CANS with FP16 GEMMs, at $8192\times8192$ with batch size four, the benchmarks (Figure 4) report approximately 197 ms for Scaled CANS, compared with 337 ms for QR and 2571 ms for eigendecomposition on an H100. The full implementation in Python is available in Appendix D.

3: Full Online KL Shampoo Algorithm

With the inverse-root bottleneck removed, Online KL Shampoo (OKLS) becomes feasible for practical large-scale training. For each parameter matrix $\theta_t\in\mathbb{R}^{m\times n}$, the exact zero-staleness algorithm is as follows.

The AdamC weight-decay correction used above follows Defazio [20]. For discussion on memory usage, see Appendix F. The complete OKLS specification, the Adam and Muon baselines, and the shared initialization scheme are provided in Appendix C.

4: muP for Online KL Shampoo

In this section, to support hyperparameter transfer, we derive a width-scaling rule for OKLS based on the spectral norm of its whitened update.

Building on Yang et al.'s muP framework [9] and the spectral condition of Yang, Simon, and Bernstein [10], feature learning across widths and zero-shot hyperparameter transfer require the update $\Delta W$ to a linear layer $W\in \mathbb{R}^{d_\text{out}\times d_\text{in}}$ to satisfy

where $|| \cdot||_2$ is the spectral norm.

Thus, we need to compute the spectral norm of the raw update produced by OKLS to determine the scaling factor required to achieve it.

Consider the idealized full AdaGrad matrix case; the raw update is

which is an i.i.d. standard Gaussian vector if $g$ is zero-mean Gaussian, a reasonable approximation in a noise-dominated environment.

When reshaping the i.i.d. vector into a matrix, random-matrix theory predicts a spectral norm of

Therefore, if $U_t$ is the raw Online KL Shampoo update, we scale it by

The resulting parameter update is

Although real gradients are not exactly zero-mean Gaussian and the Kronecker approximation is not an ideal full-matrix whitening, the observed singular-value distributions of Online KL Shampoo updates remain close to the random-matrix prediction across most layers.

Figure 5 illustrates the singular value spectrum of the raw updates across training steps 1100 through 5000 for a model with a hidden dimension of 2048 and 24 layers, comparing the empirical singular values of the 24 transformer blocks against standard Gaussian matrix references. This confirms that the empirical updates are close to the idealized case and validates the muP rule we derived.

We will also empirically validate that this optimal learning rate transfers across width in Section 5.

5: Experiments

We pretrained transformers with gated attention [21] using Online KL Shampoo on Nemotron CC v2 [11], largely following the recipe from previous research posts for apples-to-apples. Training details are in Appendix B. All optimizers tested received equal-compute Bayesian hyperparameter searches to find their optimal configurations.

5.1 Training Results

For the scaling ladder, we train all models for 20B tokens.

At the loss reached by the 1.183B-parameter Muon model, Online KL Shampoo achieves approximately 1.59× Muon's parameter efficiency. The fitted multiplier increases toward larger models.

At 150M, we additionally compare against SOAP [5], the strongest published Kronecker-factored baseline for our setting. Online KL Shampoo reaches a terminal loss of 2.6809 against SOAP's 2.7048, a margin of 0.024 nats, while keeping the smaller optimizer state reported in Appendix F.

For the production run, we train a 1.45B total-parameter model, with 1.183B non-embedding transformer-body parameters, on 100B tokens. The loss curves are shown below.

Notably, OKLS exhibits a qualitatively different loss curve from Muon, with worse early convergence but better final results over the decay. OKLS achieves a final gain of approximately 0.05 nats over Muon. We leave the diagnosis of the unique shape of OKLS' loss curve to future work.

OKLS-1.45B improves upon our previous approaches on the token-efficiency frontier for small-scale pretraining evaluations. It exceeds Aurora [12] on all four displayed benchmarks and Qwen3-1.7B [13] on HellaSwag, despite no changes to the data mixture or architecture. The evaluations use an earlier matched checkpoint pair whose terminal losses are within 0.003 nats of the matched pair in Figure 8.

For MMLU and Commonsense QA in particular, Online KL Shampoo improves eval scores significantly. This is consistent with the spectral-coverage hypothesis discussed in Section 5.2, although the evaluation gains do not identify a causal mechanism.

5.2 Mechanisms of OKLS

We then sought to uncover the mechanisms by which OKLS yields pretraining gains. We analyze stable rank, a continuous measure of spectral concentration defined as the squared Frobenius norm divided by the squared spectral norm:

Stable rank increases when spectral energy is spread across more singular directions. It measures spectral coverage, not memory capacity by itself.

When comparing the two production runs, we find that Online KL Shampoo consistently has a higher stable rank across parameters and activations.

We observe that the residual-stream stable rank is higher for all layers after the first under OKLS. Among parameters, the largest gains include the attention value projections and MLP matrices, which prior work identifies as associative-memory parameters [14] [15].

Wang et al. [16] among others describe the connection between higher-effective-rank spectra in value-output and FFN weights and factual recall (though their metric is the slightly different entropy-based effective rank). Kim et al. [19] show in a power-law linear associative-memory model that one Muon step stores $\widetilde{\Theta}(d^{1+1/(2\alpha)})$ associations versus $\widetilde{\Theta}(d^{1/(2\alpha)})$ for SGD by amplifying bulk singular directions.

Together, these results suggest that spectrally broader updates can produce higher-rank memory weights, which may improve factual storage. Our stable-rank measurements and performance on memorization-based benchmarks are consistent with that mechanism, but causally establishing whether increasing stable rank improves factual memory remains future work.

5.3 Ablations

In this section, we present four ablations we ran for the core components of our OKLS recipe. They are presented in a joint panel below.

OKLS muP: We verify that our new muP rule transfers the learning rate from 78.7M to 295.8M non-embedding parameters.

Matrix Root Strategy: We compare matrix-root strategies at a fixed 10-step budget. Unscaled FP16 Newton–Schulz diverges immediately, while naive coupled Newton–Schulz remains finite but fails to converge. Scaled CANS is the only Newton–Schulz variant that trains stably in mixed precision.

NS vs QR: We find that Online KLS and online QR converge to similar loss, while infrequent QR (the conventional implementation) performs slightly worse. Online KLS is significantly faster and directly approximates the inverse root with an explicit residual rather than treating one QR factor as an exact eigenbasis (Appendix E).

Stale Preconditioner Updates: Across two independently swept run populations, stale preconditioners are associated with substantially more unstable runs.

6: Discussion

Above, we showed that OKLS can significantly improve language-model training. We now discuss two implications: how to execute it at distributed scale, and why online AdaGrad better explains its behavior than slow curvature estimation. Detailed optimizer-state accounting is in Appendix F, the distributed ownership design is in Appendix G, and the curvature-versus-AdaGrad argument is developed in Appendix H.

6.1 Distributed execution without persistent optimizer-state GPU memory

Although lighter than other Kronecker product based optimizers, OKLS maintains larger optimizer states than simpler methods such as Muon. To make these states practical at distributed scale, we developed an ownership-based optimizer framework using one-sided communication and CPU offloading.

Each parameter is assigned to an owner GPU, which maintains the complete optimizer state for that parameter. During the optimizer step, the owner gathers the required gradient shards, computes the update, and scatters the resulting update shards back to the participating GPUs. One-sided communication allows these transfers to proceed without requiring collective synchronization or communicating optimizer states.

Besides, because the arithmetic cost of the matrix operations grows as $O(d^3)$, while transferring the optimizer state grows as $O(d^2)$, state movement can be overlapped with computation for reasonably large matrices. This allows the persistent optimizer state to be stored in CPU memory without increasing the critical-path latency in the regimes we evaluate.

6.2 OKLS behaves more like online AdaGrad than slow curvature estimation

Shampoo-style optimizers are often described as curvature approximations. This perspective motivates slowly changing preconditioners, infrequent decomposition updates, and second-moment EMA coefficients close to one. However, our results suggest a different interpretation. Across our hyperparameter sweeps, Online KL Shampoo performs best when

This differs from the conventional setting, in which $\beta_2$ is often substantially larger than $\beta_1$. The $\beta_1^2 = \beta_2$ setting is not a random point; instead, it can be motivated from an online optimization perspective. In the scalar case, setting $\beta_1^2 = \beta_2$ has been proven to yield an optimal regret bound when facing an oblivious adversary, which means the adversary fixes the sequence of gradients regardless of the optimization algorithm's choices. [18] In the multivariable case, this setting allows the optimizer to act like a full-matrix AdaGrad, which Online KL Shampoo is designed to simulate.

With $\beta_1^2 \approx \beta_2$, the optimizer behaves more like an exponentially windowed version of online AdaGrad than a slowly evolving curvature estimator.

This interpretation is also consistent with our staleness ablations. If the preconditioner represented slowly changing curvature, delaying it by a single step should have little effect. Instead, one-step staleness can substantially destabilize training, suggesting that the preconditioner must remain synchronized with the current gradient sequence.

The weighted-AdaGrad derivation and gradient-spike analysis are developed in Appendix H.

7: Inflection

Idealized KL Shampoo offers a principled, scale-invariant approximation to full-matrix AdaGrad. However, its requirement for fresh matrix inverse square roots at every step limited its practicality. Via Scaled CANS Coupled Newton–Schulz, Online KL Shampoo achieves zero-staleness preconditioning at a fraction of the cost of traditional dense linear algebra routines. In our experiments, this translates to a 1.59× parameter-efficiency advantage over Muon at the 1.2B Muon-equivalent scale, with the advantage increasing toward larger models, while maintaining roughly 98% of the training throughput.

Five things excite us about this work:

1. Fresh dense preconditioning is practical. Scaled CANS replaces a stale, high-precision decomposition with 27 hardware-aligned FP16 GEMMs and produces a new inverse root every step. This changes zero-staleness Shampoo from an idealized algorithm into a practical training primitive.

2. Clean scaling-law advantage. Online KL Shampoo improves the scaling-law frontier from 162.9M to 1.183B non-embedding parameters, outperforms Adam, Muon, and SOAP at 150M total parameters, and establishes a new token-efficiency frontier beyond Muon and Aurora in the 1.45B, 100B-token run.

3. Nontrivial learning dynamics. OKLS benefits from a long decay phase, and its late crossover suggests qualitatively different learning dynamics from current optimizers. Understanding this behavior is an important direction for future work.

4. The optimizer behaves like online learning, not slow curvature estimation. The preferred relationship $\beta_2\approx\beta_1^2$ and the sensitivity to one-step staleness both point toward a time-weighted AdaGrad interpretation of OKLS. The results suggest that synchronization with the current gradient sequence is a core part of the algorithm.

5. Spectral coverage offers a testable mechanism. Online KL Shampoo produces broader parameter and activation spectra in the trained models and substantially improves factual benchmarks. Prior work connects spectral breadth to associative-memory capacity in specific optimizer, model, and data settings, but this relationship has not been causally demonstrated in the general case. We see this as a promising direction for future work.

@article{zhang2026onlineklshampoo,
  title   = {Online KL Shampoo},
  author  = {Zhang, Ashley and Keigwin, Ben and Pai, Dhruv and Dewulf, Alec},
  year    = {2026},
  url     = {https://tilderesearch.com/blog/online-kl-shampoo}
}

Appendices

A: Open-Source Release

Our implementation of Online KL Shampoo is open-sourced at github.com/tilde-research/online-kl-shampoo-release. The repository includes the optimizer, the Scaled CANS Coupled Newton-Schulz inverse-root implementation, and a minimal usage example.

B: Training Details

We report the exact settings used for our runs. We trained on fully open-source internet data from NVIDIA Nemotron CC v2 [11], broadly following the recipe used in our previous posts. Hyperparameters were selected by Bayesian sweeps, with comparable sweep compute spent on each optimizer.

The optimizer hyperparameters below apply only to the hidden matrix parameters. Muon uses a single momentum coefficient rather than separate first- and second-moment EMAs.

The table above gives the settings for the 20B-token runs. For the 100B-token runs, the learning rate was scaled by $5^{-1/4}$ [22] and the warmup was set to 125 steps.

C: Optimizer Algorithms

We specify Online KL Shampoo, the Adam and Muon baselines, and the initialization scheme shared by all three optimizers in the same format. ScaledCANS is defined in Appendix D.

D: Scaled CANS Coupled Newton-Schulz Code

L ← 1.01 · UpperBound(S)
Y₀ ← S / L,  Z₀ ← I
for k = 0, …, 9:
    Pₖ ← ZₖYₖ
    Qₖ ← aₖI + bₖPₖ
    Yₖ₊₁ ← YₖQₖ,  Zₖ₊₁ ← QₖZₖ
return sym(Z₁₀) / √L

Per-step rescaling is algebraically folded into the coefficients below, keeping FP16 intermediates in range without changing the iteration.

"""
Pure-torch 10-step coupled Newton-Schulz for S^{-1/2} via CANS polynomials.
"""

import torch
import math

# ─────────────────────── CANS polynomial coefficients ───────────────────────
CANS_COEFFS = [
    (5.182503604966906, -5.178098480082684),
    (2.586120737395915, -0.6479542005271643),
    (2.567364126726186, -0.6454968804392178),
    (2.520560084348265, -0.6393528082067044),
    (2.410759275435182, -0.6248683598710716),
    (2.1883348130094173, -0.5952022073798908),
    (1.8595760874873613, -0.5504490972723968),
    (1.589020160467417, -0.5126569802066718),
    (1.5051653981684994, -0.5007377068751799),
    (1.5, -0.5),
]

# ──────────────── Per-step safe-scale ceilings (eigs in [0, 1]) ─────────────
_M = 16384.0
_Z_MAX = [
    1.0, 5.183, 13.403, 34.409, 86.731, 209.09,
    457.55, 850.85, 1352.0, 2035.0, 3052.5,
]
_Y_MAX = [1.0, 1.297, 1.726, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 1.0, 1.0]
_P_MAX = [1.0, 3.98, 3.95, 3.88, 3.71, 3.32, 2.60, 1.73, 1.15, 1.008, 1.0]
_S_Z = [_M / z for z in _Z_MAX]
_S_Y = [_M / y for y in _Y_MAX]
_S_P = [_M / p for p in _P_MAX]

# ── Step-0 precomputed constants ──
_ALPHA_Y_0 = CANS_COEFFS[0][1] * _S_Y[1] / (_S_Y[0] ** 2)
_BETA_Y_0 = CANS_COEFFS[0][0] * _S_Y[1] / _S_Y[0]
_Z_SCALE_0 = CANS_COEFFS[0][1] * _S_Z[1] / _S_Y[0]
_Z_DIAG_ADD_0 = CANS_COEFFS[0][0] * _S_Z[1]

# ── Steps 1–9 precomputed constants ──
_ALPHA_P = []
_ALPHA_Y = []
_BETA_Y = []
_ALPHA_Z = []
_BETA_Z = []
for _k in range(1, 10):
    _a, _b = CANS_COEFFS[_k]
    _ALPHA_P.append(_S_P[_k] / (_S_Z[_k] * _S_Y[_k]))
    _ALPHA_Y.append(_b * _S_Y[_k + 1] / (_S_Y[_k] * _S_P[_k]))
    _BETA_Y.append(_a * _S_Y[_k + 1] / _S_Y[_k])
    _ALPHA_Z.append(_b * _S_Z[_k + 1] / (_S_Z[_k] * _S_P[_k]))
    _BETA_Z.append(_a * _S_Z[_k + 1] / _S_Z[_k])

# Absorb _S_Z[10] into last-step coefficients so runtime only needs 1/√L
_ALPHA_Z[8] /= _S_Z[10]
_BETA_Z[8] /= _S_Z[10]

# ────────────────────────────── Helpers ──────────────────────────────────────
def _estimate_max_eigenvalue(A: torch.Tensor) -> torch.Tensor:
    """Strict upper bound via min(Wolkowicz-Styan, Minc-Sainte-Marie). Cost: O(n²)."""
    n = A.size(-1)
    diag = A.diagonal(dim1=-2, dim2=-1)
    m = diag.sum(dim=-1) / n
    sq_norm = torch.sum(A**2, dim=(-2, -1))
    s_sq = torch.clamp((sq_norm / n) - (m**2), min=0.0)
    ws_bound = m + torch.sqrt(s_sq) * math.sqrt(n - 1)

    abs_A = torch.abs(A)
    d = torch.sum(abs_A, dim=-1)
    d_clamped = torch.clamp(d, min=1e-12)
    y = torch.einsum("...ij,...j->...i", abs_A, d_clamped)
    minc_bound = torch.max(y / d_clamped, dim=-1).values

    return torch.minimum(ws_bound, minc_bound)

# ────────────────────────── Main function ────────────────────────────────────
def scaled_cans_coupled_ns(
    S: torch.Tensor,
) -> torch.Tensor:
    """
    Computes S^{-1/2} via 10-step coupled Newton-Schulz with CANS polynomials.

    Pure torch implementation — FP16 GEMM inputs with FP32 accumulation.

    Args:
        S: (B, d, d) symmetric positive-definite matrix.
    Returns:
        W: (B, d, d) approximate S^{-1/2} in S.dtype.
    """
    B, d, _ = S.shape
    orig_dtype = S.dtype
    S = S.float()

    # ── Normalize: eigenvalues of Y₀ ∈ (0, 1] ──
    L = _estimate_max_eigenvalue(S) * 1.01                # (B,)
    Y = S * (_S_Y[0] / L.view(B, 1, 1))

    # ── Step 0: Z₁ = s_z₁·(a₀I + b₀Y₀),  Y₁ = s_y₁·(a₀Y₀ + b₀Y₀²) ──
    Z = Y.mul(_Z_SCALE_0)
    Z.diagonal(dim1=-2, dim2=-1).add_(_Z_DIAG_ADD_0)
    Y_half = Y.half()
    Y = torch.baddbmm(Y, Y_half, Y_half, torch.float32, beta=_BETA_Y_0, alpha=_ALPHA_Y_0)

    # ── Steps 1–8 (full coupled update) ──
    for k in range(8):
        Y_half = Y.half()
        Z_half = Z.half()
        # P = alpha_p · (Z @ Y)   (alpha fused into GEMM, FP16 output)
        P_half = torch.baddbmm(Z_half, Z_half, Y_half, beta=0.0, alpha=_ALPHA_P[k])
        # Y_{k+1} = alpha_y · (Y_half @ P_half)  +  beta_y · Y
        # Z_{k+1} = alpha_z · (P_half @ Z_half)  +  beta_z · Z
        Y, Z = (
            torch.baddbmm(Y, Y_half, P_half, torch.float32, beta=_BETA_Y[k], alpha=_ALPHA_Y[k]),
            torch.baddbmm(Z, P_half, Z_half, torch.float32, beta=_BETA_Z[k], alpha=_ALPHA_Z[k]),
        )

    # ── Step 9: Z only ──
    Y_half = Y.half()
    Z_half = Z.half()
    P_half = torch.baddbmm(Z_half, Z_half, Y_half, beta=0.0, alpha=_ALPHA_P[8])
    inv_scale = torch.rsqrt(L).view(B, 1, 1)             # (B, 1, 1)
    W = torch.baddbmm(Z, P_half, Z_half, torch.float32, beta=_BETA_Z[8], alpha=_ALPHA_Z[8])
    W.mul_(inv_scale)
    W = (W + W.mT) / 2.0
    return W.to(orig_dtype)

E: One-Shot QR Is Not an Exact Eigenbasis

In most QR-based Shampoo implementations, "QR" means a single factorization used to orthogonalize the current basis update at each refresh, not running the classical QR iteration until it converges to an eigenbasis.

For a symmetric preconditioner $A=V\Lambda V^\top$, an exact eigenbasis is an orthogonal matrix $V$ that diagonalizes $A$:

A one-shot basis refresh instead factors some update matrix $B$ as $B=QR$ and retains $Q$. The factorization is exact and $Q$ is orthogonal, but this does not imply that $Q^\top A Q$ is diagonal. A single QR factor therefore does not generally recover the eigenbasis of the preconditioner.

The iterative QR algorithm, by contrast, repeatedly factors the transformed preconditioner and applies similarity updates:

The accumulated transform $\bar Q_k=Q_0Q_1\cdots Q_k$ converges toward an eigenbasis under the usual assumptions. At any finite iteration count, however, it remains an (often poor) approximation. For the unshifted iteration, convergence of adjacent eigendirections is governed by ratios of eigenvalue magnitudes, with a characteristic error factor

Consequently, there is no uniform accuracy guarantee for a fixed number of QR iterations. As adjacent eigenvalues approach one another, this ratio approaches one, and the iteration can make arbitrarily little progress. For any fixed iteration budget, a sufficiently clustered spectrum can leave the estimated basis nearly as inaccurate as its initial orientation.

In particular for QR-based Shampoo and SOAP implementations:

• A one-shot QR factorization enforces orthogonality but not diagonalization.
• Iterated QR is an eigensolver, but its finite-step error depends on spectral separation.
• Infrequent basis refreshes introduce an additional and separate source of temporal staleness.

While OKLS' Scaled CANS is still a finite-precision polynomial approximation, its mathematical target is precise. It directly approximates $W=A^{-1/2}$, and its error is measurable without introducing a provisional eigenbasis:

A one-shot QR factor is exact only as a factorization and can be a poor eigenbasis approximation. Scaled CANS instead directly targets the required inverse-root equation with an explicit, measurable residual.

F: Memory Usage for Kronecker Product Optimizers

While Kronecker product optimizers must maintain left and right preconditioners ($m \times m$ and $n \times n$), we can reduce their memory footprint using Symmetric Packing. Instead of allocating full $d \times d$ dense tensors, we only store the upper triangular elements in a flattened 1D array of size $d(d+1)/2$ to slice the memory overhead of the Kronecker state exactly in half.

Specifically, all the optimizer states in Online KL Shampoo, except the momentum, $S_a$, $S_b$, $P_a$, and $P_b$, are SPD matrices, in contrast to QR-based methods, which maintain non-symmetric eigenbases.

Assuming a parameter matrix of shape $m \times n$, we can quantify the number of state elements each optimizer must persist in memory, as shown below. We assume symmetric packing is applied wherever mathematically possible.

By using direct root-finding rather than QR decomposition, Online KL Shampoo eliminates the need to store dense orthogonal matrices, resulting in minimal memory overhead compared to existing Kronecker product optimizers.

Zero-Delay CPU Offload

Although maintaining the optimizer state in GPU memory is traditionally preferred for speed, the arithmetic intensity of Online KL Shampoo allows us to offload this state to CPU memory via PCIe without incurring any latency penalties during training.

This zero-delay offload relies on the hardware scaling law: as matrix sizes grow, the computational requirement of GEMMs (scaling at $O(d^3)$) outpaces the memory bandwidth required to transfer them (scaling at $O(d^2)$), allowing us to hide communication behind computation.

Assuming we are operating on an H100 (989 FP16 TFLOPs) with a $d \times d$ matrix, the theoretical computation latency for 60 FP16 GEMMs (which is required per Online KL Shampoo step) is

Utilizing the symmetric packing strategy, the optimizer states contain $3d^2$ elements. With FP32 optimizer states and 64 GB/s bandwidth with PCIe Gen5 x16, the loading to GPU latency is:

Thus, for $d \ge 1546$, $T_{\text{transfer}}<T_{\text{compute}}$, enabling zero-delay CPU offload with proper pipelining. In reality, the exact number could differ because 6 GEMMs in the preconditioner EMA and whitening are in TF32 instead of FP16; the GEMMs and transfer cannot fully saturate the TFLOPS and PCIe, etc. But for reasonably large models, where the optimizer's state memory matters most, we comfortably achieve massive reductions in GPU memory usage without incurring additional latency.

G: Distributed Training Design

While the full implementation details for our distributed framework will be covered in a future blog post, we want to briefly outline the design that allows us to execute Online KL Shampoo with zero cross-GPU communication for optimizer states.

At a high level, the framework avoids sharding the optimizer state by assigning ownership of each weight to a specific GPU. Here is how the pipeline operates during the optimizer step:

1. Ownership System: Every individual weight is assigned to a single GPU. This owning GPU stores all optimizer state for that parameter locally and permanently, without sharding.
2. The Gather Phase: During the optimizer step, the owning GPU gathers the gradients from all FSDP shards (FSDP) or Expert Holders (EP).
3. Computation and Scatter: The owning GPU updates the optimizer state, computes the update, and then scatters the update shards directly to all GPUs that require them.
4. Local Application: Each GPU subsequently applies the update shards locally.

To completely eliminate synchronization delays, all of these gather and scatter operations utilize NVSHMEM one-sided peer-to-peer operations, and the communication and computation are overlapped via CUDA streams.

H: Rethinking the AdaGrad Family: Are Preconditioners Curvature or AdaGrad Matrices?

Shampoo-style optimizers are often referred to as "second-order" or "curvature-based" optimization methods. However, claiming that these optimizers genuinely capture the curvature of the loss landscape is theoretically flawed for two primary reasons:

• The Empirical Fisher Approximation is Flawed: Many of these methods use the empirical Fisher information matrix, which differs from the true Fisher information matrix and does not generally recover the natural gradient. It is not guaranteed to capture useful second-order information and can produce badly scaled or misdirected updates. [17]
• The Preconditioner Root Mismatch: True second-order methods (like Newton's method or Natural Gradient Descent) precondition the gradient using the inverse of the Hessian matrix (the $-1$ root). In contrast, Shampoo-style algorithms use the $-1/2$ root, which aligns with the AdaGrad root rather than the Newton root.

This distinction is not merely conceptual; it has important implications for how we understand practical trade-offs, algorithm design, and hyperparameter tuning.

For example, a pervasive belief in the Shampoo-style optimizer community is that the curvature of the loss landscape, approximated by the preconditioners, evolves slowly, allowing you to update Shampoo-style preconditioners infrequently (e.g., every few hundred steps) with little performance loss. Also, the optimal $\beta_2$ should be close to 1 (e.g., 0.99), or at least $\beta_2 \ge \beta_1$.

However, as our ablation on preconditioner staleness demonstrated in Section 5.3, even a single-step delay in updating the preconditioner severely destabilizes training. If the preconditioner were genuinely capturing a slow-moving curvature, a one-step delay would be mathematically imperceptible.

This conceptual mismatch extends to our extensive Bayes search for optimal KL-Shampoo hyperparameters, where we found that the classic setting of $\beta_1 < \beta_2$ generally failed to produce the best results.

Instead, the lowest-loss runs (represented by green data points) cluster around the dotted line denoting $\beta_2 = \beta_1^2$ rather than the diagonal $\beta_2=\beta_1$. This is a two-dimensional projection of a larger Bayesian sweep, so $\beta_1$, learning rate, weight decay, and other covarying hyperparameters also contribute to the observed loss differences.

The $\beta_1^2 = \beta_2$ setting is not a random point; instead, it can be motivated from an online optimization perspective. In the scalar case, setting $\beta_1^2 = \beta_2$ has been proven to yield an optimal regret bound when facing an oblivious adversary, which means the adversary fixes the sequence of gradients regardless of the optimization algorithm's choices. [18] In the multivariable case, this setting allows the optimizer to act like a full-matrix AdaGrad, which Online KL Shampoo is designed to simulate.

True AdaGrad, in online optimization, maintains an unweighted, infinite-memory accumulator of past gradient outer products, defined as

and scales updates using its inverse square root

This unweighted accumulation means the preconditioner grows indefinitely, causing the effective step size to decrease strictly monotonically over time until learning eventually halts.

To address it, AdaGrad can be transformed into a weighted version by an exponential discount factor. By defining a time-decayed gradient as $\tilde{g}_i = \beta^{t-i} g_i$, the update becomes

Notice that if we apply EMA accumulation for the momentum and preconditioner, respectively, with $\beta_1$ and $\beta_2$,

then

By setting $\beta_1=\beta$, $\beta_2 = \beta^2$, you recover the weighted AdaGrad update.

This result suggests that OKLS performs best when it behaves like a time-weighted full-matrix AdaGrad, which has a distinct dynamic from curvature estimation. It also explains why even 1-step staleness is fatal in the ablation: 0 staleness is a strict requirement for AdaGrad to give a valid regret bound.

Interestingly, the $\beta_1^2=\beta_2$ setting also provides a guarantee: the update magnitude when a gradient spike hits does not explode to infinity or shrink to 0. When an abnormally large gradient $g_t$ arrives, its massive magnitude $K$ dominates the historically accumulated statistics. If $g_t$ is large enough, both accumulators effectively collapse to represent only the current step: the momentum reduces to $(1-\beta_1) g_t$, and the uncentered covariance reduces to $(1-\beta_2) g_t g_t^\top$. If we separate the massive scalar $K$ from the underlying gradient direction $v$ such that $g_t = K v$, the raw update becomes proportional to

The $K$ magnitudes perfectly cancel out, and because $\beta_2=\beta_1^2$, the remaining scalars simplify cleanly to a stable constant $\sqrt{\frac{1-\beta_1}{1+\beta_1}}$, leaving the update magnitude strictly bounded and unchanged from the stable state.