K-Maxwell Momentum

Collaborator: Jeffrey Cheng · GitHub report

We believe momentum remains an under-optimized area of pre-training. While much recent research focuses on optimizer design, momentum often still uses a single exponential moving average (EMA) buffer with a fixed decay rate. K-Maxwell combines several such buffers and gradually shifts their weights during training. [1], [2]

Results

We simultaneously set NanoGPT records for 2 different optimizers (Muon and MuonH but not SOAP) by tuning a single more-expressive momentum implementation. K-Maxwell reduces the required training steps by 90 on Muon (2.77%) and 60 on MuonH (1.92%). The same momentum method is used for both optimizers, with six buffers on MuonH and eight on Muon. On SOAP-Muon, the observed reduction is only 10 steps (0.37%) on the current world record, with no statistically significant improvement over its predecessor, Bi-Maxwell. We suspect this is due to an overlap in how momentum and SOAP-style preconditioning [6] both address oscillations at the edge of stability.

Experiment design

With K-Maxwell, we test the hypothesis that more expressive momentum can accelerate convergence on modded-nanoGPT Track 3 optimization. This track measures how many optimizer steps are needed to reach below 3.28 validation loss.

The setup is simple, we use approximately 2 billion FineWeb tokens, with 524,288 tokens per step and we train GPT-2 small, which is 124M parameters with 12 transformer blocks, using the Muon and MuonH optimizers.

In our experiments, we compare single-EMA momentum, Bi-Maxwell’s fixed two-buffer mixture, and K-Maxwell’s six-buffer mixture with the MuonH optimizer and an eight-buffer mixture on the Muon optimizer. Within each optimizer comparison, we only change the momentum configuration, keeping the architecture, data, batch size, and other training settings fixed.

To ensure our results are not just random variance, we test our final K-Maxwell configuration across eight seeds. We report the earliest shared checkpoint satisfying the benchmark’s rule, (3.28 − mean validation loss) × √8 ≥ 0.004, and separately test improvement over the baseline at matching step counts.

After this, we also tested our configuration at larger batch sizes, which implies an increase in the number of tokens trained on - which we continue to use FineWeb for. Surprisingly, K-Maxwell retains an advantage at larger batch sizes while Bi-Maxwell approaches the matched Muon control with momentum turned off (µ = 0). This suggests that the annealed momentum mixture does more than dampen gradient noise, which increasing batch size adequately does.

What is momentum?

Momentum is an exponential moving average (EMA) of our past gradients, which we use to update our model parameters, after pre-conditioning our momentum values, with Muon, Adam or other optimizers. Muon (Jordan, 2024), [1]

Single EMA momentum [1].

What does momentum do?

Momentum is commonly conceived of as a heavy ball carrying us through small divots, saddle points and across flat regions of the loss surface.

While this is intuitively useful, momentum does much more - helping us converge faster when curvature differs greatly across our gradient directions, or in other words, our loss objective is poorly conditioned. For a positive-definite Hessian, which describes our loss surface’s local curvature, the condition number is the ratio of its largest to smallest eigenvalues:

To understand why momentum helps us converge faster, we can observe a simple example. For full-batch optimization of a convex quadratic loss, tuned Polyak momentum helps accelerate progress along directions of low curvature while damping oscillations along directions of high curvature, which occur at the edge of stability. In this simplified setting, it improves the convergence dependence on the condition number from κ to √κ. [7], [8]

For a positive-definite quadratic with optimally tuned constant parameters, the asymptotic contraction factors for parameter error are:

Smaller ρ means faster convergence. These are parameter-error rates; they are not the per-step contraction factors for the loss itself. [7]

Polyak momentum. [8]

The intuition is that gradients along a low curvature direction tend to keep the same sign: +g, +g, …, so momentum accumulates along a persistent gradient direction. However, down high curvature directions, updates can overshoot due to the discretized learning rate of gradient descent, causing gradients to alternate: −g, +g, −g, +g. Averaging these opposing gradients dampens this oscillating motion, which is exactly what momentum does.

Cohen et al. (2021) [9] observed that, during full-batch neural network training, the largest Hessian eigenvalue often rises and then hovers near $\frac{2}{\eta }$, where $\eta$ is the learning rate. This is called training at the edge of stability, and we find that momentum permits us to take on larger learning rates which can speed up convergence.

For the standard Polyak update on a positive-definite quadratic loss, the new stability condition is $0<\eta <\frac{2+2\beta }{{\lambda }_{max}}$, where $\beta$ is the momentum coefficient. [7] In the experiments of Andreyev et al., large-batch momentum dynamics approach the optimizer-specific deterministic stability boundary, while smaller batches impose tighter stability constraints. [17]

K-Maxwell momentum

