K-Maxwell momentum

A mixture of EMAs is a weighting of past gradients by age, and its mean age is linear in the weights. Let there be $K$the number of EMA buffers: 8 on Muon, 6 on MuonH and SOAP-Muon buffers $m_{k,t}=\beta_km_{k,t-1}+(1-\beta_k)g_t$ with decay rates $\beta_k$the fixed decay rate of buffer $k$, set from its mean age by $\beta_k=\tau_k/(\tau_k+1)$, mixed with weights $w_k(t)$the weight of buffer $k$ at step $t$; the weights are nonnegative and sum to 1 into $\bar m_t=\sum_kw_k(t)m_{k,t}$. Unrolling each EMA (ignoring how the buffers were started) gives the weight $h_t(a)$the memory kernel: the weight the mixed momentum $\bar m_t$ puts on the gradient from $a$ steps before step $t$ on each past gradient: $$\begin{aligned}\bar m_t&=\sum_{a\ge0}h_t(a)\,g_{t-a},\\h_t(a)&=\sum_{k=1}^{K}w_k(t)\,(1-\beta_k)\,\beta_k^{\,a}.\end{aligned}$$ Write $\tau_k$${}=\beta_k/(1-\beta_k)$ for buffer $k$'s mean age; the mixture's mean age is then $\sum_kw_k(t)\tau_k$, linear in the weights, so interpolating the weights moves it linearly, from 58 steps to 26. The mean age does not fix the shape. In the post's example a single EMA with $\beta=0.95$ and an equal mix of EMAs with mean ages 3 and 35 both have mean age 19, but the newest gradient gets 5% of the weight in the first and 13.9% in the second. The tuned mixes differ the same way (my computation from the published weights): the final mix has the single EMA's median age, 13 steps, yet puts 0.11 of its weight on gradients older than 64 steps against 0.04; the starting mix puts 0.33 there.

On tuned Muon, K-Maxwell passes 90 steps before the single EMA and 50 before Bi-Maxwell. Eight seeds were evaluated every 10 steps from 2900 to 3250. At step 3150 the statistic is $(3.28-3.27873)\sqrt8=0.0036$, short of 0.004; at 3160 the mean is 3.27794 and the statistic 0.0058. The seeds' own first crossings of 3.28 average 3141 steps; the reported step is later because the mean must sit $0.004/\sqrt8\approx0.0014$ nats below the target. Against the single-EMA baseline's ten runs the gap is significant at equal steps by the track's two-sample form of the rule, $(3.27866-3.27328)/\sqrt{1/8+1/10}=0.0113\ge0.004$ at step 3250. K-Maxwell's mean loss is also below Bi-Maxwell's at every checked step from 2900 to 3250, by 0.0072 nats at 2900 and 0.0034 at 3250.

The annealing carries the gain, transfers to MuonH, and fails on the SOAP-Muon stack. Per the experiment report, fixed-weight eight-buffer mixes did not beat the baseline in any configuration tested, even after extensive tuning; the first annealed schedules won clearly. On MuonH the recipe, lightly retuned (six buffers, switch at step 750, mean age 50 to 22), passes at 3065 against 3125. The SOAP-Muon record stack also evaluates a Tail-EMAan exponential moving average of the weights over the last few hundred steps, evaluated in place of the final weights of the weights. On it a fixed six-buffer mix with mean age 35 passes at 2680 against 2690, not significant at equal steps (two-sample statistic 0.0010). Bi-Maxwell on that stack passes at 2635 and, per the authors' companion report, beats their mix in a paired comparison; annealing made the loss worse. The post attributes the failure to overlap between momentum and SOAP's preconditioning, the companion report to Tail-EMA, which already averages the end of training; neither is tested.

At larger batches K-Maxwell keeps its gain and Bi-Maxwell loses it. From three seeds' step-2000 states, Muon with momentum off ($\mu=0$), single-EMA Muon, Bi-Maxwell and K-Maxwell were each run to step 2750 at 1, 2, 4, 8 and 16 times the standard batch. Against momentum off at the same seed and batch, K-Maxwell lowered the validation loss by 0.0047 to 0.0067 nats at every batch size, while Bi-Maxwell's gain fell from 0.0105 nats at the standard batch to 0.0020 at 8 times and $-0.0006$ at 16 times, with seed-to-seed standard deviations of about 0.0002 or less. A larger batch already averages noise within the step, so a gain that comes only from averaging noise across steps should fade as the batch grows; the post reads K-Maxwell's persistent gain as doing more than that. The post's text leaves out two features of the table, both stated in the authors' linked ablation log. At the standard batch Bi-Maxwell's gain is twice K-Maxwell's, the reverse of the full runs in Step 3 (my comparison), so these short continuations do not reproduce the full-run ordering. And single-EMA Muon is within 0.0004 nats of momentum off at every batch size: over these steps ordinary momentum gives no measurable gain, and what the mixtures add is something a 19-step EMA does not provide.

The post's explanations are untested. They are that long memory keeps persistent gradient directions early and short memory reacts faster to oscillations late, near the edge of stability; that mixing very different timescales cancels recurring errors, loosely like Richardson extrapolationcombining estimates computed at different step sizes so that their leading error terms cancel; and that at large batches K-Maxwell lets training reach sharper regions. None of the experiments measures curvature or oscillation. The direction of the schedule is also open: the Bi-Maxwell paper's model predicts that the best memory grows over training, because gradient directions change more slowly later, and AdEMAMix lengthens its memory, as the post notes; K-Maxwell's memory jumps from the single EMA's 19 steps to 58 at step 1000 and then shortens.