Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data

TL;DR

This paper develops a unified framework to characterize the implicit bias of NSD and NMD (including Spectral-GD, Muon, and Adam) in multiclass linear classification with cross-entropy loss. It proves that iterates converge to max-margin solutions with respect to the chosen norm (entrywise or Schatten), and provides non-asymptotic rates for margin convergence that cover both static and momentum-based updates. The framework hinges on a proxy function G(W) that connects gradient magnitude and loss progression across norms, enabling a seamless translation of results between max-norm, l2-norm, and spectral-norm geometries. Empirically, the authors demonstrate norm-specific margin preferences on synthetic data and a two-layer neural network, supporting the theory and suggesting that linear-margin insights extend to simple nonlinear regimes. This work advances understanding of how optimizer geometry shapes generalization in multiclass settings and offers a pathway to analyzing broader optimization algorithms in deep learning contexts.

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.

Figures (6)

• Figure 1: (a) We normalize the iterates of SignGD w.r.t. the max-norm (denoted as $\bar{{\bm{W}}}_t$), compute the margin (denoted as $\gamma_{\bar{{\bm{W}}}_t}$), then plot its difference to data margins $\gamma_{\lVert \cdot \rVert_{\infty}}$, $\gamma_{\lVert \cdot \rVert_{2}}$, and $\gamma_{\left|\left|\left|\cdot\right|\right|\right|_{\infty}}$ (note that the margin difference is further divided by the corresponding data margin for comparisons). SignGD favors the margin defined w.r.t. the max-norm. (b, c, and d) Same as (a) with SignGD (max-norm) replaced by NGD (2-norm), Spectral-GD (spectral-norm), and Muon (spectral-norm), respectively. NGD favors the 2-norm margin, while Spectral-GD and Muon favor the spectral-norm margin.
• Figure 2: (a) Correlations between the iterates of SignGD (${\bm{W}}_t$) and max margin separators ${\bm{V}}_{\infty}$, ${\bm{V}}_{2}$, and ${\bm{V}}_{\mathop{\mathrm{spec}}\nolimits}$ against iterations (correlation defined as $\frac{\langle {\bm{W}}, {\bm{V}} \rangle}{\lVert {\bm{W}} \rVert_2 \lVert {\bm{V}} \rVert_2}$). (b, c, and d) Same as (a) with SignGD replaced by NGD, Spectral-GD, and Muon, respectively. SignGD and NGD correlate well with ${\bm{V}}_{\infty}$ and ${\bm{V}}_{2}$, respectively, while Spectral-GD and Muon correlate well with ${\bm{V}}_{\mathop{\mathrm{spec}}\nolimits}$.
• Figure 3: (a) Spectral-norm margin $\gamma_a^{{\bm{V}}_t}$ as a function of $t$ for SignGD, NGD, and Spectral-GD. (b) Same as (a) with algorithms replaced by Signum, NMD-GD, and Muon, respectively. (c) Spectral-norm margin $\gamma_b^{{\bm{V}}_t, {\bm{W}}_t}$ as a function of $t$ for SignGD, NGD, and Spectral-GD. (d) Same as (c) with algorithms replaced by Signum, NMD-GD, and Muon, respectively.
• Figure 4: Implicit bias of Signum and NMD-GD on multiclass separable data. (a) Relative margin gap of Signum's iterates against iterations. (b) Correlation of Signum's iterates to ${\bm{V}}_{\infty}$, ${\bm{V}}_{2}$, and ${\bm{V}}_{\mathop{\mathrm{spec}}\nolimits}$ against iterations. See Figure \ref{['fig:main_margin']} and \ref{['fig:main_cor']} for the definitions of relative margin and correlation. (c) and (d) Same as (a) and (b) with Signum replaced by NMD-GD.
• Figure 5: Implicit bias of NGD, SignGD, and Adam on multiclass separable data ($k=5$,$d=25$, and 50 data points in each class). (a,b) Loss and gradient 2-norm vs. iterations: SignGD converges faster than others. (c) We normalize the iterates w.r.t. 2-norm (aka Frobenius), compute the margin, then plot its difference to the dataset's max-margin w.r.t. 2-norm (given in captions). Only NGD converges to the max 2-norm margin. (d) Same as (c) with 2-norm replaced by max-norm. Margins of SignGD/Adam(with $\epsilon=0$) converge to max-margin w.r.t max-norm. For SignGD, the training is stopped after $10^{4}$ iterations due to the numerical instabilities caused by the small gradient norm.

Theorems & Definitions (85)

• Lemma 1: Lower bounding the gradient dual-norm
• Lemma 2
• proof
• Lemma 3: Properties of ${\mathcal{G}}({\bm{W}})$ and $\mathcal{L}({\bm{W}})$
• Lemma 4: Ratio of ${\mathcal{G}}({\bm{W}})$
• Lemma 5: NSD Unnormalized Margin
• Theorem 1
• Remark 1
• Lemma 6
• Lemma 7