K-Maxwell momentum extends Bi-Maxwell [10] and AdEMAMix [2] using eight EMA buffers with logarithmically spaced timescales and fixed decay rates (β). AdEMAMix and Bi-Maxwell both use two EMA buffers, but AdEMAMix anneals from a shorter to a longer memory — the reverse of our annealing schedule and AdEMAMix also uses the Adam optimizer. AdEMAMix further anneals the β₃ decay rate itself over training, along with the relative weight of the longer memory EMA buffer to curb instability early on in training. Bi-Maxwell however, has a fixed memory kernel, and instead of just adding a weighted, longer memory EMA buffer, it does a weighted average of both the short and longer term memory.

AdEMAMix

Here m₁ and m₂ are the fast and slow gradient averages. ν averages squared gradients. Hats denote correction for zero initialization. α(t) controls the slow buffer’s contribution, β₃(t) controls its memory, and ε is a small constant. [2]

Bi-Maxwell, in the tuned Muon comparison

The fast and slow buffers use fixed decay rates and mixture weights after activation. µ is the unchanged Nesterov mixing coefficient. Bi-Maxwell implementation

In our variant, K-Maxwell, training begins with a single EMA buffer and switches to K-Maxwell at step 1000, skipping the instability AdEMAMix sees early on in training from a longer memory EMA buffer. From the switch step, we linearly interpolate the mixture weights from a mean age of 58 steps toward a mean age of 26 steps, at step 3250. Mean age describes on average, the gradient age in the EMA buffer, weighted by each past gradient’s contribution. [3]

The following pseudocode illustrates the annealed momentum after the switch step. Initialization and the surrounding optimizer are defined in the source trainer [11].

# per Muon 2-D param, step >= 1000:
for k, beta in enumerate(kmaxwell_decay_rates):  # 8 log-spaced τ in [3, 64]
    m[k].lerp_(g, 1 - beta)
frac = (step - 1000) / 2250
w = (1 - frac) * w_age58 + frac * w_age26        # mixture mean age 58 → 26
m_eff = sum(w[k] * m[k] for k in range(8))
update = g.lerp_(m_eff, mu)                       # Nesterov mix unchanged

We use logarithmically spaced timescales to capture meaningfully different ranges of gradient history. Slower EMAs can act as low-pass filters [12], smoothing out high frequency changes while retaining persistent gradient directions. Since we know characteristic frequency $\omega$ is approximately inversely proportional to memory $\tau$ (i.e. wavelength), we adopt a logarithmic scale.

However, mean age alone does not capture the expressivity allowed by multiple EMA buffers. As discussed in AdEMAMix [2], an EMA mixture can give substantial weight to the newest gradients while retaining a long tail of older gradients. The example below shows how two memory profiles can have the same mean age but have vastly different gradient ages present:

We suspect that combining very distinct timescales at each training step helps reduce recurring errors, including oscillations at the edge of stability. This is loosely analogous to Richardson extrapolation [13], which combines estimates at different step sizes to remove predictable error. Our mechanism is more closely associated with the idea of multi-timescale temporal filtering, however we presume it goes beyond this as it retains a stable advantage over normal single EMA Muon even with larger batch size. For related mathematical background on stochastic momentum dynamics, see Li, Tai, and E (2019). [18]

Annealing momentum

We also observe substantial gains from annealing momentum, and see a step-function improvement over configurations that kept the mixture weights fixed. [14] We suspect that longer memory helps maintain persistent gradient directions during the extended transient phase, when training is still making broad progress. Shorter memory may then improve responsiveness during the late-stage convergence phase, when the optimizer needs to adapt quickly to oscillatory changes in gradient direction. [15]

Another possible explanation is that long momentum memory becomes less useful late in training as retaining older gradient directions can delay the optimizer’s response to changes in the loss landscape. Shortening that memory may improve responsiveness, although it also reduces noise smoothing. Optimizer families such as RMSProp and ADAM, attempt to solve related problems by adjusting updates responsively, using running squared-gradient estimates to scale updates. This changes the effective learning rate of each component, potentially slowing convergence when gradient norm grows large - which is likely to occur when training at the edge of stability. On the other hand, K-Maxwell instead adjusts how much gradient history contributes to momentum, which appears to strictly improve convergence in our experiments.

Extending to larger batch sizes

We also varied batch size in matched continuation experiments and found that K-Maxwell’s advantage remains robust at larger batch sizes, while Bi-Maxwell’s performance approaches the control: Muon with momentum turned off (µ = 0). Ordinary single-EMA Muon (µ = 0.95) is included as a separate comparison. We suspect this reflects a shift in momentum’s role as larger batches make gradients less noisy. Recent work on batch sharpness, which describes the curvature encountered along each mini-batch’s gradient direction, finds that momentum imposes a tighter stability constraint at small batch sizes. At larger batch sizes, training approaches the optimizer’s deterministic stability boundary, that is the limit associated with full-batch gradients, allowing it to reach sharper regions than in the small-batch regime, which allows us to continue to converge despite high curvature. [17]We hypothesize that K-Maxwell retains its advantage at larger batch sizes by extending this effect to even higher curvature regions.

