Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The authors derive exact and non-asymptotic (in size or time) expressions for the expected test loss of a linear model during the SGD training dynamics.
The time-dependent average test loss is given as function of the eigenvectors and the eigenvalues of the covariance matrix of the features when the features are Gaussian.
For non-Gaussian features also forth-moments are involved and the Gaussian formula can be adapted to provide an upper bound. In any case, the formula for the 
Gaussian case seems to match very well the experiments on real world datasets.
The formalism covers both the online setting and the setting of multiple passes over a fixed training set. 
main_review: 
The paper is generally well written, although exposition in the main text or in the appendix could be made a bit more pedagogical. 
The results it provides are very interesting. The implication of the main theorems on batch sizes and learning rates and the connection with neural tangent kernel theory are carefully explored. Numerical experiments are extensive and well thought. 


Some additional comments:

- Use \lVert \lambda \rVert  instead of |\lambda| at the top of pag. 4 for consistency. 

- Could discuss the relation with the hidden manifold model where a linear teacher acts on the hidden space, y(x) = w*\cdot x, if the authors think there is anything interesting to 
report.

- Theorem 3.2: make explicit t depence  in c_k
- Theorem 3.4: this is assuming Gaussian features as in Theorem 3.1, right? If so, should be clearly stated. 

- Theorem 3.5: Why these expressions don't involve 4-th moments? Is some Gaussianity assumption made also here?
- Theorem 3.5: make t dependence explicit in Ltrain and Ltest
- In the MNIST and CIFAR tasks, the regression targets are 0s and 1s?
 - Related Work: "Their predicted exponent
in the small learning rate limit agrees with the exponent we derive with the saddle point approximation in Section 3.1.2.". I find no mention in the main nor the appendix of
saddle point approximations.

- Appendix A: "Since the basis of kernel eigenfunctions
(including the zero eigenvalue functions) is complete over the space of square integrable functions." should end with a comma instead?
- Appendix A: the first N eigenvalues of the kernel K are the same as the eigenvalues of N? The connection could be made more clear and in general this appendix expanded a bit to be more pedagogical. 

summary_of_the_review: The paper is a nice contribution to the theory of SGD in the quite fundamental linear regression setting. It is well written, contains original and relevant results, and provides good experimental verification. Therefore I recommend acceptance on ICLR without esitation.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper addresses the problem of characterising analytically the dynamics of stochastic gradient descent (SGD) in problems where the data have features with arbitrary covariance structure. All the results are valid for linear models (e.g., random features, neural tangent kernel) trained on the mean-squared error loss. First, the authors consider one-pass SGD and Gaussian features, and derive an exact closed-form expression for the time-evolution of the expected test loss, i.e., averaged over the data distribution and the sampling sequence. Then, they turn to generic non-Gaussian feature maps with arbitrary covariance and similarly compute an exact closed-form expression for the expected test loss of one-pass SGD. They also derive an upper bound that only depends on the features covariance, provided that a regularity condition holds on the fourth moments of the features. This result shows that non-Gaussian effects are negligible in the settings under consideration. Finally, they extend their results to multi-pass SGD and provide expressions for the expected test and training losses over time. The authors apply their theoretical findings to derive heuristic estimates of the optimal batch size and learning rate and study their dependence on the features and target distributions. The numerical simulations show a good agreement with the theoretical predictions.
main_review: 
I really liked reading this paper. Overall, it is well written and the results are presented in a clear way. It points towards a very important direction of investigation that is attracting increasing interest in machine learning theory. However, the discussion is missing some very relevant references that already cover some of the results, and therefore the authors over-claim some of their contributions.
Missing references:

(i)  Ref. [1] derives rigorous conditions under which learning with data coming from a single-layer generator can be analysed via an equivalent Gaussian model. Moreover, they show for a broader class of data distributions — in particular, pre-trained deep generators — that the theoretical predictions for both online-SGD and full-batch learning match the experimental curves for the loss;

(ii) Ref. [2] extends the above work to include target distributions generated from a pre-trained teacher with a static feature map.

