Physical Muon: Orthogonalization as an Equilibrium Computation
cs.LG, cs.ET, stat.ML
Submitted: 2026-09-29
Updated: 2026-09-29
Comments: Accepted at 18th annual workshop on Optimization for Machine Learning
License: http://creativecommons.org/licenses/by/4.0/
The gist: Physical neural networks and analog in-memory computing could reduce the energy cost of neural network training.
Terminology
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.
Sources
- Dion: Distributed Orthonormalized Updates
- The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm
- The Spectral Dynamics and Noise Geometry of Muon
- Old Optimizer, New Norm: An Anthology
- Symbolic Discovery of Optimization Algorithms
- CacheMuon: Using Temporal Preconditioning To Approximate Polar Factor
- Understanding and Embracing Imperfection in Physical Learning Networks
- Adam: A Method for Stochastic Optimization
- Cautious Optimizers: Improving Training with One Line of Code
- Muon is Scalable for LLM Training
- Decoupled Weight Decay Regularization
- Move on Muon : A Hamiltonian probability gradient flow perspective of Muon optimizer
- The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale
- Muon Dynamics as a Spectral Wasserstein Flow
- Adafactor: Adaptive Learning Rates with Sublinear Memory Cost
- Generative Models on Analog Hardware with Dynamics
- OLion: Approaching the Hadamard Ideal by Intersecting Spectral and $\ell_{\infty}$ Implicit Biases
- Why Transformers Need Adam: A Hessian Perspective
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks