A physics-motivated two-timescale tweak to Muon ('Bi-Maxwell') reaches the NanoGPT target loss notably earlier than the single-timescale baseline, with theory linking the gain to a two-timescale stress relaxation and empirical probes showing the optimal memory length grows during training.

Main Finding

The paper builds a physical model treating a trainable weight matrix as a responsive medium with memory, and uses that model to (1) derive why Muon’s semi‑orthogonalized momentum direction is effective (it is the maximally dissipative response under a constrained output‑side safety budget), and (2) explain the role and optimal timescale(s) of momentum (momentum = accumulated internal stress whose relaxation spectrum sets the optimal memory kernel). From this they propose Bi‑Maxwell, an optimizer that replaces the single‑timescale EMA momentum with a two‑timescale (fast + slow) stress memory. In controlled experiments on the public NanoGPT optimizer benchmark, swapping only the memory kernel to the two‑timescale form reduced steps-to-target: Bi‑Maxwell hit the target at step 2635 vs the standing record 2690 (55 steps earlier), and ~40 steps earlier than a tuned Muon baseline without additional add‑ons. Read‑only probe measurements across 8 independent runs are consistent with the model’s prediction that the optimal memory length should grow with training stage.

Key Points

• Physical picture
  • Treat a weight block W as a responsive medium; define a mismatch potential Φ (batch loss), with conjugate driving force X = −∂Φ/∂W (the backprop gradient).
  • Structural rearrangement rate V = ˙W produces an output perturbation Vx; impose an output‑side safety budget (a bound on Aout(Vx) relative to Ain(x)).
  • Under that constraint, the direction of V that maximizes instantaneous dissipation (drop in Φ) is the semi‑orthogonalized momentum direction — providing a principled derivation of Muon’s update direction.
• Momentum and memory
  • Momentum is interpreted as internal stress of the medium accumulated over past gradients.
  • The duration / kernel of averaging (EMA) is determined by the stress relaxation spectrum of the medium. Real media typically relax on multiple timescales; the minimal nontrivial discretization is two timescales (fast + slow).
• Bi‑Maxwell optimizer
  • Keeps Muon’s spatial components (orthogonalization, scaling, etc.) intact and replaces the single‑timescale momentum with a two‑timescale (bi‑exponential) stress memory.
  • This modification alone (kernel swap) produced measurable speedups on a heavily tuned public benchmark.
• Empirical support
  • Controlled kernel‑swap experiments held all other optimizer components constant and replaced only the memory kernel; Bi‑Maxwell reached the benchmark target sooner (2635 steps vs 2690 record).
  • A read‑only probe measured a proxy for the optimal memory length across 8 independent training trajectories; results are consistent with the prediction that optimal memory length increases later in training.
• Scope and caveats
  • The derivation and conclusions are conditional on explicit model assumptions (responsive medium, output‑side safety budget, timescale separation). The authors discuss boundaries and relaxations in appendices.
  • Experiments are on the NanoGPT optimizer public benchmark and on the Muon stack; generalization to other architectures/datasets at large scale remains to be validated.

Data & Methods

• Theoretical derivation
  • Define mismatch potential Φ (e.g., 1/2 Σ eα^2) and driving force X = −∂Φ/∂W (rank‑one for linear output; equals minibatch backprop gradient for general blocks).
  • Constrain per‑input output gain: require Aout(Vx) ≤ ε Ain(x) for all input directions x (an output‑side "safety budget").
  • Solve maximization of instantaneous dissipation P = Σij Xij Vij under that constraint → yields semi‑orthogonalized update direction (polar factor of momentum matrix) familiar from Muon.
  • Model stress relaxation with a relaxation spectrum H(τ); discretize to a two‑timescale kernel (fast + slow Maxwell elements) to obtain Bi‑Maxwell memory.
• Experiments
  • Controlled kernel‑swap: on the modded‑NanoGPT public optimizer benchmark (community‑tuned Muon + add‑ons baseline), the authors swapped only the memory kernel from single EMA to two‑timescale EMA while holding all other components constant.
  • Measured steps‑to‑target loss and compared to the standing record and tuned baselines; also tested the swap on a bare tuned‑Muon baseline (no SOAP/add‑ons).
  • Read‑only probe: step‑by‑step measurement of a proxy for the optimal memory length across 8 independent training runs to observe stage‑dependent changes; cross‑hardware transfer and readout artefacts checked.
  • Appendices and Methods contain proofs, protocol details, per‑seed data and checks for alternative explanations.
• Quantitative results reported
  • Bi‑Maxwell reached target at step 2635 vs the standing record 2690 (difference 55 steps).
  • On the bare tuned‑Muon baseline, Bi‑Maxwell reached the target ~40 steps earlier than the single‑timescale momentum.

Implications for AI Economics

• Direct compute and cost savings
  • Fewer optimization steps to reach a target loss translate linearly into less compute, energy, and wall‑time for a training run. Even modest percentage reductions in steps scale to substantial dollar and energy savings at frontier scale (thousands of GPUs for weeks).
  • Because the benchmark baseline is heavily community‑tuned, measurable gains from a simple kernel change imply room for further low‑cost efficiency improvements in production training stacks.
• Operational and stability benefits
  • A physically principled optimizer that explains stability and memory behavior may reduce ad‑hoc hyperparameter tuning costs and the number of failed or unstable runs, decreasing waste and operational overhead.
  • Stage‑dependent optimal memory suggests adaptive or probe‑driven momentum scheduling: implementing an adaptive memory length could yield additional runtime savings and reduce human tuning time.
• Broader economic impacts
  • Small per‑run savings compound across many models and repeated experiments (research + production retraining), improving ROI on compute infrastructure and potentially lowering barriers to experimentation for smaller labs.
  • The physical modeling approach provides a scalable design pathway: principled optimizer innovations (e.g., multi‑timescale memory) are likely cheaper to deploy than large architecture/capacity scaling, so they can be high‑leverage for cost‑constrained organizations.
• Cautions for deployment
  • Current empirical evidence is on a public NanoGPT benchmark and Muon stacks; firms should validate Bi‑Maxwell and the stage‑adaptive memory idea at their target scale and datasets before relying on the gains for cost projections.
  • Theoretical assumptions may not map exactly to all architectures and training regimes; further replication and stress‑testing (different model sizes, datasets, batch sizes, distributed settings) are advisable.

Overall, the paper offers a compact physical theory linking the geometry of Muon updates and the temporal design of momentum, produces a minimal and implementable optimizer modification (two‑timescale memory), and shows measurable, benchmarked step savings that have direct economic relevance when scaled to large training runs.