These works also consider other losses than the MSE (e.g., logistic loss). I believe that the authors should clarify the significance of their contributions with respect to (i) and (ii).

The paper presents interesting results on the optimal hyper-parameters, i.e., learning rate and batch size, which could lead to useful guidance for practitioners. In particular, the authors highlight the dependence of the optimal batch size on the power-law decay of the feature distribution and on the noise level of the target distribution.

Minor comment: Figure 2(a) would be more readable if there was a legend indicating the correspondence between the values of $b$ and the different colors.

Figure 3(a): can the authors provide an intuitive explanation of why the mismatch between theory and simulations increases with the learning rate? In particular, why do simulations seem to consistently stay below theoretical curves? On how many seeds have the simulations been averaged over? 

In this paper the compute budget $C=tm$ is somewhat assumed to be a good measure of computational limitations, and a tradeoff between performance and compute budget implicitly emerges as a suggested goal. However, gradient evaluations can be efficiently parallelized for large batch sizes, and it is often argued that settings where large batches still generalize well are a desirable goal (see, for instance, [3]). I would like to have some clarifications on this point.

The most interesting and novel aspect of the work in my opinion is the study of multi-pass gradient descent, which — to the best of my knowledge — has only been previously done in [4,5]. The setting in [4,5] encompasses more general sampling procedures and generic (possibly non-convex) loss functions, at the price of a more complicated system of equations for the learning curves. On the contrary, in this paper the final expressions for the test and train losses are nice and simple. It seems to me that equations (71) can only be obtained for sampling with replacement since the weights ${\bf w}(t)$ at time $t$ must be independent on the mini-batch drawn to compute ${\bf w}(t+1)$. I would appreciate if the authors could elaborate more on the limitations/potential extensions of their analysis to more complicated loss functions and different sampling procedures/algorithms (e.g., SGD with momentum, sampling without replacement, curriculum learning…).

Minor comment: I was surprised to see that the scaling of fluctuations turns out to be $\eta^2/m$, at variance with what is usually assumed ($\eta/m$) in works that model the flow dynamics of SGD as a Langevin-like process (see, e.g., [6]). I do not know if this point is interesting/worth of comments.

Some typos are present both in the text and in the equations. For instance:
-  The left-hand side of eq. (4) is missing $y_{\perp}(x)$;
-  Right below eq. (6) the average brackets are missing: $L_t$ should be $\langle L_t \rangle$;
-  The left-hand side of eq. (67), Appendix I should be ${\bf w}_{t+1}$;
-  Appendix I ends with an incomplete sentence.

References:

[1] arXiv:2006.14709

[2] arXiv:2102.08127

[3] arXiv:1811.03600

[4] arXiv:2006.06098

[5] arXiv:2103.04902

[6] arXiv:1711.04623
summary_of_the_review: Although the paper is well written and points towards a relevant direction of research, I believe that the contributions are rather incremental with respect to previous literature that is not taken into account. Therefore, I think that at the present stage this work is not ready for publication.

*** Update after reading the authors response***

The authors have carefully addressed my concerns and improved their manuscript accordingly. I recommend the updated version for publication and I have therefore raised my score to 6.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper studies the relation between the test error of stochastic gradient descent on training linear models and the structure of the data distribution, the iteration number, and the batch size. The analyses are first done on one-pass SGD and then extended to multi-pass SGD. Results in theory and real data experiments are presented to illustrate the learning curves of the SGD algorithm.
main_review: I find the following pros and cons of the paper.

Pros: 

- The paper is written clearly. The problem setup is described rigorously and cleanly, and the content of the paper is well-organized.

- The results of the paper are solid. The theoretical results are correct, and the experiments are also reasonable.

Cons:

- The theoretical results in this paper may be a bit incremental given existing results. This paper misses some important related works on learning overparameterized linear models in linear regression, ridge regression, and logistic regression [2,3,4,5,6]. Notably, all the works above also focus on the relation between the test loss/error and the structure of the data. Therefore the authors should discuss how the results in this paper are different from these existing works.

