Online KL Shampoo

Online KL Shampoo (OKLS) is a zero-staleness, Kronecker-factored optimizer that approximates full-matrix AdaGrad at language-model scale. It updates the KL-optimal left and right covariance factors, computes fresh inverse-square-root preconditioners, and applies the whitened update within the same training step.

The implementation combines:

• KL-optimal Kronecker preconditioning across both matrix axes.
• Scaled CANS Coupled Newton–Schulz, a 10-step inverse-root method using 27 FP16 GEMMs with FP32 accumulation.
• Zero-staleness updates: covariance factors and their inverse roots are produced and consumed in the same step.
• muP shape scaling, Nesterov momentum with variance correction, and AdamC decoupled weight decay.

In our scaling experiments, OKLS achieves 1.59× Muon's parameter efficiency at the 1.2B Muon-equivalent scale, with the advantage increasing toward larger models, while maintaining roughly 98% of its training throughput.

See the blog for the derivation, system design, and experiments: https://blog.tilderesearch.com/blog/online-kl-shampoo

How it works

For a matrix parameter with gradient $G_t \in \mathbb{R}^{m \times n}$, OKLS maintains left and right covariance factors $S_{a,t}$ and $S_{b,t}$. Using the previous preconditioners $P_{a,t-1}$ and $P_{b,t-1}$, it updates

$$S_{a,t} = \beta_2 S_{a,t-1} + \frac{1-\beta_2}{n} (G_t P_{b,t-1})(G_t P_{b,t-1})^\top,$$ $$S_{b,t} = \beta_2 S_{b,t-1} + \frac{1-\beta_2}{m} (P_{a,t-1}G_t)^\top(P_{a,t-1}G_t).$$

Fresh inverse roots are then computed immediately:

$$P_{a,t} = S_{a,t}^{-1/2}, \qquad P_{b,t} = S_{b,t}^{-1/2}.$$

The Nesterov momentum $N_t$ is whitened from both sides,

$$U_t = P_{a,t} N_t P_{b,t},$$

and scaled using the matrix-shape-dependent muP multiplier

$$c_{\mathrm{shape}} = \frac{\sqrt{m/n}}{\sqrt{m}+\sqrt{n}}.$$

The inverse roots are evaluated with Scaled CANS Coupled Newton–Schulz. The iteration uses Chebyshev-optimized coefficients and deterministic per-step scales to keep FP16 GEMM inputs within range while retaining FP32 accumulation and persistent state.

Requirements

• A CUDA-capable GPU.
• A recent CUDA-enabled PyTorch build with FP16 torch.baddbmm and FP32 out_dtype support.

This repository currently provides the optimizer as a compact Python package without a build configuration. Run examples from the repository root, or add the repository root to PYTHONPATH.

Usage

OnlineKLShampoo is a standard torch.optim.Optimizer. Pass only matrix parameters that should receive OKLS updates, and use a separate optimizer for embeddings, output heads, norms, biases, and other parameters.

import torch

from okls import OnlineKLShampoo

# Choose the matrix weights that should use OKLS explicitly.
okls_params = []
for block in model.layers:
    okls_params.extend([
        block.attn.q_proj.weight,
        block.attn.k_proj.weight,
        block.attn.v_proj.weight,
        block.attn.o_proj.weight,
        block.mlp.gate_proj.weight,
        block.mlp.up_proj.weight,
        block.mlp.down_proj.weight,
    ])

# Manage all remaining parameters separately.
okls_param_ids = {id(p) for p in okls_params}
other_params = [p for p in model.parameters() if id(p) not in okls_param_ids]

okls = OnlineKLShampoo(
    okls_params,
    lr=0.09434,
    beta1=0.9684,
    beta2=0.9482,
    eps=1e-9,
    weight_decay=0.0303,
)
adamw = torch.optim.AdamW(other_params, lr=3e-4)

for input_ids, labels in dataloader:
    loss = model(input_ids, labels=labels).loss
    loss.backward()

    okls.step()
    adamw.step()

    okls.zero_grad()
    adamw.zero_grad()

Learning-rate schedule

We use linear warmup followed by a constant phase, then linear decay in the last decay_ratio fraction of training:

