Paper: High-dimensional SGD aligns with emerging outlier eigenspaces (ICLR-2024)

Reviewer comment (Theory):
In the equation above (C.14) (the one that starts with \(\nabla_c H \otimes \nabla_d H\)), I think many terms are missing in the RHS. Please check that part carefully.

Author response:
*"I checked the proofs corresponding..."*: We are grateful to the referee for pointing these out, especially the missing terms in (C.14) and its effect on Theorem C.7. We have fixed all of these in the supplementary material (they do not any of our affect our main results).

Reviewer comment (Theory):
Because of the previous point, Eq. (C.14) should be incorrect, which eventually makes the correctness of the correction terms in Theorem 5.7 questionable. Can you check this part again as well?

Author response:
*"I checked the proofs corresponding..."*: We are grateful to the referee for pointing these out, especially the missing terms in (C.14) and its effect on Theorem C.7. We have fixed all of these in the supplementary material (they do not any of our affect our main results).

Reviewer comment (Theory):
Your results crucially depend on the high-dimensional limit for SGD proven in [1]. In that work, it was shown that there is a critical learning rate that causes additional correction terms in the limit. Can you comment on how the existence of the correction terms affects your results?

Author response:
*"As the last question for the authors..."*: This is a great question; the correction term at critical scaling also scales as $1/\lambda$, i.e., the variance of the noise. Thus, the effect of the critical scaling is getting absorbed as one of the sources of the $O(1/\lambda)$ errors that we sustain in our main results. It would require a more refined analysis of these errors to understand better the role of the critical scaling.

Reviewer comment (Theory):
The main limitation is that the two classification tasks have a particular structure. Indeed, in both the linear case and the two-layer case considered here, the Hessian/G-matrix of the first layer weights is proportional to $Y\otimes Y$. For a Gaussian mixture $Y = y\mu + Z_{\lambda}$, indeed the relevant subspace spanned by the means will naturally pop as an outlier of this matrix when the SNR $\lambda$ is large enough: this is the classical BBP transition. Therefore, it is not surprising this also happens at the level of the Hessian/G-matrix. Maybe proving this requires a tour de force, but it is fair to doubt on the generality of the conclusion. For instance, it would be desirable to show a similar result for tasks where the structure is on the labels, and not on the inputs, e.g. teacher-student models.

Author response:
We thank the referee for their careful reading and helpful comments. We address these below. 

*Weaknesses*: We agree with the referee that, in the case of multiclass logistic regression and the first layer of the XOR model, the results we present are natural. That said, we would like to emphasize, here, the following subtle points: 
1. The spectrum of these matrices are spatially dependent and, in particular, their outlier eigenvectors, their spectra and even the “critical” lambda is itself vary in space. Note that even for the first layer, the spatial dependence of the coefficients of $Y\otimes Y$ can affect not only the eigenvalues but also the eigenvectors. 
2. For the second layer, the matrix is not proportional to the data matrices: It is a spatially (and thus time) dependent, non-linear transformation of the data. Moreover the second layer is O(1) and thus distinct from the regime of the classical BBP transition. 
3. It is not crucial to our analysis that the structure is in the feature data. In particular, similar arguments should extend to single index models/teacher-student networks.

Reviewer comment (Theory):
In Thm 3.1, how does the critical variance $\lambda_{0}$ for SGD alignement compares with the BBP threshold for the recovery of the means $\mu_{a}$ from $Y = y\mu + Z_{\lambda}$?

Author response:
*[Q1]*: We work here in the regime where $\lambda_0$ is large enough (though $O(1)$); it is a very interesting question to understand the critical $\lambda$ for the dynamical BBP transition.

Reviewer comment (Theory):
In both examples considered here, the timescale for the trajectory to correlate with the outlier subspace is linear in the dimension. How general should we expect this to be? For instance, would the authors expect this to hold in problems where escaping a fixed point would take longer than linear time in the dimension?

Author response:
*[Q2]*: We expect similar behavior to hold for problems where it takes more than linear time to escape a fixed point, except that the alignment will itself take a diverging time to appear. A good example to have in mind for this is a problem like tensor PCA or certain single-index models, and is an interesting avenue of future investigation.

Reviewer comment (Theory):
On the same spirit of the question above, would the authors expect the theorem to remain true at longer timescales? In general, should we expect these subspaces to be stable or the SGD trajectories to eventually escape these subspaces?

Author response:
*[Q3]*: The alignment should hold for much longer timescales, and depending on how strong alignment one desires, that can be proven in some examples. In the XOR problem with multi-layers, one must be careful since there can be cascades of transits from one fixed point to another, and during those transits the alignment could become lost as it looks for the next best fixed point, then regained once it settles there.