A related explanation is that K-Maxwell’s momentum annealing becomes more useful as the balance shifts from mostly gradient noise to predominately oscillations caused by curvature, causing us to overshoot along steep directions. Annealing the mixture may help K-Maxwell adapt to this shift more effectively than a fixed momentum mix such as in Bi-Maxwell.

Limitations

K-Maxwell has a substantial memory cost at larger model sizes as it stores seven additional EMA buffers per parameter that it is applied to. This can become prohibitively expensive at the trillion-parameter scale. We are actively working on momentum approaches that use less memory.

Conclusion

By extending single EMA and Bi-Maxwell momentum, K-Maxwell shows that mixing multiple momentum timescales and annealing their mixture during training can speed up pre-training convergence. Its robustness at larger batch sizes also suggests that these improvements extend beyond simply dampening stochastic gradient noise. We presume the advantage comes from maintaining a longer tail of gradient history and having this moving average of

Many open questions remain about why temporal filtering improves optimization, and through what mechanism mixed EMA momentum does this. We believe understanding these mechanisms could help us design more effective, memory-efficient momentum methods.

Acknowledgements

Thank you to Jerry Hong for assisting with this project. Baseline implementations and their contributors are credited in the linked experiment reports.

References

[1] Kingma, D. P., & Ba, J. (2015). Adam: A method for stochastic optimization. International Conference on Learning Representations.

[2] Pagliardini, M., Ablin, P., & Grangier, D. (2024). The AdEMAMix optimizer: Better, faster, older. arXiv.

[3] jacknzheng. (2026c, August 27). Track 3: Annealed K-Maxwell momentum on tuned Muon—3160 steps (n=8) [Pull request #357]. GitHub.

[4] jacknzheng. (2026d, August 28). Track 3: K-Maxwell on MuonH fast-slow decay, 3065 steps (n=8) [Pull request #359]. GitHub.

[5] jacknzheng. (2026a). Record: Track 3 optimization—Frozen K-Maxwell momentum on SOAP-CWD—2680 steps (n=8) [Experiment report]. GitHub.

[6] Vyas, N., Morwani, D., Zhao, R., Kwun, M., Shapira, I., Brandfonbrener, D., Janson, L., & Kakade, S. (2024). SOAP: Improving and stabilizing Shampoo using Adam. arXiv.

[7] Goh, G. (2017). Why momentum really works. Distill, 2(4), Article e6. doi:10.23915/distill.00006

[8] Polyak, B. T. (1964). Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5), 1–17.

[9] Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z., & Talwalkar, A. (2021). Gradient descent on neural networks typically occurs at the edge of stability. International Conference on Learning Representations.

[10] Hu, Y., Xiang, H., Gong, X., & Yu, H. (2026). A physical response-and-memory model for Muon optimization. arXiv.

[11] jacknzheng. (2026e). train_gpt_kmaxwell_anneal.py [Python source code]. GitHub.

[12] Smith, J. O., III. (2007). Introduction to digital filters: With audio applications. W3K Publishing.

[13] Israel, R. (2002, January 29). Richardson extrapolation. University of British Columbia.

[14] jacknzheng. (2026b). Record: Track 3 optimization—K-Maxwell annealed momentum on the tuned Muon baseline—3160 steps (n=8) [Experiment report]. GitHub.

[15] Sutskever, I., Martens, J., Dahl, G., & Hinton, G. (2013). On the importance of initialization and momentum in deep learning. In S. Dasgupta & D. McAllester (Eds.), Proceedings of the 30th International Conference on Machine Learning (Vol. 28, pp. 1139–1147). PMLR.

[16] Hinton, G., Srivastava, N., & Swersky, K. (n.d.). Neural networks for machine learning: Lecture 6e—RMSProp: Divide the gradient by a running average of its recent magnitude [Lecture slides]. University of Toronto.

[17] Andreyev, A., Ananthkumar, A., Walden, M., Poggio, T., & Beneventano, P. (2026). Momentum further constrains sharpness at the edge of stochastic stability. arXiv.

[18] Li, Q., Tai, C., & E, W. (2019). Stochastic modified equations and dynamics of stochastic gradient algorithms I: Mathematical foundations. Journal of Machine Learning Research, 20(40), 1–47.

[19] Cohen, J., Damian, A., Talwalkar, A., Kolter, J. Z., & Lee, J. D. (n.d.). Part I: How does gradient descent work? Understanding Optimization in Deep Learning with Central Flows [Companion website].