🤖 machine learning

A Physical Response-and-Memory Model for Muon Optimization

This paper proposes a physical response-and-memory model for training dynamics that explains the efficacy of the Muon optimizer and introduces the Bi-Maxwell optimizer, which utilizes a dual-timescale memory kernel to achieve faster convergence on large language model benchmarks.

Original authors: Yinze Hu, Hongjun Xiang, Xingao Gong, Hongyu Yu

Published 2026-08-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yinze Hu, Hongjun Xiang, Xingao Gong, Hongyu Yu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast, silent architecture of modern artificial intelligence, the most expensive and time-consuming task is not the design of the model itself, but the act of teaching it. To train a large language model, engineers feed it billions of words and adjust its internal settings, known as weights, to reduce errors. This process relies on a mathematical tool called an optimizer, which acts as the teacher's hand, deciding exactly how much to shift those weights after every single lesson. For years, the most effective teachers have been chosen through trial and error, guided by engineering intuition rather than a deep understanding of why they work. The goal is simple: with a fixed amount of computing power, how low can the error rate go, and how quickly can it get there? The answer depends entirely on the rules the optimizer follows to turn a mistake into a correction.

A recent study by researchers at Fudan University and Tongji University offers a new way to see this process. Instead of viewing the training of a neural network as a series of abstract calculations, they treated the network's weight matrix as a physical object with memory, much like a piece of metal or a gel that responds to force. They proposed that when the network makes an error, it creates a kind of internal stress, similar to the tension in a stretched spring. The optimizer's job is to release this stress efficiently. By applying the laws of physics to this digital medium, the researchers derived a new set of rules for how the network should learn. They found that the most effective way to update the weights is to move in a direction that dissipates this stress as quickly as possible, while respecting a safety limit on how much the network's output can change at once. This physical perspective explained why a specific, recently popular method called Muon works so well, and it revealed a flaw in how that method remembers past lessons.

The researchers discovered that the standard way these optimizers remember the past is too simple. They usually average out recent mistakes using a single, uniform timescale, as if the material they are made of relaxes its tension at only one speed. However, the researchers argued that real materials, and by extension, complex neural networks, have a more complicated internal structure. They possess both fast, local reactions and slow, deep structural changes. To test this, they replaced the single-speed memory of the standard optimizer with a new version they called Bi-Maxwell, which uses two different speeds: one fast to catch immediate changes and one slow to hold onto long-term patterns. When they tested this new approach on a public benchmark for training language models, the results were immediate and measurable. The new optimizer reached the target level of accuracy in 2,635 steps, beating the previous record of 2,690 steps. This might sound like a small difference, but in the world of training massive models where thousands of computers run for weeks, saving even a few dozen steps translates to significant savings in time and energy.

To ensure this improvement was not a fluke, the team performed a rigorous series of checks. They kept every other part of the training process exactly the same, changing only the memory kernel, the part responsible for remembering past gradients. They ran the experiment on different types of computer hardware and with different random starting points, and the new optimizer consistently outperformed the old one. They also investigated why the two-speed memory worked better by measuring how the network's needs changed over time. They found that early in training, the network learns quickly and needs a short memory to stay agile. As training progresses, the learning slows down, and the network benefits from a longer memory to smooth out the path. The single-speed optimizer cannot adapt to this shift; it is stuck with one setting for the entire journey. The two-speed version, however, naturally accommodates both the fast reactions needed at the start and the deep, steady memory required later.

The study also clarified what the optimizer is actually doing. It showed that the most efficient update direction is not just a random guess but a specific geometric path that maximizes the release of error potential without causing the network to become unstable. This direction is what the Muon optimizer uses, but the researchers provided the physical reasoning for why it works: it is the path of maximum energy dissipation under a safety constraint. Furthermore, they demonstrated that the "memory" in these systems is not just a mathematical trick to smooth out noise, but a representation of accumulated internal stress. Just as a physical object takes time to settle after being pushed, the neural network needs time to integrate the lessons from its mistakes. By treating the network as a responsive medium with a spectrum of relaxation times, the researchers were able to design an optimizer that respects the natural rhythm of learning.

