Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data

TL;DR

This paper characterizes the implicit bias of spectral descent and Muon algorithms in multi-class linear classification, showing they converge to max-margin solutions with respect to various matrix norms, including spectral and max-norm.

Contribution

It provides the first complete theoretical analysis of the implicit bias of p-norm normalized steepest descent and momentum methods in multi-class classification, unifying several special cases.

Findings

Algorithms converge to max-margin solutions with respect to the p-norms.

Spectral Descent and Muon specifically converge to spectral norm max-margin solutions.

Adam converges to the max-norm solution with preconditioning.

Abstract

Different gradient-based methods for optimizing overparameterized models can all achieve zero training error yet converge to distinctly different solutions inducing different generalization properties. We provide the first complete characterization of implicit optimization bias for p-norm normalized steepest descent (NSD) and momentum steepest descent (NMD) algorithms in multi-class linear classification with cross-entropy loss. Our key theoretical contribution is proving that these algorithms converge to solutions maximizing the margin with respect to the classifier matrix's p-norm, with established convergence rates. These results encompass important special cases including Spectral Descent and Muon, which we show converge to max-margin solutions with respect to the spectral norm. A key insight of our contribution is that the analysis of general entry-wise and Schatten p-norms can be reduced to the analysis of NSD/NMD with max-norm by exploiting a natural ordering property between all p-norms relative to the max-norm and its dual sum-norm. For the specific case of descent with respect to the max-norm, we further extend our analysis to include preconditioning, showing that Adam converges to the matrix's max-norm solution. Our results demonstrate that the multi-class linear setting, which is inherently richer than the binary counterpart, provides the most transparent framework for studying implicit biases of matrix-parameter optimization algorithms.