△ Mechanism confirmed, baseline not beaten 2026

Bi-Maxwell Muon

Implementation & benchmark of arXiv:2608.22994 — A Physical Response-and-Memory Model for Muon Optimization

Usefulness7/10

Difficulty4/10

Novelty6/10

Optimization Speedup Stability Accuracy Optimizer Transformer Training dynamics Dynamical systems Optimization Linear algebra

Source paper: A Physical Response-and-Memory Model for Muon Optimization arXiv:2608.22994 ⓘ · analyzed Sep 1, 2026

AI-generated research hypothesis, automatically tested. Not peer-reviewed.

Idea description

Replace Muon's single momentum matrix with a weighted mixture of fast and slow relaxation modes. The fast mode tracks rapidly changing gradients while the slow mode preserves a longer-horizon direction; their mixture is semi-orthogonalized and applied as the matrix update.

Formulas

$$\tau_j\frac{dM_j}{dt}=X(t)-M_j(t).$$

$$M_{\mathrm{BM}}=wM_f+(1-w)M_s.$$

$$M_{j,t}=\beta_{j,t}M_{j,t-1}+(1-\beta_{j,t})X_t,\qquad \beta_{j,t}=\exp(-\eta_t/\tau_j).$$

$$W_{t+1}=W_t-\alpha_t\,\mathcal{O}\!\left(wM_{f,t}+(1-w)M_{s,t}\right).$$

Mathematical statement

The paper models each internal stress mode by the relaxation equation $\tau_j\frac{dM_j}{dt}=X(t)-M_j(t)$, where $X(t)$ is the matrix force supplied by the current gradient, $M_j(t)$ is the stress/momentum matrix of mode $j$, and $\tau_j>0$ is its relaxation time. For two modes, the macroscopic stress is $M_{\mathrm{BM}}=wM_f+(1-w)M_s$, where $M_f$ and $M_s$ are fast and slow stresses and $0\leq w\leq1$. Discretizing with step size $\eta_t$ gives $M_{j,t}=\beta_{j,t}M_{j,t-1}+(1-\beta_{j,t})X_t$, with $\beta_{j,t}=\exp(-\eta_t/\tau_j)$; therefore $\tau_f<\tau_s$ implies $\beta_f<\beta_s$. Use $M_{\mathrm{BM},t}=wM_{f,t}+(1-w)M_{s,t}$ as Muon's pre-orthogonalization matrix. If $\mathcal{O}(A)$ denotes Muon's semi-orthogonalization or polar-factor approximation, the parameter update is $W_{t+1}=W_t-\alpha_t\mathcal{O}(M_{\mathrm{BM},t})$, with $W$ the weight matrix and $\alpha_t$ the parameter learning rate. The physical interpretation is that several coupled internal degrees of freedom diagonalize into independent first-order modes with different relaxation times.

Implementation notes

Integrate this at the optimizer state for every trainable 2-D weight matrix that currently uses Muon. Keep two momentum buffers, Mf and Ms, with the same shape and dtype policy as Muon's existing momentum buffer. At step t, obtain the matrix gradient G; optionally apply the same gradient normalization or distributed aggregation already used by the Muon baseline, and call the resulting matrix X. Compute betaf = exp(-lr_tau / tauf) and betas = exp(-lr_tau / taus), where lr_tau is the optimizer time increment used to interpret the relaxation equation. A simpler implementation can use fixed coefficients directly: betaf < betas, for example betaf=0.9 and betas=0.99. Update Mf = betaf*Mf + (1-betaf)*X and Ms = betas*Ms + (1-betas)*X; then form M = w*Mf + (1-w)*Ms. Pass M through exactly the same semi-orthogonalization routine as baseline Muon, such as its Newton-Schulz/polar approximation, and apply W -= alpha * O(M). Do not add Adam-style coordinatewise second moments in the first experiment, so the effect of memory is isolated. The paper's mathematics determines the exponential recursion and the relation between relaxation time and beta; the values of tauf, taus, and w are empirical hyperparameters. First run a cheap ablation on a 100M-parameter decoder-only transformer trained on a fixed subset of C4 or WikiText-103, comparing Muon, Bi-Maxwell Muon, and AdamW at equal tokens and equal optimizer FLOPs. Try (betaf, betas, w)=(0.9,0.99,0.5), (0.95,0.995,0.5), and weights favoring the fast mode early. Record training loss versus optimizer steps, validation perplexity, gradient-update cosine similarity, and wall-clock overhead. Success means lower loss at equal steps/FLOPs, or reaching a fixed validation perplexity in fewer steps without instability; failure is no improvement over single-buffer Muon after tuning each baseline fairly.

