Muown Implicitly Performs Angular Step-size Decay

2026-06-22

AI Summary

This work addresses the instability in training Transformers caused by the coupling of magnitude and direction in conventional optimizers, which leads to erratic step sizes and poor control over angular updates. To resolve this, the authors propose AngularMuON, an algorithm that decomposes weight matrices into row-wise magnitudes and unnormalized directions, optimizing them separately with Adam and Muon, respectively. The method reveals, for the first time, an implicit angular step-size decay mechanism inherent in Muon and is shown to be equivalent to Riemannian gradient updates on the unit sphere, enabling explicit scheduling of angular steps. Empirical results demonstrate that AngularMuON outperforms MuON, achieves state-of-the-art speed in the nanoGPT benchmark, and exhibits strong scalability on Qwen2 models of 0.5B and 1.1B parameters, including MoE variants.

Abstract

Matrix-aware optimizers such as Muon and Muown have recently shown strong empirical performance for pre-training Transformers. In particular, Muown separates each weight matrix into row magnitudes and an un-normalized direction variable, updating the former with Adam and the latter with Muon. We show that the directional update of Muown is equivalent to a Riemannian step on the normalized directions, while the magnitude of the un-normalized parameterization only modulates the angular step size. This explains the step-size stability of Muown and suggests making the angular step size explicit. The resulting method, AngularMuown, optimizes directly over the normalized directions and uses a schedulable angular multiplier decoupled from the radial magnitude update. AngularMuown improves over Muown and, at the time of writing, a preliminary version is leading the per-optimizer category of the modded nanoGPT speedrunning competition. Further experiments on Qwen2-0.5B, and 1.1B parameter mixture-of-experts models confirm the algorithm scales beyond small models. An implementation of the algorithm is available at https://github.com/fhueb/angular-muown

Problem

• angular step-size
• matrix-aware optimizer
• directional update
• Riemannian optimization
• step-size stability

Innovation

• Angular step-size decay
• Riemannian optimization
• Matrix-aware optimizer
• Direction-magnitude decomposition
• AngularMuown