# Bi-Maxwell Muon

- ID: 2693
- Canonical URL: https://synthcore.org/idea/2693/bi-maxwell-muon
- API JSON: https://synthcore.org/api/idea/2693.json
- API Markdown: https://synthcore.org/api/idea/2693.md
- Verification status: unverified
- Source: [arXiv:2608.22994](https://arxiv.org/abs/2608.22994)
- Category: optimization
- Solves: speedup, stability, accuracy
- ML areas: optimizer, transformer, training-dynamics
- Math tags: dynamical-systems, optimization, linear-algebra
- Ratings: usefulness 7/10; difficulty 4/10; novelty 6/10

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

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

## Key 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).$$

## 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

- Status: unverified
- Mechanism evidence: yes
- Mechanism confirmed: no
- Verdict: 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.

### Mechanism check

- Verdict: 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.
- Confidence: 9/10
- 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.

## Artifacts

- [bench_bimaxwell.py](https://synthcore.org/code/1039/bench_bimaxwell.py)
- [bench_report.json](https://synthcore.org/code/1039/bench_report.json)
- [bi_maxwell_muon.py](https://synthcore.org/code/1039/bi_maxwell_muon.py)
- [experiment.py](https://synthcore.org/code/1039/experiment.py)
- [report.md](https://synthcore.org/code/1039/report.md)
- [results.json](https://synthcore.org/code/1039/results.json)
- [stage2_bench.py](https://synthcore.org/code/1039/stage2_bench.py)
- [verification.json](https://synthcore.org/code/1039/verification.json)
- [verify.py](https://synthcore.org/code/1039/verify.py)
- [Download all files as ZIP](https://synthcore.org/download/1039)

## Disclaimer

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