This work suggests that the future of training artificial intelligence may lie in borrowing more deeply from the physical sciences. By viewing the learning process through the lens of materials science and thermodynamics, the researchers were able to move beyond trial and error and derive rules that are grounded in the fundamental behavior of the system. They did not just find a faster way to train a model; they provided a framework for understanding why certain methods work and how they can be improved. The success of the Bi-Maxwell optimizer, which simply added a second timescale to the memory, proves that even small, physically motivated changes can yield significant gains. As language models grow larger and more complex, the efficiency of their training becomes increasingly critical, and this physical approach offers a promising path forward for making that training faster, more stable, and more efficient.

Technical Summary: A Physical Response-and-Memory Model for Muon Optimization

Problem Statement

Training large language models (LLMs) incurs immense computational costs, making the efficiency of the optimization process critical. The choice of optimizer determines the training speed, stability, and the final loss achievable under a fixed compute budget. While modern optimizers like Muon have achieved state-of-the-art performance by semi-orthogonalizing the momentum matrix, their design has largely been driven by engineering intuition and empirical benchmarking rather than a systematic physical theory. Two fundamental questions remain unanswered by existing theory:

Why does the semi-orthogonalized direction (used in Muon) work effectively?
How long should the momentum memory average over past gradients? Currently, this is determined by a single exponential timescale (a single relaxation time), selected heuristically.

Methodology: A Physical Model of Training

The authors propose a physical model treating the weight matrix during training as a responsive medium with memory. This framework derives the update rules from physical principles rather than heuristic selection.

1. Spatial Response: The Medium and the Update Direction

The weight matrix $W$ is modeled as a linear responsive medium where inputs $x$ produce outputs $y = Wx$ . The training objective is defined as a mismatch potential $\Phi$ (e.g., squared error between output and target), analogous to the potential energy of a spring.

Driving Force: The gradient of the mismatch potential with respect to the weights acts as a conjugate driving force $X = -\partial \Phi / \partial W$ .
Output-Side Safety Budget: The authors impose a constraint on the output-side perturbation. Specifically, the rate of change of the output $\dot{y} = Vx$ (where $V$ is the structural rearrangement velocity) is bounded by an output-gain limit $\epsilon$ for all input directions.
Maximal Dissipation: The goal is to maximize the rate of potential release ( $P = -\dot{\Phi}$ ) subject to this output budget.
Derivation of Semi-Orthogonalization: By decomposing the driving force into independent collective modes (via Singular Value Decomposition) and applying the output-gain constraint, the authors prove that the optimal rearrangement velocity $V$ must update all driven modes at the same maximum safe speed. This results in a velocity matrix proportional to the orthogonal polar factor of the driving force. This provides a physical derivation for Muon's semi-orthogonalized update rule: it is the maximally dissipative response under an output-side safety budget.

2. Temporal Response: Internal Stress and Memory

The authors argue that the momentum matrix represents the internal stress accumulated by the medium, rather than the inertia of a particle.

Stress Relaxation: When an external force (gradient) is applied, internal stress does not adjust instantly; it relaxes over time. The standard momentum update ( $M_t = \beta M_{t-1} + (1-\beta)X_t$ ) is identified as a first-order relaxation equation (Maxwell model) where the momentum coefficient $\beta$ corresponds to a single relaxation time $\tau$ .
Limitation of Single Timescale: Real materials (and complex neural networks) possess multiple internal rearrangement modes with different relaxation times. A single exponential kernel forces a trade-off: a short $\tau$ allows fast response but poor memory filtering, while a long $\tau$ provides deep memory but slow response to changes.
Bi-Maxwell Model: To resolve this, the authors propose a two-timescale kernel (Bi-Maxwell). This models the internal stress as a superposition of a fast relaxation mode ( $\tau_f$ ) and a slow relaxation mode ( $\tau_s$ ). The total memory is a weighted sum: $M_{BM} = w M_f + (1-w) M_s$ . This allows the optimizer to respond quickly to persistent signals (via the fast mode) while retaining long-term structural memory (via the slow mode).

