Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data

We sincerely thank you for your positive assessment of our work. Here, we address your comments/questions in sequence:

Linear setting We provide a first complete characterization of the implicit bias of NSD/NMD algorithms on multiclass data, and we believe the insights from our work can be helpful for future studies on the non-linear models.

Feature-norm upper bound You are right: Here, we introduce $B$ as the upper bound on the maximum $L_1$-norm of the features. For finite dataset in a finite-dimensional space, such a $B$ always exists. We choose to state this as an assumption for clarity (particularly since $B$ appears throughout the analysis). We will add a remark about the existence of $B$ as mentioned above.

Constant step-size For GD[1,2] and NGD[3], it is indeed true that the implicit bias can be shown with constant step-size for binary data, utilizing the underlying Euclidean geometry. However, the class of algorithms (including Schatten norms) that we consider is much wider and can have different geometries. It is challenging to unify their analyses under the same framework for multiclass data. We manage to do this by constructing a proxy function (that works for all p-norms) to bound both the first and second-order terms in the Taylor expansion of the loss. This results in a bound on the margin gap with explicit dependence on the learning rates as shown in Theorem 3 in Appendix. From there, we observe that the step size needs to satisfy Assumption 2 to ensure convergence. It remains interesting and (potentially) challenging to show the implicit bias of NSD/NMD with constant step-size for the multiclass setting and for the range of p-norms that we consider (we will add this in the future work discussion). Note also that decreasing step-size has been used in other implicit bias works [4,5].

Additional comments Thank you for pointing out these typos. We will fix them in the final version of the paper.

[1] The Implicit Bias of Gradient Descent on Separable Data, Soudry et al

[2] The implicit bias of gradient descent on nonseparable data, Ji and Telgarsky

[3] Characterizing the implicit bias via a primal-dual analysis, Ji and Telgarsky

[4] Convergence of Gradient Descent on Separable Data, Nacson et al

[5] The Implicit Bias of Adam on Separable Data, Zhang et al

Thank you again for your time and suggestions. Please let us know if there are any further questions.

We sincerely thank you for your positive assessment of our work. Here, we address your comments/questions in sequence:

Comparison to previous works We have provided discussions on the convergence rates of Adam (and comparisons to Zhang et al.) in Appendix G, see Remark 3. Here, we reemphasize the key differences on the results/scope of our work and others. The work by Zhang et al. studied Adam in the binary setting [1], instead we study a broad family of NSD/NMD algorithms in the multiclass setting. As a byproduct of our NMD analysis, we get a rate for Adam in the multiclass setting that matches the rate of Zhang et al. In fact, direct use of their result would introduce an extra factor of $k$ (number of classes) into the bound, we are able to avoid this by the tight decomposition of the proxy function (below Lemma 6). Regarding Xie and Li [2], their non-linear setting and asymptotic convergence (to a KKT point of a constrained optimization problem) differ from our setting and the non-asymptotic rate.

Proxy function In the binary case, the proxy function is defined using gradients. As loss decreases, the proxy function converges to $0$ as gradients become small. In the multiclass setting, the logit corresponding the true class saturates as loss decreases. Hence, the proxy function for the multiclass setting also converges to $0$.

Parameter convergence Parameter convergence (and its rate) is more general to show and we can’t conclude formally from our proof although the simulation seems to suggest that this is the case. Note that the more subtle issue is that for certain norms of interest (e.g., max-norm (SignGD) and spectral-norm (specGD/Muon)), the max-margin problem does not necessarily admit a unique minimizer, which further complicates parameter convergence. Moreover, the parameter convergence of NGD (i.e. NSD w.r.t 2-norm) utilizes the properties of Euclidean geometry [3]. Here, the algorithms (and their associated norms) that we consider are much wider and beyond NGD.

Additional comments Thank you for your suggestions regarding line 158 and 211. We will revise them in the final version of the paper.

[1] The Implicit Bias of Adam on Separable Data, Zhang et al