- This paper claims to cover the training of deep neural networks. This may be overselling. The connection between the analysis in this paper and deep neural networks is mainly through the recent results on Neural Tangent Kernel. However, it is believed in the literature that the NTK setting cannot represent deep neural networks. Moreover, according to [1], neural networks under the NTK setting corresponds to a random feature model, and in literature, the random feature models are commonly studied based on its Gram matrix, just like what has been done in [1,2,3,4,5,6]. However, the results in this paper are all given in terms of the data covariance matrix, which does not seem appropriate for possibly infinite-dimensional feature maps or random feature maps.

- The study in this paper is not very comprehensive either. Although the authors present theories as well as experiments, they are not exactly to back up each other. Instead, it seems that many of the experiments are presented to handle results that cannot be covered by theories, such as the optimal mini-batch size under a fixed compute budget. I think it could greatly improve the quality of theories of the paper if the authors can use them to more comprehensively predict/explain the experiment results.


[1] Cao and Gu. "Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks." NeurIPS 2019

[2] Bartlett et al. "Benign overfitting in linear regression.", PNAS, 2020

[3] Tsigler and Bartlett. "Benign overfitting in ridge regression." arXiv:2009.14286, 2020

[4] Chatterji and Long. "Finite-sample analysis of interpolating linear classifiers in the overparameterized regime." JMLR, 2021

[5] Zou et al. "Benign Overfitting of Constant-Stepsize SGD for Linear Regression", COLT 2021

[6] Cao et al. "Risk bounds for over-parameterized maximum margin classification on sub-gaussian mixtures." NeurIPS 2021

summary_of_the_review: Based on the cons listed in the Main Review section, I would like to give a rating of 5 for the current version of the submission.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper studies the learning dynamics of stochastic gradient descent on simple linear models with structured features. In particular, the paper discussed the influence of data structure on learning dynamics and optimal batch sizes. In practice, this model seems to be able to predict the training/test error of small neural networks on real data.

main_review: Strengths:
- The paper analyzes the influence of data structure (e.g., the spectrum of covariance) on the loss dynamics, which is very interesting.

Weaknesses:
- While this paper studies an interesting problem, the contributions and novelty of the paper are unclear. Similar models have been studied extensively for either least-squares regression or neural network training. For all the subsections in section 3, there are similar works done before. I believe the presentation of the paper could be largely improved by first explicitly stating the existing work (rather than deferring the related work to the end of the paper) and then introducing the novel contributions of this work. For example, [1, 2] did very similar analyses about the learning dynamics of SGD or other algorithms. In addition, [1] also looked at power-law distribution.
- For the discussion on power-law eigenspectrum, it seems that the authors assume the dimensionality is infinite which leads to the power-law error dynamics. However, the features are often finite dimension (other than kernel methods), so the motivation of assuming infinite dimensionality is unclear to me. It is important to clarify this assumption since the authors use this model to predict neural network training dynamics. To my knowledge, the power-law error in practice is due to the label noise, rather than the artificial result of infinite dimensionality.
- The discussion on the relationship between optimal batch size and feature covariance is new. However, I don’t find it surprising as many empirical investigations have shown that the optimal batch size depends on the specific dataset.


[1] Which algorithmic choices matter at which batch sizes?   
[2] Accelerating Stochastic Gradient Descent For Least Squares Regression


summary_of_the_review: This paper studies an interesting problem, but its contributions are unclear to me as similar models have been studied before. Rederiving the loss dynamics with different data structures makes the paper a bit incremental. Overall, this is a borderline paper and I'm inclined to vote for a reject. That being said, I'm happy to increase the score if the authors could reorganize the paper and emphasize their novel contributions.
review_title: null

Meta-review: The work presented in this study gives a theoretical finite-sample generalisation performance of stochastic gradient descent on linear models, for different batch-sizes and feature structures. This approach enable the authors to predict the training and test losses of neural networks on real data.

While there were some parts that were initially mis-understood by some reviewers in the initial version of the papers, the extensive discussions between the authors and the reviewers led to several updates, both in the reference to prior work, but also in the presentation clarity. The wide impact and relevance to ICLR of this type of contribution made us recommend this work for acceptance at ICLR.
Decision: Accept (Poster)