Muon variant cuts diffusion transformer training time 24%

A periodic Muon variant cuts diffusion transformer training time up to 24%, per a paper on periodic row-wise Muon.

The tradeoff matters to anyone training diffusion transformers past a few billion parameters. Muon already beats AdamW on generative quality by 12.9% to 19.1%, and that edge holds from 1.3 billion to 15 billion parameters. But it recomputes a full spectral correction on every step, a five-step Newton-Schulz iteration plus full-momentum storage. That overhead grows enough at scale to cut into the value of the gain.

The row-wise substitute swaps in a cheaper constrained update most steps, saving the full computation for once every K steps. Quality held within 0.5% of standard Muon at 1.3 billion to 4 billion parameters, and improved 4.5% at 9 billion.

Step time fell 15.7% to 24.3% end to end. Communication volume fell 66.7%, since the row-wise update moves far less data between steps than the full spectral computation.