total_steps = ...   # total training steps
warmup_steps = 250
decay_ratio = 0.7319   # decay begins at step (1 - 0.7319) * total_steps
min_lr_factor = 0.0    # final lr = min_lr_factor * peak lr

def lr_lambda(step):
    # Warmup: linear ramp from ~0 to 1
    if warmup_steps > 0 and step < warmup_steps:
        return (1 + step) / warmup_steps
    # Decay: linear ramp from 1 to min_lr_factor in last decay_ratio fraction
    progress = step / total_steps if total_steps > 0 else 1.0
    if progress > 1.0 - decay_ratio:
        decay_progress = (progress - (1.0 - decay_ratio)) / decay_ratio
        return 1.0 - decay_progress * (1.0 - min_lr_factor)
    return 1.0

scheduler = torch.optim.lr_scheduler.LambdaLR(okls, lr_lambda)

# After okls.step():
scheduler.step()

The initial learning rate is retained internally as lr_peak for the AdamC weight-decay correction, while the scheduler updates the current lr.

Batched matrices

The optimizer accepts both:

• 2D parameters with shape (m, n).
• 3D parameters with shape (N, m, n).

A 3D parameter is treated as a batch of N independent matrices. Each matrix receives its own (m, m) and (n, n) covariance factors and preconditioners, while all matrices share the optimizer hyperparameters.

Hyperparameters

Argument	Default	Description
lr	0.09434	Current learning rate and fixed initial lr_peak used by AdamC.
beta1	0.9684	Nesterov momentum EMA coefficient.
beta2	0.9482	Kronecker-factor EMA coefficient.
eps	1e-9	Stability term added during factor initialization and updates.
weight_decay	0.0303	Decoupled AdamC weight-decay coefficient.

The implementation stores momentum, covariance factors, and preconditioners in FP32. For a matrix of shape (m, n), the persistent state contains one (m, n) momentum matrix and two copies each of the (m, m) and (n, n) factor shapes.

Weight initialization

We use Spectral muP Gaussian Initialization, which samples i.i.d. Gaussian weights calibrated so the operator norm (spectral norm) of each matrix is a controlled constant at init. This is consistent with the muP shape scaling used in the optimizer update.

Given a weight matrix $W \in \mathbb{R}^{n_\text{out} \times n_\text{in}}$ and a target operator norm $\sigma^*$ (init_operator_norm):

1. Compute the RMS-to-RMS scale:

$$s = \sigma^* \sqrt{n_\text{out} / n_\text{in}}$$

2. Sample from a Gaussian at the Marchenko–Pastur edge:

$$W \sim \mathcal{N}\!\left(0,\;\left(\frac{s}{\sqrt{n_\text{in}} + \sqrt{n_\text{out}}}\right)^{\!2}\right)$$

The denominator $\sqrt{n_\text{in}} + \sqrt{n_\text{out}}$ is the Marchenko–Pastur edge for an i.i.d. Gaussian matrix scaled to unit variance per entry, ensuring the spectral norm concentrates around $s$.

Our sweep found the optimal value to be init_operator_norm = 0.07539.

import torch
import math

def spectral_mup_gaussian_init(weight: torch.Tensor, init_operator_norm: float = 0.07539):
    """Spectral muP Gaussian initialization for a 2D weight matrix."""
    n_out, n_in = weight.shape
    s = init_operator_norm * math.sqrt(n_out / n_in)
    std = s / (math.sqrt(n_in) + math.sqrt(n_out))
    weight.data.normal_(0, std)

Code structure

okls/
├── __init__.py                    # public OnlineKLShampoo export
├── okls_optim.py                  # torch.optim.Optimizer wrapper and state
├── okls_step.py                   # factor, momentum, muP, and AdamC updates
└── scaled_cans_coupled_ns.py      # scaled 10-step inverse-root iteration

This is the compact, pure-PyTorch optimizer implementation. It does not include the distributed optimizer-state offload and fused production training system described in the blog post.

References

• Online KL Shampoo
• Understanding and Improving Shampoo and SOAP via Kullback–Leibler Minimization
• Accelerating Newton–Schulz Iteration for Orthogonalization via Chebyshev-type Polynomials

Citation

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