Reviewer comment (Theory):
Are the authors aware of converse examples to their result. For instance, a problem where the underlying target function does depend only on a few relevant directions, but that the outliers in the Hessian/G-matrix are not necessarily aligned with it?

Author response:
*[Q5]*: We thank the referee for this question. In the XOR model, while the population loss has only four relevant directions, one for each hidden class, we show that there are regimes of initialization where the spectrum that emerges misses some of these directions. Please see Figure 4.3 and the related discussion.

Reviewer comment (Writing):
The reference in the top of Page 2 (“Liao & Mahoney (2021) studied Hessians of some non-linear models beyond generalized linear models, also at initialization.”) is misleading. The main result in Liao & Mahoney (2021) concerns the Hessian of a loss function of the type (see eq. 3): $$L(w)=\frac{1}{n}\sum\limits_{i=1}^{n}\ell(y_{i}, w^{\top}x_{i})$$ under the assumption $y_{i}\sim f(y|w_{\star}^{\top}x_{i})$ (see eq. 1). Despite their unusual choice of terminology for this model ("G-GLM"), this is mostly commonly known as a generalized linear model in the literature... when just not simply refereed to as "linear model".

Author response:
*"Reference in top of page 2"*: We reworded this to match the terminology from Liao & Mahoney (2021).

Reviewer comment (Novelty):
The following references — “ODE limits of the SGD training with single layer networks have been derived and numerically solved in Mignacco et al. (2020); Loureiro et al. (2021).” — are not accurate. Mignacco et al. (2020) and Loureiro et al. (2021) study the classification of Gaussian mixtures with single layer networks. However, they derive exact expressions for the training and test errors of the minimiser of the empirical risk. Indeed, this should correspond to the asymptotic performance of one-pass SGD when $t\to\infty$ of the example considered in Section 3.

Author response:
*"The following references..."*: Thank you for pointing this out, we have corrected this.

Reviewer comment (Presentation):
The figures in the manuscript are not very readable. First, even if this is described in the caption, it would be good to label the (x,y)-axis for clarity. Second, they are small, and since the format is not vectorial they get pixelated when zoomed over. Third, the meaning of the colour scheme is unclear — maybe add a colour bar.

Author response:
*"The figures are not very readable..."*: We have added to the captions descriptions of what the different colors are denoting in each figure. Due to space constraints, adding legends/axis labels becomes even harder to parse compared with having them in the caption, but we will try to vectorize the figures for the final version so that they are more clear when zooming.

Reviewer comment (Novelty):
While the authors might not want to dwell on 30 years of literature on SGD scaling limits, it is recommended to mention at least some of the recent works contemporary to [Ben Arous et al. 2022] to represent the diversity of this literature. An example list is provided in the review.

Author response:
*"30 years of literature on SGD scaling limits"*: We have included the references as requested, except for Chen et al (2019) which appears to be less relevant as it is focused on sample complexity bounds for gradient descent.

Reviewer comment (Novelty):
Comparison and contribution over the prior works are not very clear.

Author response:
We thank the referee for their careful reading and helpful comments. 

*Weaknesses*: We have added further references to previous work as per this and other referees' requests to better contextualize our work.

Reviewer comment (Theory):
Is it possible to extend the theory based on cross-entropy loss to general function class?

Author response:
*Question 1*: Our focus in this work was to build a theoretical foundation for the experimental work, which was largely in the context of classification tasks. We expect that many of the phenomena we find extend to more general settings, and these are certainly interesting avenues for future work.

Reviewer comment (Theory):
The authors mention that SGD finds the subspace generated by the outlier eigenvalues for any uninformative initializations with norm $O(1)$ in Section 3.2, does the similar property hold for XOR-type mixture models via two-layer networks?

Author response:
*Question 2*: Indeed, Gaussianity of the initialization is similarly not essential to the XOR theorems.

Reviewer comment (Theory):
The results are proved for two somewhat restricted settings, and it seems not clear how to extend the argument to more general settings. It would be helpful if the authors can comment on the limitations of the current proof.

Author response:
*“Two rather restricted settings”*: We agree that these two settings are rather restricted, and it would indeed be of interest to extend them to general, multi-class, non-linearly classifiable mixture models. We expect the random matrix concentration to be amenable to this greater level of generality. The biggest hurdle to such extensions is analyzing the limiting dynamical system for the stochastic gradient descent. These can be incredibly complex, even in simple settings like XOR labelings.

Reviewer comment (Novelty):
It seems that the main technical tool is from Ben Arous et al. (2022). It would be helpful to add some discussion on the technical novelty.