Key Contributions

Physical Derivation of Muon: The paper derives the semi-orthogonalized update direction of Muon from first principles, identifying it as the solution to a constrained optimization problem (maximizing potential dissipation under an output-gain budget).
Physical Interpretation of Momentum: Momentum is reinterpreted as the internal stress of a viscoelastic medium. The paper argues that the single-timescale assumption is a "single-pole approximation" of a broader relaxation spectrum.
The Bi-Maxwell Optimizer: Based on the two-timescale physical model, the authors introduce the Bi-Maxwell optimizer. It replaces the single-timescale momentum kernel of Muon with a two-timescale kernel (fast and slow modes) while keeping all other spatial components (Newton-Schulz orthogonalization, etc.) unchanged.
Aging Phenomenon in Training: The paper identifies that the "optimal memory length" is not static but grows as training progresses (analogous to physical aging in glassy materials). Early in training, gradient directions change rapidly (requiring short memory); later, they change slowly (requiring longer memory).

Experimental Results

The authors validate the Bi-Maxwell model on the modded-nanogpt public benchmark (124M parameter GPT-2 model) using two configurations: a "bare tuned-Muon" stack and a "record stack" (including SOAP, Tail-EMA, and other advanced components).

Performance Gain:
- On the bare tuned-Muon stack, the Bi-Maxwell optimizer reached the target validation loss (3.28) in 3210 steps, compared to 3250 steps for the single-pole baseline (a 40-step improvement).
- On the record stack (which held the state-of-the-art record of 2690 steps as of July 2026), the Bi-Maxwell optimizer reached the target in 2635 steps, breaking the existing record by 55 steps.
- These results were consistent across multiple random seeds (n=8) and different hardware (A800 and H100 GPUs).
Controlled Experiments:
- Kernel Shape vs. Mean Age: Experiments showed that the gain comes from the two-timescale shape itself, not merely from increasing the average memory depth. A single-pole kernel matched to the same mean age as the Bi-Maxwell kernel performed significantly worse (crossing 140 steps later).
- Enabling Window: The slow mode should not be enabled from step 0. The optimal performance was achieved when the two-timescale kernel was activated at step 1000 ( $T_{on}=1000$ ), suggesting the medium needs to "age" before the slow memory becomes beneficial.
- Read-Only Probes: Using a diagnostic probe, the authors measured the "optimal memory length" ( $n^*$ $n^{*}$ ) across different modes (ranked by strength). They found that:
  - The optimal memory length increases from early to late training.
  - Different modes (strong vs. weak) exhibit opposite trends in memory preference depending on the learning rate schedule, confirming that a single fixed timescale cannot satisfy all modes simultaneously.

Significance and Claims

The paper claims to provide the first systematic physical theory connecting the spatial update rule (semi-orthogonalization) and the temporal memory kernel of modern optimizers.

Theoretical Unification: It unifies the understanding of Muon's spatial efficiency and the role of momentum as internal stress relaxation.
Algorithmic Improvement: It demonstrates that replacing the memory kernel with a minimal two-timescale form (Bi-Maxwell) yields significant, reproducible speedups on public benchmarks without altering the core spatial structure of the optimizer.
Physical Analogy: The work draws a parallel between training dynamics and the "aging" of glassy materials, suggesting that the optimal memory length is a state variable that evolves with the training stage.
Modesty: The authors note that the physical model relies on specific assumptions (isotropic, worst-case output budget) and that the Bi-Maxwell parameters (weights and timescales) are currently static approximations. They acknowledge that the true temporal constitutive law may depend on the specific training state and mode strength, and that Bi-Maxwell represents the "lowest-order static basis" for this more complex reality.

In summary, the paper argues that viewing the weight matrix as a viscoelastic medium with a multi-scale relaxation spectrum provides a principled path to improving optimizers, moving beyond heuristic tuning to physically grounded design.