Organizations: Key Laboratory of Computational Physical Sciences (Ministry of Education), Institute of Computational Physical Sciences, State Key Laboratory of Surface Physics, and Department of Physics, Fudan University, Shanghai 200433, China · School of Physics Science and Engineering, Tongji University, Shanghai 200092, China
Training large language models is costly. How low a loss the same compute can ultimately reach depends on how each step's gradient is converted into a weight update; the rule that performs this conversion is the optimizer. From SGD and AdamW to the recent Muon, effective update rules have mostly been shaped by engineering intuition and then selected on benchmarks. Muon semi-orthogonalizes the momentum matrix before applying the update and has kept breaking records on public training benchmarks; yet why the semi-orthogonalized direction works, and over how long a history the momentum should average, are two questions at present answered mainly by experience. Here we treat the weight matrix during training as a responsive medium with memory and build a physical model for it, in which both questions find answers: the semi-orthogonalized direction is the maximally dissipative response under an output-side safety budget, which explains why it works; momentum is the internal stress accumulated by the medium; how long it should average is set by the relaxation of this stress, and a real medium relaxes on more than one timescale, the simplest form being one fast and one slow. On this basis we propose the Bi-Maxwell optimizer. The framework further yields a testable consequence: gradient directions change fast early in training and more slowly later, so the optimal memory length should grow with training stage; step-by-step measurements of a proxy for it by a read-only probe across 8 independent training trajectories are consistent with this consequence. Replacing the memory kernel alone, from a single timescale to two, brings training to the target loss in noticeably fewer steps on a public large-language-model optimizer benchmark.
Figures & tables
Figure 1: Construction of the responsive medium and the mismatch potential (schematic). An input x passes through the weight matrix W to produce the output y=Wx ; the mismatch between the output and the target y⋆ , eα=yα−yα⋆ , is measured by the mismatch potential Φ=21∑αeα2 , so that each output channel is equivalent to a linear spring whose elongation is the mismatch. For a selected coupling element Wαi , only eα contributes to the first-order variation, giving the conjugate driving force Xαi=−eαxi ; assembling all components gives X=−∂Φ/∂W . The linear readout and quadratic mismatch shown here are an illustrative special case chosen so that every step can be drawn; for a general nonlinear network the same definition is carried by the backpropagation gradient and does not depend on this special case (full derivation in Sec. 3 ).
Figure 2: Glassy picture and temporal picture of the Bi-Maxwell memory (schematic). (a) Microscopic analogy of the fast and slow internal relaxation modes in a glassy medium. The same matrix drive X(t) excites both types of response: the fast mode corresponds to local relaxation of a single particle inside the cage of its neighbours, in which no contact of the cage is broken, so that the mode follows the drive fast and forgets fast (analogous to β relaxation in glasses); the slow mode corresponds to a cooperative structural rearrangement of a group of particles, in which existing contacts must first break and then reconnect (grey dashed and blue solid lines in the panel), so that the mode builds and relaxes slowly and leaves long-lived structural memory once the contacts have changed (analogous to α relaxation). The contrast between a single particle and a group of particles is drawn only to separate the two types of motion; the fast and slow modes of this paper are internal modes of one and the same medium and differ only in relaxation time (Sec. 3.3 ). The panel provides intuition by analogy, not a microscopic derivation of the optimizer dynamics. (b) Temporal picture (illustrative parameters: τs=5τf , w=0.44 ; the recipe used in the experiments is βf=0.85 , βs=0.98 , w=0.4385 , see Methods). A constant external force is applied at t=0 and removed at t=4τf (grey). The fast mode follows quickly and also forgets quickly (orange); the slow mode builds up slowly and relaxes slowly (blue); the macroscopic internal stress MBM=wMf+(1−w)Ms (black dashed) is fast first and slow later, combining response speed with memory depth.
Figure 3: Two-timescale structure of Bi-Maxwell (schematic). The same matrix force X(t) drives the fast and the slow relaxation mode; the two combine with weights w and 1−w into the macroscopic internal stress MBM=wMf+(1−w)Ms , which then enters the spatial law to complete the weight-matrix rearrangement.
Figure 4: Thin lines are per-seed curves, thick lines arm means; the dashed line is the 3.28 target, and the dotted lines are the pass thresholds implied by the margin criterion (3.28−Lˉ)n≥0.004 , that is Lˉ≤3.28−0.004/n , which depends on the seed count: 3.27859 at n=8 (orange) and 3.27874 at n=10 (blue). (a) Bare tuned-Muon stack (raw validation loss): the two-timescale arm (orange, A800, n=8 ) first crosses at step 3210 ; the tuned reference (blue, H100, n=10 , by its n=10 criterion) first crosses at step 3250 . (b) Record stack (Tail-EMA readout): both arms have n=8 , so only the 3.27859 threshold applies; the two-timescale arm (orange, A800) at step 2635 , the existing record (blue, A40, PR #328 public logs) at step 2690 . The spread of the thin lines is the within-arm seed spread; no separate error range is drawn. The pass criterion is judged one-sided (see Methods).
Stack
Memory kernel
Hardware
First crossing
n
Per-seed and source
Bare tuned-Muon
single-pole
H100
3250
10
official #36 tuned log
Bare tuned-Muon
two-timescale
A800
3210
8
3190.6±15.2 ; this work (PR #340)
Record stack
single-pole
A40
2690
8
existing record (PR #328 logs)
Record stack
single-pole
A40
2735
8
same logs, recomputed with Tail-EMA off
Record stack
two-timescale
A800
2635
8
2616.9±17.9 ; main result (PR #339), margin 0.00419
Record stack
two-timescale
A800
2690
8
as above, Tail-EMA off; ordering preserved
Table 1: Summary of the evidence boundary. Smaller first-crossing steps are better; the existing records in the table are the state as of July 2026 (see also the Note added in the Discussion). The hardware column gives the GPU on which each arm actually ran; rows on different hardware are not compared directly. The per-seed column gives the mean ± standard deviation of the individual seeds’ own crossing steps (data in Appendix C); a blank entry means this paper has no per-seed data for that arm. Step counts can be verified against the public benchmark and the corresponding PR logs; results introduced in this work are documented in PR #339/#340.
Protocol (trajectories)
b1
b2
b3
b4
b5
Original schedule (8)
1.1828(8/8)
1.3519(8/8)
1.4646(8/8)
1.8132(8/8)
2.3130(8/8)
Constant high LR (16)
1.0558(14/16)
0.9798†(3/16)
0.9187(0/16)
0.9063(0/16)
0.9584(2/16)
Constant low LR (8)
1.0689(7/8)
1.1077(8/8)
1.0587(8/8)
1.1168(8/8)
1.1260(8/8)
Table 2: Kernel-age ratio Rb per mode band under the three learning-rate protocols (the ratio of the median n⋆ in the late window [2200,2700] to that in the early window [700,1500] , taken as the median across trajectories). Rb>1 means a longer memory is preferred late; the count in parentheses is the number of trajectories that rise in that band. b1 is the strongest 5% and b5 the weakest stretch. The b2 of the constant high learning rate is marked † : computed directly from n⋆ it shortens slightly, while the per-tensor decomposition under a different aggregation order gives a slight increase, and this paper judges it near-neutral. The preregistered primary criterion of the three-arm experiment is the b1 of the two frozen arms; the complete five-band profile was identified only after seeing the data.
Figure 5: Band-resolved optimal-memory profile under the three learning-rate protocols. (a) The kernel-age ratio Rb across mode bands; faint small markers are the per-trajectory readings ( 8 each for the original schedule and the constant low learning rate, 16 for the constant high learning rate), thick lines the median across trajectories, the grey dashed line Rb=1 , and the vertical axis is logarithmic. The original schedule lengthens in all five bands, by more towards the weaker bands; under the constant high learning rate only b1 lengthens while b3 – b5 shorten; under the constant low learning rate all five lengthen again. (b) Per-tensor paired ΔlogD^b (down triangles) and ΔlogT^b (dashes); the two points of one band are joined by a thin line, coloured as in (a). Both are negative, that is, drift and batch noise weaken together in every protocol and every band; what decides whether the optimal memory lengthens or shortens is which of the two falls more. (c) The near-cancellation in b1 of the constant high learning rate: late-window drift and batch noise are 0.283 and 0.298 times their early-window values (falls of about 72% and 70% ), leaving the net increase Rb1=1.0558 after cancellation. The preregistered primary criterion of the three-arm experiment is the b1 of the two frozen arms; the rest are exploratory results.
Figure 6: Age dose curve (fork@700, seed 0, window [2650,2700] , zeroed at the nˉ=19 matched point). The vertical axis ΔL is the fixed-step validation-loss difference relative to the matched point (lower is better); the curve is U-shaped with its minimum near nˉ≈30 , degrading at both ends.
Figure 7: Cross-seed replication of aging. On 8 independent from-scratch trajectories, the ratio of the late-window [2200,2700] to the early-window [700,1500] median n⋆ for each mode band (8 points per band, one per trajectory, slightly offset horizontally for readability). All 8×5=40 ratios exceed 1 (grey dashed line; range 1.13 – 2.43 ); the five bands within one trajectory are not mutually independent, and the independent unit of replication is the 8 trajectories: all five bands rise from the early to the late window in every one of the eight trajectories. Each point is one trajectory’s ratio and the eight points are the entire sample; no summary or error bar is drawn. This figure is descriptive; the significance test (one-sample t test, two-sided) is reported in the main text and Methods.
Figure 8: Local kernel-age proxy n⋆ (defined in Methods) versus training step (seed-0 trajectory). (a) Each curve is the median over the 72 hidden-layer matrices of n⋆ for one mode band (band 1 is the strongest 0−5% , band 5 the weakest 65−100% ); the vertical axis is logarithmic; the first 250 steps are estimator warm-up and are not shown; the two vertical strips are the pre-specified windows [700,1500] and [2200,2700] . (b) Median n⋆ per band within the two windows (horizontal axis logarithmic), open markers for the early window, filled for the late window. The shaded bands in (a) are the interquartile range ( Q1 – Q3 ) over the 72 hidden matrices within each band, smoothed in the same rolling window as the median lines; their mutual overlap shows that the dispersion of individual channels is much larger than the separation between bands, so what is resolved is the ordering of the medians rather than the value of any single channel. This figure shows a single trajectory; across-trajectory statistics are in Table 2 of the main text, and the per-point replication over 8 trajectories is in Fig. 7 .
Training large language models is constrained by both wall-clock time and memory footprint, and recent matrix-aware optimizers sharpen this tradeoff: Muon can improve optimization but its dense polar update remains expensive. We propose Monarch-Muon, a structured matrix-aware optimizer for weights parameterized by Monarch factors. The resulting structure reduces weights, gradients, and optimizer state, and exposes small independent update subproblems in place of a single dense orthogonalization; we formalize this through a blockwise trust-region interpretation. In quality runs, Monarch-Muon improves over Monarch-AdamW at fixed structure: 3.06 vs. 3.09 in the dense 257M model and 2.68 vs. 2.77 in the 1B MoE model; the two-block MoE variant also matches dense AdamW while using 15% fewer active and 25.4% fewer total parameters. In 6.89B-parameter systems measurements, Monarch-Muon reduces peak memory from 68.9 to 52.9 GB, Muon optimizer state from 14.60 to 6.77 GB, and the Muon optimizer step from 1288 to 304 ms.
Alexander Miasnikov, Ruslan Kabirov, Gleb Molodtsov +7
The Muon optimizer has received considerable attention for its strong performance in training large language models, yet the design principle behind its matrix-gradient orthogonalization remains largely elusive. In this paper, we introduce a surrogate model that not only sheds new light on the design of Muon, but more importantly leads to a new optimizer. In the same spirit as the derivation of Newton's method, the surrogate approximates the loss as a quadratic function of the perturbation to a weight matrix W using only three matrices: the gradient G, an output-space curvature matrix H, and the data matrix Z that stacks the layer inputs. By minimizing this surrogate in one step and adopting a certain isotropic assumption on the weights, we obtain the closed-form update rule (up to momentum and weight decay) W←W−η⋅msgn(G(ZZ⊤)−1), where η is the learning rate and msgn(X)=UV⊤ if X=USV⊤ is a compact singular value decomposition. This new optimization method, which we refer to as Newton-Muon, shows that standard Muon can be interpreted as an implicit Newton-type method that neglects the right preconditioning induced by the input second moment. Empirically, on a reproduction of the earliest publicly released Modded-NanoGPT speedrun configuration using Muon for GPT-2 pretraining, Newton-Muon reaches the target validation loss in 6% fewer iteration steps and reduces wall-clock training time by about 4%.
The Muon optimizer orthogonalizes each update matrix, setting all of its singular values to one, and has proven highly effective for training large language models. It remains unclear, however, whether this success comes from amplifying small singular directions that gradient descent neglects or from suppressing large, degenerate directions that disrupt training. We introduce Spectrally Targeted Muon, which orthogonalizes only the singular values above or below a threshold τ, so that varying τ interpolates between normalized SGD and Muon. It isolates the relevant singular subspaces with projections computed by Newton-Schulz iteration on a shifted Gram matrix, so no SVD is needed. We evaluate these variants on the CIFAR-10 and NanoGPT speedruns, tracking the effective rank of gradient, update, and weight matrices and a new metric, the alignment of updates with the tangent space of the weight matrix's isospectral manifold. We find three things. First, the small singular values of the momentum are not noise. Orthogonalizing everything except the few largest singular values of each matrix nearly matches Muon while touching only a small fraction of the momentum, whereas orthogonalizing only the top falls well short even though it holds almost all of it. On language models every momentum singular value is far below one, so targeted orthogonalization can only amplify, and Muon wins by making directions that are too small to train on at their raw scale trainable. Second, shrinking the largest singular values is what keeps the parameter spectrum flat; this is cheap, and it is not what drives the loss. Third, AdamW differs from Muon mainly in how slowly it builds structure, which explains its slower start and why warmup helps AdamW but only hurts Muon.
Vishrut Goyal, Rohan Ramkumar
Department of Computer Science, Stanford University · Department of Electrical Engineering, Stanford University.