Muon optimizer tweak with a polynomial grid saves ~0.5% of LLM training compute

burny_tech · x · 2026-10-20

Developer Vuk Rosic shares a Muon optimizer improvement that saves up to 0.5% of LLM training compute:

Problem: standard Muon always uses 5 rounds of a degree-5 polynomial (Newton-Schulz / Polar Express) to equalize gradient direction strengths, but one fixed formula isn't optimal everywhere.
Method: build a grid over three parameters — weakest-direction share (0.0001–0.01), tolerance to 1 (±0.02–±0.4), matrix shape (square vs wide 1:4) — and for each setting try every mix of degree 3/5/9/13/17 rounds (1,104 sequences per row), keeping the cheapest that hits the target.
Results: e.g. 9+9+9 hits a common target with 33% fewer FLOPs than the standard 5+5+5+5+5; since Muon is 0.5–2% of training, that's 0.1–0.5% overall savings.
Reusable: open-sourced as muon-formula-grid (grid.json, coefficients, a picker page); he suggests having your coding agent measure your model's direction distribution, then pick a formula from the grid or rerun the search.

More from Infra