Physical muon enables energy efficient transformer training with orthogonalization flow
Physical Muon: Orthogonalization as an Equilibrium Computation
Machine LearningEmerging Technologies
Summary
Training large AI models like transformers can use a lot of energy, especially for certain math steps called orthogonalization. The authors looked at a new method called Physical Muon that makes these steps work using simpler, energy-saving operations. They tested it on a real transformer model and found it performs almost as well as the standard method. This approach could help future AI training run more efficiently on specialized hardware.
What this means in practice
- •For machine learning engineers: Run transformer training algorithms on analog hardware with energy-efficient orthogonalization steps using continuous-time flow computation.
- •For hardware developers: Design analog accelerator circuits that implement equilibrium-based orthogonalization to improve transformer training efficiency.
Authors
Yuren Hao
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.