Physical Muon: Orthogonalization as an Equilibrium Computation

Yuren Hao

cs.LGSep 29, 2026

Physical Muon: Orthogonalization as an Equilibrium Computation

Authors: Yuren Hao

Organizations: University of Illinois at Urbana-Champaign

Abstract

Physical neural networks and analog in-memory computing could reduce the energy cost of neural network training. Realizing this potential, however, requires optimizers that combine effective learning with physical implementability. SGD fits local analog updates but struggles on transformers, while Adam family is unstable against analog bias. Muon offers strong training performance, but its Newton--Schulz orthogonalization relies on dense matrix-matrix products. To address this obstacle, we introduce Physical Muon, which computes the orthogonalization as the equilibrium of a continuous-time flow. Random probes approximate the flow using matrix-vector products, reciprocal reads, and local rank-1 writes. To test whether this replacement preserves training performance, we evaluate it on a 10.95M-parameter transformer. The dense flow's mean validation cross-entropy is 0.0085 above Newton--Schulz across nine seeds per method; the probe implementation is 0.0188 above the control across two seeds. Circuit simulations further reproduce the flow dynamics and yield comparable training behavior.

Figures & tables

Figure 1: Physical execution of ( 4 ). Left: reciprocal reads form the row error for local rank-1 writes. Right: a portion of the

X

array and an enlarged behavioral cell. A capacitor stores

X_{ij}

as a voltage that controls its effective read conductance; signed write currents update the state. Section 4 details circuit validation.

Solver	$n$	Validation CE	$\Delta$ CE	Upper bound
NS5 control	9	$5.0124\pm 0.0081$	—	—
Dense flow	9	$5.0209\pm 0.0054$	$+0.0085$	$0.0142$

Table 1: We perform an equivalence test using the two one-sided tests (TOST) procedure and find that NS5 and the dense flow have statistically equivalent validation CE within a margin of

\pm 0.0426

(

p=2.4\times 10^{-8}

). CE is mean

\pm

sample std across

n

training seeds.

\Delta

CE is relative to NS5; the upper bound is its one-sided 95% confidence bound.

Figure 2: Left: directly reusing NS5 outputs degrades training as the reuse interval grows (one seed). Right: singular-value responses of NS5 and the flow at two integration times; vertical lines mark the approximate 95% thresholds. Shading covers the measured NS5 outputs.

Figure 3: Circuit-in-the-loop training on the four-class synthetic task. Left: training CE for the three Muon implementations defined above. Right: cosine similarity between their orthogonalized update directions and the NS5 reference for each of the two weight matrices. The momentum matrix for $W_{2}$ has rank three.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Optimizer	CE
Muon ( Jordan et al., 2024 )	4.9190
OLion ( Wang et al., 2026 )	4.9263
OLion, 3 NS rounds	4.9523
OLion, 2 NS rounds	5.0043
AdamW ( Loshchilov and Hutter, 2019 )	5.0707
OLion, 1 NS round	5.1503

Appendix

Table 2: Best validation CE for 11 optimizers under one protocol

Stale interval (steps)	5	10	20
Full matrix	0.60	0.39	0.19
Leading-16 subspace	0.89	0.78	0.52

Appendix

Table 3: Cosine between stale and fresh NS5 outputs

Figure 4: Flow–control CE gap against integration time. Filled circles:

\eta=0.25

; open circles:

\eta=0.5

; diamond:

T=400

, nine-seed mean

\pm

std; square: probe flow,

K=32

\eta=0.15

T=417

, two-seed mean. The grey band is the NS5 seed spread.

Seed	1	2	3	4	5	6	7	8	9
NS5 control	5.0048	5.0084	5.0183	5.0080	5.0233	5.0248	5.0144	5.0035	5.0064
Flow	5.0251	5.0152	5.0217	5.0185	5.0178	5.0285	5.0124	5.0220	5.0269

Appendix

Table 4: Best validation CE per seed, NS5 control and flow (

T=400

\eta=0.5

)

$\eta$	$T$	$\eta T$	$\hat{s}^{*}$	$\Delta$ CE
0.25	50	12.5	0.144	$+0.136$
0.25	100	25	0.072	$+0.076$
0.25	200	50	0.036	$+0.033$
0.5	200	100	0.018	$+0.021$
0.5	400	200	0.009	$+0.0085^{*}$

Appendix

Table 5: Flow degradation against the NS5 control across integration depths (seed 1)

Batch size	24	12	6	4	3
$\Delta(\text{flow}-\text{ctl})$	$+0.0085$	$+0.0175$	$+0.0078$	$+0.0219$	$+0.0149$

Appendix

Table 6: Flow–control gap across batch sizes

Width	128 (10.95M)	192	256 (26.4M)	384 (59M)
$\Delta$	$+0.0085$	$+0.0114$	$+0.0034$	$+0.0112$
Seeds	9	3	2	1

Appendix

Table 7: Flow–control gap across widths at

T=400

Budget $B$	2k	5k	10k	20k	40k	80k
$128\times 128$ attention	0.708 (16)	0.802 (16)	0.850 (16)	0.888 (16)	0.919 (32)	0.945 (32)
$384\times 128$	0.806 (8)	0.896 (8)	0.930 (16)	0.951 (16)	0.965 (32)	0.973 (64)
$128\times 384$	0.677 (16)	0.803 (16)	0.873 (32)	0.912 (32)	0.946 (64)	0.953 (64)

Appendix

Table 8: Best cosine against NS5 (best

K

in parentheses) across array-pass budgets; entries above the 0.9 screen in bold

	Persistent gain			Persistent offset		Write noise
	5%	10%	20%	1%	5%	10%	30%