[2] Implicit Bias of AdamW: $L_\infty$ Norm Constrained Optimization, Xie and Li

[3] Characterizing the implicit bias via a primal-dual analysis, Ji and Telgarsky

Thank you again for your time and suggestions. Please let us know if there are any further questions.

We sincerely thank you for your positive assessment of our work. Here, we address your comments/questions in sequence:

Effects of momentum For the setting considered in this work, momentum does not change the implicit bias in terms of margin convergence w.r.t a specific norm. Indeed, Figure 1,2 reveal that both spectral-GD and Muon favor the spectral-norm margin over the others. We are not making any claims regarding the non-linear setting. One thing to note is that preconditioning may also play a role in the implicit-bias difference between Adam and GD.

Normalization In Thm 1, the norm used for normalizing the weights is the one used to define the algorithm (eqn. (3)), and the margin (denoted as $\gamma$) is defined w.r.t. the same norm (eqn. (2)). For example, the margin of (normalized) iterates of SignGD (i.e. NSD w.r.t max-norm) converge to the data margin defined w.r.t the max-norm according to Thm 1. We want to emphasize that normalization in Thm 1 has nothing to do with modifying the algorithm. The algorithm runs as if it has no norm constraints. But because the norm of the parameters diverges in training, normalization is necessary to show convergence. In experiments, we compute the margin of normalized iterates (denoted as $\gamma_{\bar{W}_t}$) where the norm used in normalization and defining the data margin are the same. For example, the iterates of SignGD can be normalized w.r.t max-norm, Frobenius-norm, or spectral-norm. Regardless of which norm is used, same norm needs to be used to compute the data margin when plotting their differences. We are doing this for the purpose of demonstrating norm-preferences of a chosen NSD/NMD algorithm.

Stability Constant The stability constant ($\epsilon$) can play a role in the implicit bias of Adam. Adam can favor a $L_{2}$-norm margin with a non-negligible $\epsilon$ [1]. However, this happens when the values of gradients are below $\epsilon$ (which is often chosen to be very small such as $10^{-8}$). In practice, training is typically not long enough for this realization to occur. Thus, Adam with negligible $\epsilon$ (i.e. max-norm solution favorance) is still practically relevant. We have provided a detailed discussion on the effect of $\epsilon$ in Appendix G (see pg 39,40). Our experiments in Figure 5 (in Appendix) demonstrate the transition from max-norm to $L_{2}$-norm margin favorance (after extended training), and hence consolidate that the findings of [1] and [2] also hold in the multi-class setting.

Non-linear settings Despite the analysis is done in the linear setting, our work is the first to completely characterize the implicit bias of NSD/NMD algorithms for multiclass data. The extension to the non-linear settings is challenging and we believe that our theoretical techniques/insights can be helpful for such future studies. To motivate such theory extensions, we have performed (and will include in the final version of the paper) the following non-linear experiment: train a 2-layer MLP (with Relu activation denoted as $\sigma(\cdot)$ and hidden dimension being 100) on MNIST dataset (100 data points sampled from each class) and track the spectral-norm margin of the model defined as $\min_{i \in [n], c \neq y_i} \frac{(e_{y_i} - e_{c})^T W \sigma(H x_i)}{||| H |||_2}$ (note that $H$ and $W$ denote the weights of the first and second-layer respectively, and we only train the first layer). We observe that the spectral-norm margin of spectral-GD grows the fastest among the algorithms considered, suggesting that the implicit bias to respective norm (demonstrated with spectral-norm) is preserved in these non-linear settings.

Table: Spetral-norm margin of SignGD, NGD, and Spectral-GD at different iterations

Non-separable dataset For the non-separable data setting (e.g. [3]), we would expect margin convergence (w.r.t a specific norm) to hold in the separable subspace. This type of implicit bias has been shown for GD (as a follow-up to implicit bias on the separable binary data) [3]. Our work makes it possible to extend the analysis to NSD on multiclass data, which is an interesting future direction. We will add this in the future work discussion.

