A Physical Response-and-Memory Model for Muon Optimization

Yinze Hu; Hongjun Xiang; Xingao Gong; Hongyu Yu

cs.LGAug 24, 2026

A Physical Response-and-Memory Model for Muon Optimization

Authors: Yinze Hu, Hongjun Xiang, Xingao Gong, Hongyu Yu

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

Abstract

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^{\star}

e_{\alpha}=y_{\alpha}-y_{\alpha}^{\star}

, is measured by the mismatch potential

\Phi=\frac{1}{2}\sum_{\alpha}e_{\alpha}^{2}

, so that each output channel is equivalent to a linear spring whose elongation is the mismatch. For a selected coupling element

W_{\alpha i}

, only

e_{\alpha}

contributes to the first-order variation, giving the conjugate driving force

X_{\alpha i}=-e_{\alpha}x_{i}

; assembling all components gives

X=-\partial\Phi/\partial 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 ).

Panel of Figure 2 — 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 $\beta$ 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 $\alpha$ 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: $\tau_{s}=5\tau_{f}$ , $w=0.44$ ; the recipe used in the experiments is $\beta_{f}=0.85$ , $\beta_{s}=0.98$ , $w=0.4385$ , see Methods). A constant external force is applied at $t=0$ and removed at $t=4\tau_{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 $M_{\mathrm{BM}}=wM_{f}+(1-w)M_{s}$ (black dashed) is fast first and slow later, combining response speed with memory depth.

Panel of Figure 3 — 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 $M_{\mathrm{BM}}=wM_{f}+(1-w)M_{s}$ , which then enters the spatial law to complete the weight-matrix rearrangement.

Panel of Figure 4 — 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-\bar{L})\sqrt{n}\geq 0.004$ , that is $\bar{L}\leq 3.28-0.004/\sqrt{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\pm 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\pm 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

\pm

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)	$b_{1}$	$b_{2}$	$b_{3}$	$b_{4}$	$b_{5}$
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^{\dagger}$ $(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

R_{b}

per mode band under the three learning-rate protocols (the ratio of the median

n^{\star}

in the late window

[2200,2700]

to that in the early window

[700,1500]

, taken as the median across trajectories).

R_{b}>1

means a longer memory is preferred late; the count in parentheses is the number of trajectories that rise in that band.

b_{1}

is the strongest

5\%

and

b_{5}

the weakest stretch. The

b_{2}

of the constant high learning rate is marked

\dagger

: computed directly from

n^{\star}

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

b_{1}

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

R_{b}

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

R_{b}=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

b_{1}

lengthens while

b_{3}

–

b_{5}

shorten; under the constant low learning rate all five lengthen again. (b) Per-tensor paired

\Delta\log\hat{D}_{b}

(down triangles) and

\Delta\log\hat{T}_{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

b_{1}

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

R_{b_{1}}=1.0558

after cancellation. The preregistered primary criterion of the three-arm experiment is the

b_{1}

of the two frozen arms; the rest are exploratory results.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

seed	crossing step
0	3210
1	3180
2	3175
3	3200
4	3190
5	3210

Appendix

Table 3: Per-seed crossing steps: bare tuned-Muon stack, two-timescale kernel (

n=8

; data: bare_twounit_perseed.csv ).

seed	crossing step
0	2620
1	2620
2	2640
3	2600
4	2630
5	2585

Appendix

Table 4: Per-seed crossing steps: record stack, two-timescale kernel (A800,

n=8

; data: record_twounit_perseed.csv ).

seed	crossing step
0	2615
1	2615
2	2635
3	2640
4	2645
5	2595

Appendix

Table 5: Per-seed crossing steps: record stack, two-timescale kernel, independent H100 replication (

n=8

; data: H100-arm logs of record_st1000_submission ).

seed	crossing step
0	2635
1	2640
2	2645
3	2615
4	2630
5	2610

Appendix

Table 6: Per-seed crossing steps: record stack, two-timescale kernel,

T_{\mathrm{on}}=700

arm (A800,

n=8

; data: logs of record_2655_submission ).

seed	crossing step
0	2760
1	2725
2	2745
3	2750
4	2745
5	2740

Appendix

Table 7: Per-seed crossing steps: matched single-pole control (equal

\beta=30/31

\bar{n}=30

, record stack, A800,

n=8

; data: sp30_matched_perseed.csv ).

Panel of Figure 6 — Figure 6: Age dose curve (fork@700, seed 0, window $[2650,2700]$ , zeroed at the $\bar{n}=19$ matched point). The vertical axis $\Delta 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 $\bar{n}\approx 30$ , degrading at both ends.

Panel of Figure 7 — 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^{\star}$ for each mode band (8 points per band, one per trajectory, slightly offset horizontally for readability). All $8\times 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^{\star}

(defined in Methods) versus training step (seed-0 trajectory). (a) Each curve is the median over the

72

hidden-layer matrices of

n^{\star}

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^{\star}

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 (

Q_{1}

–

Q_{3}

) 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 .

Neural Network Optimization

Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.

See all 321 papers

Monarch-Muon: Structure-Aware Muon for Fast and Memory-Efficient Pretraining
Oct 12, 2026Alexander Miasnikov, Ruslan Kabirov, Gleb Molodtsov +7Muon Optimizer LLM Training
BRAIn Lab · HSE University · MBZUAI
The Newton-Muon Optimizer
Apr 1, 2026Zhehang Du, Weijie SuNeural Network Optimization Muon Optimizer
University of Pennsylvania
Spectrally Targeted Muon
Oct 7, 2026Vishrut Goyal, Rohan RamkumarSpectral Optimization Deep Learning Optimization
Department of Computer Science, Stanford University · Department of Electrical Engineering, Stanford University.

Oct 12, 2026J/K move · Enter open · S save

Monarch-Muon: Structure-Aware Muon for Fast and Memory-Efficient Pretraining

Alexander Miasnikov, Ruslan Kabirov, Gleb Molodtsov, Dmitrii Vlashchuk+6

BRAIn Lab · HSE University · MBZUAI

Muon Optimizer LLM Training Memory-Efficient Training Language Model Pretraining

Open paper AskPDF arXiv

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.

Abstract

Neural Network Optimization

Explore similar work

Monarch-Muon: Structure-Aware Muon for Fast and Memory-Efficient Pretraining

The Newton-Muon Optimizer

Spectrally Targeted Muon