Search papers, labs, and topics across Lattice.
This paper analyzes the implicit bias of momentum-based optimizers, including Muon, MomentumGD, Signum, and Adam, on smooth homogeneous neural networks. By extending existing results on steepest descent, the authors demonstrate that these momentum-based methods approximate steepest descent trajectories under a decaying learning rate schedule. Consequently, they prove that these algorithms exhibit a bias towards KKT points of a margin maximization problem, with the specific margin depending on the optimizer's norm.
Adam and Muon secretly optimize for different margin norms when training homogeneous neural networks, revealing a hidden bias in these popular optimizers.
We study the implicit bias of momentum-based optimizers on homogeneous models. We first extend existing results on the implicit bias of steepest descent in homogeneous models to normalized steepest descent with an optional learning rate schedule. We then show that for smooth homogeneous models, momentum steepest descent algorithms like Muon (spectral norm), MomentumGD ($\ell_2$ norm), and Signum ($\ell_\infty$ norm) are approximate steepest descent trajectories under a decaying learning rate schedule, proving that these algorithms too have a bias towards KKT points of the corresponding margin maximization problem. We extend the analysis to Adam (without the stability constant), which maximizes the $\ell_\infty$ margin, and to Muon-Signum and Muon-Adam, which maximize a hybrid norm. Our experiments corroborate the theory and show that the identity of the margin maximized depends on the choice of optimizer. Overall, our results extend earlier lines of work on steepest descent in homogeneous models and momentum-based optimizers in linear models.
