Four optimizer refreshes are not one schedule

Expensive second-order optimizers often reuse a matrix factorization for many training steps. The usual shorthand says “refresh every k steps.” Wang and colleagues ask the missing question: which steps get the fresh preconditioner? Holding the number of refreshes fixed, they move the middle two updates earlier or later and find that the winner can reverse as learning rate changes.

A Kronecker factor approximates layerwise curvature as a product of two smaller matrices. Its eigendecomposition supplies a preconditioning transform: a rescaling of the gradient intended to make optimization better conditioned. Because decompositions are costly, Mousse caches them. At step t, cache age is the distance to the most recent refresh: aₛ(t) = t − max{s ∈ S : s ≤ t}.

def cache_ages(horizon, schedule):
    if horizon < 1 or not schedule or schedule[0] != 1:
        raise ValueError("schedule must start at step 1")
    ages = []
    refresh_index = 0
    latest_refresh = schedule[0]
    for step in range(1, horizon + 1):
        while (refresh_index + 1 < len(schedule)
               and schedule[refresh_index + 1] <= step):
            refresh_index += 1
            latest_refresh = schedule[refresh_index]
        ages.append(step - latest_refresh)
    return ages

This page runs the TypeScript version. The curves below are exact consequences of the paper's 591-step schedules, not measured language-model losses. The Python and C++ variants expose the same state update line for line.

On the 350M-parameter FineWeb experiment, the paper reports paired-seed validation-loss differences as late minus early. Negative means late wins; positive means early wins. The median moves from −0.0049 at learning rate 0.012 to +0.0157 at 0.024. A single-rate comparison could make either placement look universally superior.

The main-cohort estimates put the preference boundary at progressively smaller learning rates as the model grows: from 0.029393 at 64.1M parameters to 0.013196 at 928.8M. This is not a clean scaling law because model size, tokens per parameter, and training horizon move together.

Report the complete refresh-step list, learning-rate schedule, warmup, seed pairing, and total horizon. For a causal test, hold tokens per parameter fixed across scale, measure eigenvector rotation as well as spectral width, and repeat the grid with other preconditioners.

Start with the from-scratch neural network to make gradient scaling visible, then connect the schedule audit to the Transformer training loop. Frequency is a lossy label for a time-varying intervention.