Additional comments Thank you for these comments. We will fix them in the final version of the paper.

[1] Does Momentum Change the Implicit Regularization on Separable Data, Wang et al

[2] The Implicit Bias of Adam on Separable Data, Zhang et al.

[3] The implicit bias of gradient descent on nonseparable data, Ji and Telgarsky

Thank you again for your time and suggestions. Please let us know if there are any further questions.

We sincerely thank you for your positive assessment of our work. Here, we address your comments/questions in sequence:

Other Losses The convergence rates also hold for other losses. We have discussed these extensions in Appendix F.

Dependence of p The margin $\gamma$ naturally depends on p via its definition in eqn. (2), and the convergence rates in Thm 1 and 2 hold for any p. Our current analysis does not explicitly track how the constant depends on p. However, inspecting the analysis steps in Sec. 3 (i.e. lower and upper bounding the first and second-order Taylor terms respectively) reveal some insights. First, for all values of p, the dual-norm of the gradients is lower bounded by the same proxy function multiplied by the margin $\gamma$. The proxy function itself has no p-dependence. When it is used to upper bound the second-order term (see the display after Line 198), the bound depends on the norm of the features. Currently, we are using the upper bound on $L_1$-norm of the features, denoted as $B$ (i.e. $max\_{i \in [n]} \lVert h_i \rVert_1 \leq B, \forall i \in [n]$). The constant $B^2$ can be substituted by $max\_{i \in [n]} |||h_i h_i^T|||_{p}$ (note that this reduces to $max\_{i \in [n]} \lVert h_i \rVert_p^2$ for entry-wise norms, see lines 195-197), which captures the geometry of the features. We will include this comment in the final version of the paper.

Intuition on the choice of norms This is a good question and in-part is one of the motivations behind implicit-bias analyses like our paper. Specifically, having shown that different norms in NSD/NMD maximize margins of the dataset w.r.t corresponding norms, facilitates the generalization question: which p-norm generalizes better? Of course, the answer depends on the data in complicated ways. In the linear case, this can be thought as follows: p-norm NSD favors convergence to a classifier that maximizes p-norm margin. Let the underlying data distribution be parameterized by a matrix which can represent the parameters of multinomial logistic model or the mean-vectors of multiclass Gaussian mixture (as an example). If the matrix that generates the data is sparse, SignGD is preferred; if it has a balanced spectrum, Spectral-GD is preferred; if it is low-rank, low-rank GD is preferred and so on. While we recognize the limitations of this setting, it can provide some initial intuitions towards the choices of algorithms in different cases.

Experiments (real datasets)/theory on non-linear models We performed additional experiments on MNIST dataset (100 data points sampled from each of the 10 classes) in the non-linear multi-class setting. We use a two-layer MLP with random initializations and the hidden dimension being 100. We only train the first layer and track the spectral-norm margin (as an example) of the overall network defined as follows: $\min_{i \in [n], c \neq y_i} \frac{(e_{y_i} - e_{c})^T W \sigma(H x_i)}{||| H |||_2}$, where we denote the first and second-layer weights as $H$ and $W$ respectively, and $\sigma(\cdot)$ is the Relu activation function. We observe that the spectral-norm margin of Spectral-GD grows the fastest among the three algorithms considered (SignGD, NGD, and Spectral-GD). This confirms that the norm preferences of NSD algorithms (exhibited in the linear settings) can also show up in the non-linear settings. We will incorporate these experiments in the final version of the paper. From theoretical perspectives, despite being in the linear setting, we are the first to formally prove the implicit bias of NSD/NMD for the multiclass data. The extension to the non-linear (multiclass) settings is challenging and we believe our work provides the ground for such future studies.

Table: Spetral-norm margin of SignGD, NGD, and Spectral-GD at different iterations

Thank you again for your time and suggestions. Please let us know if there are any further questions.