$128\times 128$ attention	$+0.0004$	$+0.0008$	$+0.0013$	$-0.0004$	$-0.0129$	$-0.0006$	$-0.0052$
$384\times 128$	$+0.0002$	$+0.0003$	$+0.0001$	$-0.0001$	$-0.0014$	$-0.0002$	$-0.0017$
$128\times 384$	$+0.0005$	$-0.0008$	$-0.0106$	$-0.0002$	$-0.0073$	$-0.0008$	$-0.0069$

Appendix

Table 9: Change in screen value (cosine against NS5) under injected defects

Method	5%	10%	20%
NS5	0.913–0.940	0.397–0.447	0.366–0.424
PolarExpress, 5 rounds	0.366–0.429	2 of 3 diverged; 0.389	3 of 3 diverged
Flow (probe form)	0.9986–0.9999	0.9945–0.9998	0.9763–0.9990

Appendix

Table 10: Cosine between defected and clean output under persistent per-pass gain error

Budget per matrix	$E$ (pitch)	$E$ (stress)	Digital / pitch	Digital / stress
$B=10$ k (measured, width 128)	0.076 J	0.174 J	18.3–39.7 $\times$	8.0–17.2 $\times$
$B=20$ k (measured, width 128)	0.151 J	0.348 J	9.2–19.8 $\times$	4.0–8.6 $\times$
$B=40$ k (measured, width 128)	0.302 J	0.696 J	4.6–9.9 $\times$	2.0–4.3 $\times$
$B=60$ k (rank-scaled, $r=768$ )	0.453 J	1.044 J	3.1–6.6 $\times$	1.3–2.9 $\times$
$B=240$ k (rank-scaled, $r=768$ )	1.812 J	4.177 J	0.8–1.7 $\times$	0.3–0.7 $\times$

Appendix

Table 11: Estimated orthogonalization energy per step for the GPT-2 124M workload. The 10k–40k budgets are measured on width-128 matrices; 60k–240k are rank-scaled projections. Pitch and stress denote the lower and higher device energy anchors.

Interface	Quantity	Value
Probe form vs. dense flow	code path with identity probes	bitwise equal
Quantized vs. continuous block	cosine of NS5 outputs	1.0000
Circuit vs. probe form	cosine of converged states	0.999983
	relative Frobenius distance	0.0059
	maximum element-wise distance	0.0045
Circuit vs. NS5	screen value, clean	0.9627 (behavioral 0.9624)

Appendix

Table 12: Agreement across the simulation ladder

Figure 5: Circuit transient on the $8\times 8$ block. Left: the 64 capacitor voltages against the probe-form state (cosine 0.999983). Middle: screen values of the circuit output against NS5, clean and with non-idealities. Right: convergence over the 600-period transient.

Panel of Figure 6 — Figure 6: Left: device law $g_{\mathrm{eff}}=2KV_{X}$ of the read element (200.00 $\mu$ S/V, deviation below 0.01%). Right: reciprocity of the two- and four-transistor cells and the $2\times 2$ array cosines for $q$ , $r$ , and $\Delta X$ .

Physical Neural Networks

See all 10 papers

Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates
Dec 4, 2025Thibaut Boissin, Thomas Massena, Franck Mamalet +1Newton-Schulz Iteration Orthogonal Optimization
IRIT-MISFIT · Institut de Recherche Technologique Saint-Exupery, France · DTIPG - SNCF, IRIT-MISFIT +2
Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training
Mar 18, 2026Ben S. Southworth, Stephen ThomasDeep Learning Optimization Muon Optimizer
Theoretical Division, Los Alamos National Laboratory, USA. · Lehigh University, Computer Science
AMO: Operator-level Adaptive Muon Orthogonalization
May 18, 2026Xinlin Zhuang, Panyi Ouyang, Yichen Li +7Language Model Pretraining Newton-Schulz Iteration
The Chinese University of Hong Kong · Shopee · MBZUAI +3

Dec 4, 2025J/K move · Enter open · S save

Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates

Thibaut Boissin, Thomas Massena, Franck Mamalet, Mathieu Serrurier

IRIT-MISFIT · Institut de Recherche Technologique Saint-Exupery, France · DTIPG - SNCF, IRIT-MISFIT+2

Newton-Schulz Iteration Orthogonal Optimization Neural Network Optimization Muon Optimizer Second-Order Optimization

Open paper AskPDF arXiv

Orthogonality-based optimizers, such as Muon, have recently shown strong performance across large-scale training and community-driven efficiency challenges. However, these methods rely on a costly gradient orthogonalization step. Even efficient iterative approximations such as Newton-Schulz remain expensive, typically requiring dozens of matrix multiplications to converge. We introduce a pre-conditioning procedure that improves the initialization of the Newton--Schulz iterations while incurring negligible overhead. Furthermore, our pre-conditioning reduces the initial polar error and enables the removal of one Newton-Schulz iteration (out of the five iterations usually used in practice). The resulting implementation significantly reduces Muon's overhead. At the end-to-end training level, we observe consistent runtime improvements across speed-run and standard benchmarks, including

\sim

3% reductions in training time on multiple fast training benchmarks, while matching reference performance on both language and vision tasks. Crucially, these improvements require no hyperparameter tuning and can be adopted as a simple drop-in replacement. Beyond empirical gains, we provide theoretical insight into the geometry of the update and its potential robustness against feature collapse. Our code is publicly available on github, in optax and huggingface kernels.

Abstract

Physical Neural Networks

Explore similar work

Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates

Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training

AMO: Operator-level Adaptive Muon Orthogonalization