The Implicit Bias of Adam and Muon on Smooth Homogeneous Neural Networks

We study the implicit bias of momentum-based optimizers on smooth homogeneous models. We show that *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 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.

"Generalization" is the measure of how well machine learning models perform on data they did not see during training. A prevalent problem in machine learning is "overfitting" - when models adapt too well to their training data and generalize poorly, like a student memorizing solutions to problems instead of understanding the underlying concepts. One source of good generalization is for the model to have a large "confidence margin" for its decisions: even if the input changes a bit, the decision of the model doesn't flip. Often, to ensure that models generalize well and don't overfit, it is necessary to incorporate smart training strategies like penalizing the model for being too sophisticated. However, surprisingly, sometimes models can learn to generalize even without such explicit penalizing strategies. Researchers noticed that this might be related to the training algorithm "gradient descent" and its variants; this line of research is called "implicit bias". We analyze popular variants of gradient descent called Adam and Muon, on a specific set of models called smooth homogeneous models. We prove mathematically and show experimentally that training with Adam, Muon, and other algorithms encourages the model to generalize well by increasing a sort of margin for its decisions. Also, interestingly, the characteristics of this margin depend on the algorithm used.