Author response:
*“Main technical tool”*: Our analysis consists of two stages; understanding the SGD trajectory, and understanding the matrices’ spectra along those trajectories. For the first stage, the main technical tool is indeed from Ben Arous et al (2022) (plus new stability/perturbative analysis of the limits). But the second half of the paper involves a disjoint set of tools, mainly from random matrix theory, which we describe in page 9 (Outline and Ideas of Proof).

Reviewer comment (Theory):
The results are valid for online SGD. It is worth a comment whether this is necessary, as well as what happens for multi-pass SGD.

Author response:
*“What happens for multi-pass SGD”*: Other variants of SGD would indeed be very interesting to study, but the limit theory for the SGD is undeveloped in this setting.

Reviewer comment (Theory):
The learning rate $\delta = O(1/d)$ seems to be on a very small scale, especially for large $d$. Could the authors comment on this requirement for the learning rate?

Author response:
*Question 1*: The learning rate has to be that small as a function of the dimension for the SGD to remain stable enough to converge in the right observables. In particular, if the learning rate is larger than that, then the resulting stochastic terms have diverging variances. That is the sense in which $\Theta(1/d)$ is identified as a critical step-size scaling in Ben Arous et al (2022).

Reviewer comment (Theory):
The main theorems state that the alignment happens for $\ell \in [T_0\delta^{-1}, T_f \delta^{-1}]$. What's the corresponding generalization performance of the model on the test data?

Author response:
*Question 2*: Thank you for this important question. The SGD analysis lets us understand the generalization error over the course of training. In the $k$-GMM setting, the population loss approaches $0$ on the same timescale as alignment occurs. In the XOR setting, the answer is more complicated. As shown in Ben Arous et al (2022), on this timescale the SGD can converge to a classifier with poor generalization with positive probability over the initialization. In this case, the alignment still occurs, but we emphasize that when the latter happens, the corresponding principal eigenspace has dimension less than $4$: see Figure 4.3.

Reviewer comment (Evaluation):
The major weakness of this work is that the motivation of studying the evolution of the Hessian/Fisher information is not super clear. According to the introduction and the related works, from my perspective, the best thing we know about the relationship between the Hessian and the effectiveness of SGD is that "this common low-dimensional structure to the SGD and Hessian matrix may be key to many classification tasks in machine learning" (Gur-Ari et al., 2019). From my point of view, there is no solid evidence indicating that the main results of this work imply an advantage or a disadvantage of SGD. I will consider raising my score if this question is properly settled.

Author response:
*“motivation of the work is not clear...”*: The success of SGD for high-dimensional tasks is often attributed to its ability to find low-dimensional subspaces, which captures the underlying data structure. This has not been rigorously shown, but empirically long predicted, and linked to outliers in the eigenspectra of loss Hessians and FI matrices. The goal of our work was to give this a strong theoretical foundation, demonstrating a precise manner in which it holds true in high-dimensional Gaussian mixture classification tasks, and explore refined features of this low-dimensional alignment with more complex multi-layer networks. In these tasks, we also relate the outlier eigenspaces of the matrices with the relevant coordinates of the data proving how and when the SGD succeeds at the classification task in terms of when the outlier eigenspectrum of the Hessian being of the same rank as the number of hidden classes. 

Our goal was not to compare the performance of SGD to other optimization algorithms, though understanding their alignment phenomena would also be of significant interest. In particular, that may allow one to prove relative advantages/disadvantages of different learning algorithms.

Reviewer comment (Presentation):
Figure 3.3 and Figure 4.3 are a bit confusing without legends.

Author response:
*“Figures 3.3 and 3.4”*: We have modified the captions to indicate what the different colors indicate in the figures. We will also try to vectorize the figures for the final version to make them more clear when zoomed in.

Reviewer comment (Writing):
In the paragraph following Theorem 4.2, the terms "GOE" and "Wishart matrix" may appear unfamiliar to many.

Author response:
*“GOE and Wishart unfamiliar”*: We expanded the acronym GOE, and added a footnote with a reference to a classical random matrix theory textbook.

Reviewer comment (Theory):
Is there a counterpart of Theorem 3.3 for the XOR setting?

Author response:
*Question 2*. We are not sure if this addresses your question, but in the XOR setting, we have already specialized the means to be orthonormal. It would be interesting to study the XOR phenomenology with more than two classes, and when the means are not orthonormal. We are happy to clarify further if this doesn’t address the question.

Reviewer comment (Theory):
Is there any specific reason why the authors study the Hessian and Fisher information that is obtained from the **test data** instead of from the **expectation w.r.t. the data distribution**?

Author response:
*Question 3*. Our analysis in fact contains the analogous results for the population Hessian and G-matrices. However, from a practical/data-driven perspective, the test matrices are more natural objects to analyze as those are what the practitioner has access to.