Verification

Mechanism confirmed, baseline not beaten

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Confirmed ✓

Stage 2 · Practical benchmark result: Significantly worse ✗

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 9/10

Built a reusable two-buffer Bi-Maxwell Muon implementation with fast/slow EMA relaxation and Newton-Schulz semi-orthogonalization. The mechanism manifested: discrete half-life crossings matched predictions for all tested betas, impulse log-decay slopes matched log(beta) to numerical precision, and frequency-response gains agreed within 0.02%-8.6%. In the tiny regression comparison, Bi-Maxwell did not improve optimization loss over Muon or AdamW, so this is mechanism validation rather than evidence of a practical training win.

Agent confidence: 9/10

Baseline: Muon: final loss 0.001597, step-100 loss 1.532585, mean gradient/update cosine 0.5865, 0.141 s; AdamW: final loss 0.001178, step-100 loss 0.001263, cosine 0.3449, 0.068 s.
Idea: Bi-Maxwell (beta_fast=0.90, beta_slow=0.99, w=0.5): final loss 0.004175, step-100 loss 1.533368, cosine 0.5502, 0.137 s. Bi-Maxwell (0.95,0.995,0.5): final loss 0.023674, step-100 loss 1.537552, cosine 0.5018, 0.152 s. Predicted-versus-observed half-life integer crossings matched 5/5; impulse slopes matched to numerical precision; frequency-gain relative errors were 8.6%, 2.35%, 0.22%, and 0.024%.

Limitations: The optimizer test was a small full-batch linear regression, not a decoder-only transformer or C4/WikiText experiment. It did not test distributed aggregation, mixed precision, GPU behavior, equal-FLOP transformer training, extensive hyperparameter tuning, validation perplexity, or multiple random seeds.

How to run: python3 verify.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.22994) Sep 3, 2026

Significantly WORSE than the tuned baseline (+11.1%, p=0.0081, wins 0 / 8 paired seeds).

Benchmark: Tabular regression (Friedman)
Model: mlp_tiny
Paired seeds: 8
Baseline mean: 21.5773
Idea mean: 23.9797
Effect (Δ): 2.4024 (+11.1%; negative = idea better)
Wins: 0 / 8 paired seeds
p-value: 0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect: ±2.5%
Mechanism: Confirmed ✓
Practical effect: Hurts
Baseline tuning: swept over 8 configs

Limitations:

Only the structurally matched tabular optimizer track was tested; wall-clock and optimizer FLOP overhead were not primary metrics.

How to run: python3 stage2_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Sep 3, 2026 latest Failed ✗

Agent confidence: 10/10

Bi-Maxwell Muon was evaluated on the registered tabular Friedman#1 track with identical mlp_tiny systems, a tuned baseline sweep, and 8 paired seeds. The trained-model relaxation signature was confirmed, but Bi-Maxwell was significantly worse on test MSE, so there was no practical win.

Baseline: Muon best lr=0.005, weight_decay=0.0001; mean test MSE 21.5773.
Comparison: delta_mean=+2.40236; p_value=0.0081; idea_wins=0/8.
Idea: Bi-Maxwell best lr=0.003, weight_decay=0.0001, beta_fast=0.90, beta_slow=0.99, weight_fast=0.5; mean test MSE 23.9797.

Benchmark result