Source-linked AI summary
A Physical Response-and-Memory Model for Muon Optimization
Yinze Hu, Hongjun Xiang, Xingao Gong, Hongyu Yu
TL;DR
Large-language-model optimizers remain costly to evaluate, and key questions about Muon’s update direction and momentum memory lack systematic theory. The paper models training as a responsive medium, derives Bi-Maxwell’s two-timescale memory, and reports earlier target attainment on a public benchmark alongside measurements consistent with longer memory later in training.
Problem
Optimizer rules have largely been selected through engineering intuition and benchmarks, leaving why Muon’s semi-orthogonalized direction works and how long momentum should average as open questions.
Method
The paper treats the training weight matrix as a responsive medium whose output safety budget determines the update direction and whose internal-stress relaxation determines momentum memory, then replaces Muon’s single timescale with fast and slow memory.
Results
Bi-Maxwell reached the modded-nanogpt target at step 2635 versus the 2690-step standing record, while measurements across 8 trajectories were consistent with longer optimal memory later in training.
Takeaways & Limitations
The framework supplies a physical explanation for Muon’s semi-orthogonalized direction and supports replacing single-timescale momentum memory with a two-timescale kernel.
Takeaways & Limitations
The confirmatory memory-length conclusion is limited to the Track-3 setting at 124 M parameters, and interactions between memory kernels and Tail-EMA require a same-hardware factorial experiment.
Abstract
from arXiv · showhide
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.
1 Introduction
Optimizer design determines training efficiency and attainable loss, yet key choices in Muon remain mainly empirical. The paper develops a physical memory model and proposes Bi-Maxwell, which reaches the benchmark target earlier while motivating stage-dependent memory.
- Motivation: Optimizer choice determines training speed, stability, and the loss attainable under a fixed compute budget.The optimizer converts each gradient into a parameter update.
- Existing optimizers: Muon semi-orthogonalizes the momentum matrix before applying updates and has achieved strong public language-model training results.Muon uses Newton–Schulz iterations to approach the momentum matrix’s orthogonal polar factor.
- Open questions: Existing optimizer rules were largely shaped by engineering intuition and benchmark selection, leaving Muon’s update direction and memory length without a systematic physical theory.The paper identifies these as two basic unresolved questions.
- Proposed framework: The paper models the weight matrix as a responsive medium with memory, interpreting Muon’s direction as maximally dissipative under an output-side safety budget and momentum as internal stress.The model derives optimizer laws from the medium’s response to applied forces.
- Proposed framework: Bi-Maxwell replaces Muon’s single-timescale momentum memory with two-timescale stress memory while keeping the other Muon components unchanged.The two timescales are one fast and one slow.
2 Related work
The paper places Bi-Maxwell within optimizer and physical-memory traditions, connecting single-timescale momentum to Maxwell-type relaxation and two-timescale memory to richer relaxation spectra and aging.
- Optimizer development: Optimizer development spans momentum, adaptive methods, matrix preconditioning, symbolic search, lightweight second-order estimation, and Muon’s use of matrix spectral geometry.The related-work discussion positions Muon within several established optimizer-design lines.
- Physical precedents: Momentum’s exponential average has the form of single-relaxation-time viscoelastic stress, with β related to relaxation time τ by β = e^−∆t/τ.Real materials generally exhibit many internal modes rather than one relaxation time.
- Physical precedents: A two-timescale optimizer kernel is the minimal discretization of a positive relaxation spectrum, while Mori–Zwanzig reduction supplies a precedent for convolutional memory.The physical analogy motivates memory over gradient history.
- Physical precedents: The observed lengthening of optimal memory with training stage resembles physical aging, in which older polymer glasses relax more slowly.The paper presents this as a similarity in form rather than an identity of mechanisms.
3 A physical model for the optimizer
The paper models a weight matrix as a responsive medium whose output-constrained motion maximizes mismatch-potential dissipation and whose internal stress supplies momentum memory. Because real media have multiple relaxation times, the model extends single-timescale momentum to Bi-Maxwell’s fast and slow memory modes.
- 3.1 Space: the responsive medium and the update direction: The model treats a trainable weight matrix as a linear responsive medium embedded in a nonlinear network, mapping input x to output y = Wx.Weight rearrangement is analyzed as collective motion of the matrix while W remains fixed during each fast input–output response.
- 3.1 Space: the responsive medium and the update direction: The structural velocity is constrained by output-side perturbation rather than by individual coupling-element speeds.This reflects the physically observable quantity: the change in the downstream output produced by the matrix motion.
- 3.1 Space: the responsive medium and the update direction: Among velocities satisfying the output-gain limit, the model selects the one that maximizes mismatch-potential release.The mismatch potential is built from response deviations, whose restoring force drives the matrix toward the target response.
- 3.1 Space: the responsive medium and the update direction: Under the output cap, all driven collective channels update at the same safe amplitude, allowing weak directions to contribute without increasing worst-case output perturbation.Undriven modes are set to zero; compared with viscous dynamics, the rule opens more parallel mismatch-release channels.
- 3.2 Time: internal stress memory: Momentum is interpreted as internal stress that accumulates recurring force directions and screens transient noise before the full-amplitude update.Only forces persisting across multiple steps remain in the memory; briefly appearing directions cancel rather than consuming the output perturbation budget.
- 3.3 From a single relaxation time to a two-timescale kernel: A single relaxation time ties response speed to memory depth, whereas Bi-Maxwell minimally separates them into fast and slow modes combined into one macroscopic stress.The fast mode governs short-time response, and the slow mode supplies the long-time tail before the spatial update law rearranges the weight matrix.
4 Experiment 1: the controlled kernel swap
The controlled experiment replaces only Muon’s single-timescale memory kernel with a two-timescale kernel, testing whether kernel shape—not merely longer average memory—improves training. Across frozen protocols, readouts, and hardware, the two-timescale variant reaches the target sooner, while several controls constrain the interpretation of the gain.
- 4.1 Speedup on a frozen protocol: Replacing only the memory kernel reduced bare tuned-Muon first crossing from 3250 to 3210 steps, a 40-step improvement.The two-timescale arm used A800 with n = 8, while the single-pole reference used official H100 logs with n = 10.
- 4.1 Speedup on a frozen protocol: The eight bare-stack seeds crossed between 3170 and 3210 steps without exclusions, so the advance was not driven by one trajectory.The reported per-seed crossing steps are {3210, 3180, 3175, 3200, 3190, 3210, 3170, 3190}.
- 4.1 Speedup on a frozen protocol: At steps 3200, 3225, and 3250, synchronized mean-loss statistics were 0.00436, 0.00413, and 0.00408, each above the 0.004 significance threshold.The comparison uses the reference arm minus the tested arm, so positive values indicate lower loss for the two-timescale arm.
- 4.2 Kernel shape, age and path: At matched mean age ¯n = 19, the two-timescale kernel beat the single-pole kernel, showing that kernel shape mattered beyond average memory depth.At ¯n = 30, the matched single pole crossed at 2775, 140 steps later than the two-timescale arm at 2635.
- 4.2 Kernel shape, age and path: The exploratory mean-age scan was non-monotonic, with best observed performance at ¯n = 30; both shorter and longer memories degraded the gain.The scan used ¯n ∈ {15, 19, 25, 30, 36, 42}.
- 4.3 Raw readout and hardware transfer: On the record stack, replacing only the internal kernel reached the target at step 2635 versus the 2690-step record, and the ordering persisted without Tail-EMA readout.With raw readout, the two-timescale and single-pole arms crossed at 2690 and 2735 steps, respectively.
- 4.3 Raw readout and hardware transfer: The record-stack gain reproduced across A800 and H100 arms, but the 40-versus-55-step difference between bare and record stacks cannot be attributed to one component.A same-hardware factorial experiment switching the memory kernel and Tail-EMA on and off would be required to quantify their interaction.
5 Experiment 2: mode dependence and protocol dependence of the optimal memory length
The optimal memory length depends on both mode strength and learning-rate protocol because drift and batch noise decay at different relative rates. These opposing mode preferences rule out a single-timescale memory and motivate a two-timescale kernel.
- Diagnostic: n⋆ is determined by the ratio T/D, so its change depends on the relative decay of drift and batch noise rather than drift alone.A faster fall in T decreases n⋆, whereas a faster fall in D increases it.
- Frozen protocol: All 40 band-level readings on the frozen-learning-rate baseline increased from the early to late window.The five-band n⋆ values span more than an order of magnitude across the trajectories.
- Protocol dependence: Under the constant high learning rate, Rb3 = 0.9187 and Rb4 = 0.9063, while all five ratios exceed 1 under the constant low learning rate.The high-rate protocol lengthens memory only for the strongest band and shortens it for middle and weak bands.
- Kernel implication: Opposite memory shifts across modes exclude any memory law with only one time constant at a given instant.A two-pole kernel supplies fast response and long memory simultaneously, while its branch allocation can vary with learning rate and training state.
- Mechanism: Both D and T decreased in all 160 band-protocol values, but their relative decay produced different memory profiles.In the constant-high-rate middle and weak bands, T falls faster; in the original and constant-low-rate protocols, D falls faster.
- Causal test: The fork experiment found a paired difference of +0.073 with t = 2.42, p = 0.046, and n = 8 pairs.The result supports a learning-rate effect on the preferred memory length, although the fork explains less than half of the from-scratch difference.
Discussion
The discussion interprets Bi-Maxwell as a physically derived replacement for single-pole momentum and connects training-stage memory growth to aging-like relaxation. It also identifies implementation costs and scope boundaries.
- Outcome: Replacing only the memory kernel reduced first-crossing steps from 2690 to 2635 on the frozen public benchmark.The bare tuned-Muon stack also reached the target 40 steps earlier.
- Physical interpretation: The two-timescale kernel is a positive superposition of fast and slow exponential decays, representing a coarse relaxation spectrum.The fast branch forgets within a few steps, while the slow branch persists for tens of steps.
- Training-stage dependence: The measured optimal-memory proxy lengthens from early to late training, matching the paper’s aging-like interpretation of the training medium.The five direction groups move in the same direction.
- Limitations: The paper’s interpretation of physical aging is limited because a stricter comparison would require waiting-time-dependent two-time correlations and response measurements.Fluctuation–dissipation violations and effective temperature are identified as additional nonequilibrium diagnostics.
- Practical guidance: The default Bi-Maxwell configuration adds about 680 MB of FP32 GPU memory at the 124 M-parameter scale, with wall-clock cost within rerun noise.The overhead comes from two extra momentum buffers and grows linearly with parameter count.
Methods
The methods fix the model, benchmark, statistical readout, and optimizer configurations, then compare a single-pole baseline with a two-timescale kernel while using read-only probes for memory diagnostics.
- Experimental setup: The experiments use GPT-2 with 124 M parameters on FineWeb, 524288 tokens per step, targeting held-out validation loss 3.28.The target is scored by the minimum number of steps needed to reach it.
- Statistical criterion: The benchmark requires arm-mean validation loss significantly below 3.28 at one-sided p < 0.01.The paper additionally uses the fixed readout (3.28 − L̄)√n ≥ 0.004 for comparing arms.
- Controlled comparison: The bare tuned-Muon comparison replaces only the internal memory kernel while keeping all other components unchanged.The frozen baseline uses Muon learning rate 0.025, weight decay 0.05, and an AdamW auxiliary.
- Kernel implementation: Bi-Maxwell uses two momentum states with βf = 0.85, βs = 0.98, w = 0.4385, and Ton = 1000.The corresponding fast and slow memory lengths are about 6 and 49 steps.
- Read-only probe: The K⋆ probe changes no parameters, gradients, or optimizer state while measuring 8 independent baseline trajectories.It estimates the optimal memory length from the inverse of a gain-matched quantity.
A Derivation of the independence conditions for collective modes
The derivation establishes when collective modes contribute independently to the output response. Nonzero cross terms couple modes, so independence requires their cross contribution to vanish.
- Input side: A mode-specific input can excite another mode unless the corresponding input-side coupling condition is satisfied.The derivation tests an input composed in the proportions of one mode and examines its amplitude in another.
- Output response: The output response for two modes is written as ẏ = b_k u_k + b_l u_l.The amplitudes b_k and b_l determine the two modal contributions.
- Cross contribution: The squared output amplitude contains the two modes’ own contributions plus a cross term.That cross term can make the modal responses reinforce or cancel.
- Independence condition: For arbitrary modal amplitudes, independent output contributions require the coefficient of the cross term to vanish.This condition makes the squared output amplitude equal the sum of the individual modal contributions.
B Completeness proof of the collective-mode decomposition
The appendix proves that any real driving-force matrix decomposes completely into orthogonal input-output mode pairs, including zero-force and degenerate cases. This decomposition supports the paper’s collective-channel representation and mode-wise output constraint.
- Completeness proof: Any real driving-force matrix admits a complete decomposition into orthogonal input-side and output-side mode pairs.The number of nonzero modes equals the matrix rank, representing the collective channels actually driven.
- Proof construction: The proof constructs the decomposition by orthogonally diagonalizing X^T X and retaining its positive-eigenvalue directions.The remaining basis directions carry no driving force, so the reconstructed matrix agrees with X on the full basis.
- Special cases: Zero driving force yields an empty decomposition and no driven channel requiring a rearrangement speed.For X = 0, every singular value is zero and the rank is zero.
- Special cases: Repeated eigenvalues do not change the decomposition’s form because the summed projection onto a degenerate eigensubspace is basis-independent.The maximum-dissipation condition likewise depends only on each mode’s scalar speed, not on the selected basis.
- Mode constraint: Under unit-norm mode vectors, the output-to-input amplitude ratio is |a_k|, making the output-gain limit |a_k| ≤ c mode-wise.This is the unified scale used in the main text.
C Per-seed data tables
The appendix tables report per-seed crossing steps for two-timescale and matched single-pole optimizer configurations across bare, record, replicated, and controlled arms. These are descriptive seed-level crossings, separate from arm-level benchmark pass decisions.
- Interpretation: Per-seed crossings are the first validation points at which an individual seed reaches 3.28, while arm passage uses the across-seed criterion separately.The arm criterion is (3.28 − L̄)√n ≥ 0.004.
- Table 3: Table 3 reports per-seed crossing steps for the bare tuned-Muon stack using the two-timescale kernel with n = 8.The table is paired with the bare-stack data file.
- Table 4: Table 4 reports per-seed crossing steps for the record stack using the two-timescale kernel on A800 with n = 8.The table corresponds to the record-stack data file.
- Table 5: Table 5 reports per-seed crossing steps for an independent H100 replication of the record-stack two-timescale kernel with n = 8.The entries come from H100 arm logs.
- Table 6: Table 6 reports per-seed crossing steps for the Ton = 700 record-stack arm on A800 with n = 8.These data come from the logs of the record 2655 submission.
- Table 7: Table 7 reports per-seed crossing steps for a matched single-pole control with β = 30/31 and n̄ = 30 on the record stack.The control uses A800 and n = 8.
D Methodological status of the fork protocol
The fork protocol compares trajectories sharing the same pre-fork history and correlated driving noise, making it useful for low-cost directional and ranking scans. It cannot provide effect sizes, which the paper therefore obtains from from-scratch multi-seed experiments.
- Protocol: The fork protocol changes a parameter on one branch after a shared checkpoint and compares short-window readout differences.Shared history and correlated noise make the comparison sensitive to directional and ranking differences.
- Limitations: Fork readings cannot serve as effect sizes because they inherit common-trajectory fluctuations, amplify short-window transients, and are biased by identical pre-fork histories.The paper reports that single-trajectory fork readings once overestimated the effect about 5-fold.
- Evidence discipline: The paper uses forks only for directions and parameter rankings, while effect sizes come from from-scratch multi-seed experiments.Appendix E independently confirms a fork-selected direction with 8 seeds.
E Dose scan of the kernel mean age and its multi-seed confirmation
The kernel mean-age scan finds a non-monotonic dose response, with an exploratory optimum near n̄ = 30 and degradation for shorter or longer memory. An independent record-stack check confirms that increasing n̄ from 30 to 41.8 delays first crossing and produces a positive paired loss difference.
- Dose scan: The fixed-step validation-loss difference is U-shaped across mean age, reaching its minimum near n̄ ≈30 and worsening at both ends.Figure 6 is a fork scan at fork@700 over window [2650, 2700], zeroed at the n̄ = 19 matched point.
- Multi-seed confirmation: The independent confirmation changes w from 0.4385 at n̄ = 30 to 1/6 at n̄ = 41.8, with the record-stack arm crossing at step 2735 instead of 2635.The paired difference at step 2635 is +0.006874 across 8 seeds, closely matching the curve’s predicted +0.00708.
- Scope and caveat: The two stacks produce opposite signs with magnitudes differing by a factor of 5, and the measurements are compatible with a Tail-EMA redundancy conjecture.The stacks differ in more than Tail-EMA alone, so isolating its contribution would require a factorial experiment not performed here.
F Assumption boundary of the output-perturbation budget and anisotropic generalizations
The main-text output budget is an explicit isotropic, Euclidean, worst-case constraint applied to layer outputs. Under these assumptions, the semi-orthogonalized direction is optimal, but other output measures need not yield the same response.
- The adopted budget uses Euclidean RMS amplitudes, applies to every input direction, and assigns one common safety threshold.
- These choices weight channels equally, protect against the most dangerous input direction, and make the safety boundary direction-independent.
- Under this isotropic worst-case budget, Muon’s ideal semi-orthogonalized direction releases mismatch potential fastest, but this conclusion is conditional.
F.2 An average output budget defined by the actual activation distribution
An average output budget weights structural motion by activation statistics and can produce an anisotropic response. The appendix treats this as a generalization separate from the paper’s fixed spatial baseline and temporal-kernel experiments.
- An average budget uses the activation second-moment matrix C to measure output perturbations over the training distribution rather than across every possible input direction.
- If C or S has zero modes, pseudo-inverses require compatibility conditions; otherwise the budget must be regularized.
- The resulting response is covariance-preconditioned: frequently occurring, high-variance input directions incur greater average output cost than rarely occurring directions.
- Downstream functional sensitivity can also be represented through a Jacobian-derived matrix S, making motion cost depend on both input activation and output amplification.
- The paper does not use the generalized anisotropic response and keeps the Euclidean worst-case spatial response fixed while testing temporal memory.
G.4 Robustness of the readings
Robustness checks support the broad aging pattern and the positive strength-rank tilt, while showing that the strongest band is sensitive to late-window placement. The sensitivity is linked to a synchronized late rise in estimated noise.
- Across nine late-window positions, the weak band remains below one under high learning rate, while the strength-rank tilt stays positive.
- The strongest band exceeds one only for sufficiently late windows, with the preregistered [2200, 2700] window occurring after its crossing.
- Across five ratio computations, bands 2–5 remain below one and the strength-rank tilt remains positive, but the strongest-band result fails when the late window moves earlier.
- After step 2505, the strongest band’s n⋆ rises sharply while the other bands barely move, and this pattern appears synchronously across trajectories.
- The estimator is especially sensitive in the strongest band because it has the fewest directions and therefore the largest noise-subtraction term.
- Across 8 trajectories, all 40 late-to-early band ratios exceed one, ranging from 1.13–2.43, though bands within a trajectory are not independent.
G.5 Supplementary decompositions
Supplementary decompositions separate common profile shifts from relative mode reallocation and reveal module-specific behavior. They also delimit what the kernel-age probe and frequency-domain argument can establish.
- The fork experiment attributes relative strength-rank reallocation to learning rate, but does not establish the cause of the common whole-profile shift.
- Late-window departures from the original schedule are much larger than early-window departures and increase toward weaker bands.
- Under constant high learning rate, attention output projection drives shortening in middle and weak bands, while MLP matrices move toward longer memory.
- The aggregate strength profile cannot separate these module-level contributions because band and module groupings cut across tensors differently.
- The frequency-domain analysis demonstrates a nontrivial mechanism but does not show that the chosen kernel parameters match measured mode demands or locate modes on that frequency axis.
- A direction-matched comparison found the non-convex α = 5 arm worse than the convex control by +0.00061 at step 3210.
- Because bands are recomputed rank quantiles and estimator biases affect absolute K⋆, the analysis interprets relative within-band changes rather than calibrated absolute memory lengths.