Learning Curves for SGD on Structured Features
Authors: Blake Bordelon
Cengiz Pehlevan
Venue: ICLR 2022 Poster (2022)
Keywords: Stochastic Gradient Descent
Generalization

Abstract: The generalization performance of a machine learning algorithm such as a neural network depends in a non-trivial way on the structure of the data distribution. To analyze the influence of data structure on test loss dynamics, we study an exactly solveable model of stochastic gradient descent (SGD) on the square loss which predicts test error when training on features with arbitrary covariance structure. We solve the theory exactly for both Gaussian features and arbitrary features and we show that the simpler Gaussian model accurately predicts test loss of nonlinear random-feature models and neural networks in the kernel regime trained with SGD on real datasets such as MNIST and CIFAR-10. We show that the optimal batch size at a fixed compute budget is typically small and depends on the feature correlation structure, demonstrating the computational benefits of SGD with small batch sizes. Lastly, we extend our theory to the more usual setting of stochastic gradient descent on a fixed subsampled training set, showing that both training and test error can be accurately predicted in our framework on real data.

Published as a conference paper at ICLR 2022
LEARNING CURVES FOR STOCHASTIC GRADIENT DE- SCENT ON STRUCTURED FEATURES
Blake Bordelon & Cengiz Pehlevan
John A. Paulson School of Engineering and Applied Sciences
Center for Brain Science Harvard University Cambridge, MA 02138, USA
{blake eg hwrvard. edu
ABSTRACT
The generalization performance of a machine learning algorithm such as neu- ral network depends in intricate way the structure of the data distribution. To analyze the influence of data structure on test loss dynamics, we study an ex- actly solveable model of stochastic gradient descent (SGD) which predicts test loss when training on features with arbitrary covariance structure. We solve the theory exactly for both Gaussian features and arbitrary features and show that the simpler Gaussian model accurately predicts test loss of nonlinear random- feature models and deep neural networks trained with SGD on real datasets such as MNIST and CIFAR-10. We show that the optimal batch size at a fixed com- pute budget is typically small and depends on the feature correlation structure, demonstrating the computational benefits of SGD with small batch sizes. Lastly, we extend our theory to the more usual setting of stochastic gradient descent on a fixed subsampled training set, showing that both training and test error can be accurately predicted in our framework on real data.
INTRODUCTION
Understanding the dynamics of SGD on realistic learning problems is fundamental to learning the- ory. Due to the challenge of modeling the structure of realistic data, theoretical studies of general- ization often attempt to derive data-agnostic generalization bounds or study the typical performance of the algorithm on high- digh-dimensional, factorized data distributions (Engel Van den Broeck (001). Realistic datasets, however, often lie on low dimensional structures embedded in high dimensional ambient spaces (Pope et al. 202T). For example, MNIST and CIFAR-10 lie on surfaces with intrin- sic dimension of ~ 14 and ~ 35 respectively (Spigler et al. 2020). To understand the average-case performance of SGD in more realistic learning problems and its dependence on data, model and hyperparameters, incorporating structural information about the learning problem is necessary.
In this paper, we calculate the average case performance of SGD on models of the form f(x) = w- w-xb(x) for nonlinear feature map &a trained with MSE loss. We express test loss dynamics in terms of the induced second and fourth moments of sb. Under a regularity condition on the fourth moments, show that the test error can be accurately predicted demonstrate the accuracy of our theory on random feature models and wide neural networks trained on MNIST and CIFAR-10 10 and accurately predict test loss scalings on these datasets. We explore in detail the effect of minibatch size, m, on learning dynamics. By varying m, we can interpolate our theory between single sample SGD (m 1) and gradient descent on the population loss (m -> oc). To explore the computational advantages SGD compared to standard full batch gradient descent we analyze the loss achieved at a fixed compute budget C' = fm for different minibatch size m and number of steps f, trading off the number of parameter update steps for denoising through averaging. We show that generally, the optimal batch size is small, with the precise optimum dependent on the learning rate and structure of the features. Overall, our theory shows how learning rate, minibatch size and data structure interact with the structure of the learning problem to determine generalization dynamics. It provides predictive account of training dynamics in wide neural networks.
terms of second moments alone. We

---

Published as a conference paper at ICLR 2022
1.1 OUR CONTRIBUTIONS
The novel contributions of this work are described below.
We calculate the exact expected test error for SGD on MSE loss for arbitrary feature struc- ture in terms of second and fourth moments as we discuss in Section 4.2 We show how structured gradient noise induced by sampling alters the loss curve compared vanilla GD.
For Gaussian features (or those with regular fourth moments), we compute the test error in Section I This theory shown to be accurate in experiments with random feature models and wide networks in the kernel regime trained on MNIST and CIFAR-10.
We show that for fixed compute/sample budgets and structured features with power law spectral decays, optimal batch sizes are small. We study how optimal batch size depends on the structure of the feature correlations and learning rate.
We extend our exact theory to study multi-pass SGD on a fixed finite training set. Both test and training error can be accurately predicted for random feature models on MNIST.
2 RELATED WORK
The analysis stochastic gradient descent has a long history dating back to seminal works of Polyak & Juditsky (1992) and Ruppert (1988), who analyzed time-averaged iterates in noisy problem. Many more works have examined a similar setting, identifying how averaged and accelerated ver- sions of SGD perform asymptotically when the target function is noisy (not deterministic function of the input) (Flammarion & Bach] 2015 2017: Shapiro 1989 Robbins & Monro 1951, Chung Yu al. 2020 Anastasiou et al. 2019) Gurbuzbalaban et al.) 2021).
1954 Duchi & Ruan 2021
Recent studies have also analyzed the asymptotics of noise-free MSE problems with arbitrary fea- ture structure see what stochasticity arises from sampling. Prior works have found exponen- tial loss curves problems as an upper bound (Jain et al. 2018) or as typical case behavior for SGD on unstructured data (Werfel et al., 2004). A series of more recent works have considered the over- parameterized (possibly infinite dimension) setting for SGD, deriving power law test loss curves emerge with exponents which are better than the O(t-1) rates which arise in the noisy prob-
lem (Berthier et al. 2020. Pillaud- Vivien et al.) 2018; Dieuleveut et al.) 2016 Varre et al. 2021 Dieuleveut Bach| 2016 Ying & Pontil 2008] Fischer & Steinwart] 2020) Zou et al. [002]) These
works provide bounds of the form Olt ?) for exponents / which depend on the task and feature distribution.
Several works have analyzed average case online learning in shallow and two-layer neural networks.
Classical works often analyzed unstructured data (Heskes & Kappen 1991; Bichl & Riegler 1994. Mace & Coolen 1998 Saad & Solla 1999: LeCun et al. 1991 Goldt et al. 2019), but recently the hidden manifold model enabled characterization of learning dynamics in continuous time when trained on structured data, providing an equivalence with Gaussian covariates model (Goldt al.) 2020 202T). In the continuous time limit considered in these works, SGD converges to gradient flow on the population loss, where fluctuations due sampling disappear and order parameters obey deterministic dynamics. Other recent works, however, have provided dynamical mean field frameworks which allow for fluctuations due to random sampling of data during a continuous time limit of SGD, though only on simple generative data models (Mignacco et al. 2020 2021).
Studies of fully trained linear (in trainable parameters) models also reveal striking dependence on data and feature structure. Analysis for models trained on MSE (Bartlett et al. 2020) sigler &
Bartlett 2020; Bordelon et al. 2020 Canatarr et al. 2020), hinge loss (Chatterji Long 2021 Cao
al., 2021; Cao & Guj 2019) and general convex loss functions (Loureiro al. 202 have now
been performed, demonstrating the importance of data structure for offline generalization.
Other works have studied the computational advantages of SGD at different batch sizes Ma et al. (0018) study the tradeoff between taking many steps of SGD at small m and taking a small number of steps at large After critical m, they observe a saturation effect where increasing m provides diminishing returns. Zhang et al. (2019) explore how this critical batch size depends on SGD and momentum hyperparameters in noisy quadratic model. Since they stipulate constant gradient noise
2

---

Published as a conference paper at ICLR 2022
induced by sampling, their analysis results in steady state error rather than convergence at late times, which may not reflect the true noise structure induced by sampling.
PROBLEM DEFINITION AND SETUP
We study stochastic gradient descent on a linear model with parameters and feature map (x( E R. (with possibly infinite). Some interesting examples of linear models are random feature models, where (xx = o OGx for random matrix G and point-wise nonlinearity & (Rahimi & Recht 2008 Mei & Montanari 2020). Another interesting lincarized setting is wide neural net- works with neural tangent kernel (NTK) parameterization (Jacot et al.
2020 Lee et al.
2020). Here 0, at initial-
the features are parameter gradients of the neural network function ab x
Vol f
ization. We will study both of these special cases in experiments.
We optimize the set of parameters w by SGD to minimize a population loss of the form
L(w)
(ww <>x b(x))² x-<p(x)
(1)
where x are input data vectors associated with probability distribution p(x), x6(x) is a nonlinear feature map and y(x) is target function which we can evaluate training samples. We assume that the target function is square integrable (u(x)²) oo over p(x). Our aim is to elucidate how this population loss evolves during stochastic gradient descent on w. We derive a formula in terms of the cigendecomposition of the feature correlation matrix and the target function
Σ Σ= (xx((xxxxxxx
Σ Au v(x) Σ Ur (x(x) + y 1i (x):
where (x))))(xx) We justify this decomposition of (xx) in the Appendix A using an cigendecomposition and show that it is general for target functions and features with fimite variance.
During learning, parameters w are updated to estimate a target function y which, as discussed above, can generally be expressed as a linear combination of features y = w *bb + M1- At each time step t, the weights are updated by taking a stochastic gradient step on a fresh mini-batch of m examples
W1+1
The (W, abt. Vt..s)>
where each of the vectors tit,, are sampled independently and Dt,a w' tb... The learning rate controls the gradient descent step size while the batch size m gives a empirical estimate of the gradient at timestep At each timestep, the test-loss, or generalization error, has the form
L, = ((w. (x)
<xxx y (x)))) (w.- w²) E(w, w*) (x))),
(4)
which quantifies exactly the test error of the vector w,. Note, however, that L; is a random variable since w, depends on the precise history of sampled feature vectors D; = { Our theory, which generalizes the recursive method of (Werfel et al., 2004) allows us compute the expected test loss by averaging over all possible sequences to obtain (L.),.- Our calculated learning curves are not limited to the one-pass setting, but rather can accommodate sampling minibatches from a finite training set with replacement and testing on separate test set which we address in Section 4.4
In summary, we will develop theory that predicts the expected test loss (Ia) averaged over training sample sequences D, in terms of the quantities {Ak, Uks (v=())²) J. This will reveal how the structure in the data and the learning problem influence test error dynamics during SGD. This theory is a quite general analysis of linear models on square loss, analyzing the performance of linearized models on arbitrary data distributions, feature maps 8/, and target functions yy(x).
ANALYTIC FORMULAE FOR LEARNING CURVES
4.1 LEARNABLE AND NOISE FREE PROBLEMS
Before studying the general case, we first analyze the setting where the target function is learnable, meaning that there exist weights w such that y(x) w' (x).. For many cases of interest, this
3

---

Published as a conference paper at ICLR 2022
is a reasonable assumption, especially when applying our theory to real datasets by fitting atomic measure on P points p E, 6(x x"). We will further assume that the induced feature distribution is Gaussian so that all moments of aa can be written in terms the covariance Σ. We will remove these assumptions in later sections.
Theorem 1. Suppose the features db follow a Gaussian distribution sb ~ NO Σ) and the target function is learnable in these features y = ww- After t steps of SGD with minibatch size m and learning rate n,, the expected (over possible sample sequences De) test loss (LiYD. has the form
of m
112
(L4) A A A= A- diag x))²
diag (22)
AX
(5)
m
where X is vector containing the eigenvalues of Σ and v² is vector containing elements (v²)s. o (ug w² for eigenvectors us of Σ. The function diag(-) constructs diagonal matrix with the argument vector placed along the diagonal.
Proof. See Appendix B for the full derivation. We will provide brief sketch of the proof here. The strategy of the proof relies on the fact that (Le) = Tr Σ C, where C, ((W:- w') (We w*) >')T>,, We derive the following recursion relation for this error matrix
C++1 (I- CT, 1/E) + z ECC ETr(EC,)
(6)
The loss only depends CKtt /n C.ug- Solving the recurrence, C₂ A'v? and using (Ix) Σx CR A'v?, , we obtain the desired result.
Σx u C.u,
Below we provide some immediate interpretations of this result.
The matrix A contains two components; a matrix (I dia) which represents the time-evolution the loss under average gradient updates. The remaining matrix (A3) + XAT) arises due to fluctuations in the gradients, a consequence of the stochas-
z m
tic sampling process.
The test loss obtained when training directly on the population loss can be obtained by taking the minibatch size m + co. In this case, A + (I diag (A))² and one obtains the population loss I, = Σ. Ax 77 )² This population loss can also be obtained by considering small O(v²).
learning rates, i.e. the 1 + limit, where A (I 1/ diag(A))*
For general À and 17" /m > 0, A is non-diagonal, indicating that the components {ui.... u are not learned independently as increases like for L but rather interact during learning due to non-trivial coupling across eigenmodes large r7² ///m. This is unlike offline theory for learning in feature spaces (kernel regression), (Bordelon et al. 2020; Canatar et al.) (020). this observation of mixing across covariance eigenspaces agrees with recent analysis of SGD, which introduced recursively defined "mixing terms" that couple each mode's evolution (Varre et al. (002))
Though increasing m always improves generalization at fixed time (proof given in Appendix D) learning with fixed compute budget (number of gradient evaluations) C' tm, can favor smaller batch sizes. We provide an example this in the next sections and Figure (d)-(f).
The lower bound (L.) > AT 2
17) + n² |A|2]
can be used to find necessary stability con-
ditions on m, 17. This bound implies that (L.) will diverge if m < -n- A| The learning rate must be sufficiently small and the batch size sufficiently large to guarantee convergence. This stability condition depends on the features through |A| = Σx X One can derive heuristic optimal batch
sizes and optimal learning rates through this lower bound. See Figure 2 and Appendix C
4.1.1 SPECIAL CASE 1: UNSTRUCTURED ISOTROPIC FEATURES
This special
previously analyzed by Werfel et al. (2004) which takes Σ = I NxN and m = 1. We extend their result for arbitrary m, giving the following learning curve
(L) = (11 (1-n)) 1 1+ *+NN)) II, (TG) = =(1
m
1| 12,
(7)
m m+ N+N++
4

---

Published as a conference paper at ICLR 2022
10
10c
19
10-4
10
N-100
-0
100
20200
N- 200
MNIST 8-9
N-500
Gauss. Theory
500
1000
1500
2000
2500
10'
10
16
10"
10a
10
10'
(a) Isotropic Features
(b) Power Law Features
(c) MNIST Random ReLU Fea-
tures
G- 0.1e"
A- 0.5
Q- 0.3
30.10"
60x10
0.5.5"
1or4
12x1000>
q 1.0g"
0.1
40-3
0-5.5
10r4
10:1
16a
10
(d) Fixed Compute Isotropic
(e) Fixed Compute Power Law
(l) Fixed Compute ReLU MNIST
Figure 1: Isotropic features generated 2b ~ NO I) have qualitatively different learning curves than power-law features observed in real data. Black dashed lines are theory. (a) Online learning with N-dimensionall isotropic features gives a test loss which scales like L, ~ e -/NN for any target function, indicating that learning requires t ~ N steps of SGD, using the optimal learning rates m
17" N4m44 (b) Power-law features &b ~ NN A) with Agi = ōx, have non-extensive give a
power-law scaling L, - with exponent / On (1). (c) Learning discrimninate MNIST 8's and 9's with N = 4000 dimensional random ReLÚ features (Rahimi & Recht 2008). generates a power law scaling at large t, which is both quantitatively and qualitatively different than the scaling predicted by isotropic features e -I/N, (d)-(f) The loss a fixed compute budget C' = tyn 100 for (d) isotropic features, (e) power law features and (f) MNIST ReLU random features with simulations (dots average and standard deviation for 30 runs). Intermediate batch sizes are preferable on real data.
where the second expression has optimal First, we note the strong dependence on the ambient dimension N: as >> m, learning happens at a rate (Ix) ~ e -tm/NN Increasing the minibatch size m improves the exponential rate by reducing the gradient noise variance. Second, we note that this feature model has the same rate of convergence for every learnable target function y. At small m, the convergence any learning rate 7) is much slower than the convergence of the m do limit, Lpop = (1 [[]]][ which does not suffer from a dimensionality dependence due to gradient noise. Lastly, for fixed compute budget C' tm, the optimal batch size is
see Figure (d). This can be shown by differentiating (Le/m) with respect to m (see Appendix E). In Figure
(a) we show theoretical and simulated learning curves for this model for varying values of N the optimal learning rate and in [ we show the loss as function of minibatch size for fixed compute budget fm 100. While fixed represents fixed sample complexity, we stress that it may not represent wall-clock run time when data parallelism is available (Shallue et al. 2018
4.1.2 SPECIAL CASE 2: POWER LAWS AND EFFECTIVE DIMENSIONALITY
Realistic datasets such as natural images or audio tend to exhibit nontrivial correlation structure, which often results in plwer-law spectra when the data is projected into a feature space, such randomly intialized neural network (Spigler et al.) 2020 Canatar et al.) 2020) Bahri et al.) 2021).
In the N² < limit. if the feature spectra and task specra follow power laws, Aς- k and m da k-a with b 1, then Theorem implies that generalization error also falls with a power where C is constant. See Appendix G for derivation with
law: (L,) Ct-8, B & saddle point integration. Notably. these predicted exponents we recovered as a special case of our theory agree with prior work on SGD with power law spectra, which give exponents in terms of the
5

---

Published as a conference paper at ICLR 2022
Verying
Optimal Stales with [2]
log cco
107
True Optimum
2.
0.90
10a2
Z
0.00
$10
0.0.8
6.0.5
-0.00
-445
10
1.80
-225
10-1
10
10
(a) Power Law Exponents b
(b) Optimal Batchsize vs
(c) Hyper-parameter Dependence
Figure 2: Optimal batch size depends on feature structure and noise level. (a) For power law features
A b-b X k the m dependence of the loss T.c/m depends strongly the feature exponent b. Each color is a different value, evenly spaced in [0.6, 2.5] with = 2.5, C 500. Solid lines show exact theory while dashed lines show the error predicted by approximating the mode
coupling term z²_ XAT T with decoupled - diagg(A2) )²)). Mode coupling is thus necessary to accurately predict optimal m. (b) The optimal m scales proportionally with |사| T We plot the lower bound mmiin (black), the heuristic optimum (m which optimizes lower bound for T., green) and 2 |X|² (red). (c) The loss at fixed compute C' 150, 2, b 0.85, optimal batchsize for 2-n
m
each 17 shown in dashed black. For sufficiently small 77, the optimal batchsize is = 1. For large 17, is better to trade off update steps for denoised gradients resulting in m >
feature correlation structure (Berthier et al. 2020, Dieuleveut et al. 2016 Velikanov & Yarotsky, 2021 Varre al., 2021). Further, our power law scaling appears to accurately match the qualitative behavior of wide neural networks trained on realistic data (Hestness et al. 2017 Bahri et al. 2021). which we study in Scction
We show example of such power law scaling with synthetic features in Figure (b). Since the total variance approaches a finite value as - oo, the learning curves are relatively insensitive to N, and are rather sensitive to the eigenspectrum through terms like | and 1 1 X, etc. In Figure (c), we see that the scaling of the loss is more similar to the power law setting than the isotropic features setting in random features model of MNIST, agreeing excellently with our theory. For this model, we find that there can exist optimal batch sizes when the compute budget C = Im is fixed (Figure (e) and (f)) In Appendix C.1 heuristically argue that the optimal batch size for power law features should scale as, m A
Figure
tests this result.
126-1]
We provide further evidence of the existence of power law structure on realistic data in Figure (a)-(c), where we provide spectra and test loss learning curves for MNIST and CIFAR-10 on ReL. random features. The eigenvalues Åς k- and the task power tail sums Σ > k-++1 both follow power laws, generating power law test loss curves. These learning curves are contrasted with isotropically distributed data in 784 passed through the same ReLU random feature model and we see that structured data distributions allow much faster learning than the unstructured data. Our theory is predictive across variations in learning rate, batch size and noise (Figure 3).
4.2 ARBITRARY INDUCED FEATURE DISTRIBUTIONS: THE GENERAL SOLUTION
The result in the previous section was proven exactly for Gaussian vectors (see Appendix B). For arbitrary distributions, we obtain a slightly more involved result (see Appendix()
Theorem 2. Let (xx) E N be arbitrary feature map with covariance matrix Σ En da uu u. After diagonalizing the features Φu(x) = un (x), introduce the fourth moment tensor KEjR = ( The expected loss exactly (Lt) = En (xC,(A.K.V,t)
We provide exact formula for C in the Appendix F We that the test loss dynamics depends only on the second and fourth moments of the features through quantities A and Kijke respectively. We recover the Gaussian result special case when Keyht is simple weighted of these three Gause Xzdzōst + AsAadi dRi + Ned5ōudze. As alternative to
products of Kronecker tensors K iikl

---

Published as a conference paper at ICLR 2022
10"
10F
10
-
|| 10-
MNIST
NNIST ReLU
MNST RaLU
CIFAR ReLU
10
CIFAR ROLU
CIFAR RELU
10a
1sp. ReLU
30a
10
10
10
10
10
10
10
10
10'
(a) Feature Spectra
(b) Task Power Tail Sum
(c) Learning Curves m
10a
10
10
3x10--
000.06
7-0.30
C 0.00
2x10-1
n= 20
0.50
1.00
10
10a
10
10
10
107
10
10 10a
10
(d) Vary 7 (20 trials)
(e) Vary m (20 Trials)
(1) Vary a (20 Trials)
Figure 3: Structure in the data distribution, nonlinearity, batchsize and learning rate all influence learning curves. (a) ReLU random feature embedding in NN = 4000 dimensions of MNIST and CIFAR images have very different eigenvalue scalings than spherically isotropic vectors in 784 dimensions. (b) The task power spectrum decays much faster for MNIST than for random isotropic vectors. (c) Learning curves reveal the data-structure dependence of test error dynamics. Dashed lines are theory curves derived from equation. (d) Increasing the learning rate increases the initial speed of learning but induces large fluctuations the loss and can be worse at large t. Experiment curves averaged over 20 random trajectories of SGD. (e) Increasing the batch size alters both the average test loss L, and the variance.
(f) Noise in the target values during training produces an asymptotic error Loo which persists even t + co.
the above closed form expression for (Lz), recursive formula which tracks N mixing coefficients has also been used to analyze the test loss dynamics for arbitrary distributions (Varre et al. 2021).
Next we show that regularity condition, similar to those assumed in other recent works (Jain et al., 2018 Berthier et al., 2020; Varre al., (202T), . on the fourth moment structure of the features allows derivation of an upper bound which is qualitatively similar to the Gaussian theory.
Theorem 3. If the fourth moments satisfy (xpup T Gubup
(a 1)ΣGE + &EETE for any
positive-semideffiite G. then
+ am? m
L, 5 A (I n dia)
[diag (2 ))+ AAT
We provide this proof in Appendix We note that the assumed bound on the fourth moments is tight for Gaussian features with the fourth moments satisfied, then the loss for the non-Gaussian features is upper bounded by the Gaussian test loss theory with the batch size effectively altered m = m/a.
recovering our previous theory. Thus, if this condition on
The question remains whether the Gaussian approximation will provide an accurate model on re- alistic data. We do not provide proof of this conjecture, but verify its accuracy in empirical experiments on MNIST and CIFAR-10 as shown in Figure 3 In Appendix Figure F.T we show that the fourth moment matrix for ReLU random feature model and its projection along the eigenbasis of the feature covariance is accurately approximated by the equivalent Gaussian model.
4.3 UNLEARNABLE OR NOISE CORRUPTED PROBLEMS
In general, the target function y(x) may depend on features which cannot be expressed linear combinations of features ab (x), v(x) w +b(x) + vi(x). Let (x))) o². Note that yı need not be deterministic, but also be stochastic process which is uncorrelated with vb(x).
7

---

Published as a conference paper at ICLR 2022
(a) MNIST Training Error
(b) MNIST Test Error
Figure 4: Training and errors of model trained a training of size M be computed with the C, matrix. Dashed black lines are theory. (a) The training error for MNIST random feature model approaches zero asymptotically. (b) The test error saturates quantity dependent on M.
Theorem 4. For a target function with unlearnable variance (v1) the expected test loss has the form
trained on Gaussian sb.
1 1 A) - -I-A1 AD)A (L)) --=
(9)
m
which has an asymptotic, irreducible error (Loo) o² m 1 (I A) A OC.
See Appendix H for the proof. The convergence to the asymptotic error takes the form (L, Loo) A A' (v² - o A)) We note that this quantity is not necessarily monotonic in and can exhibit local maxima for sufficiently large o², as in Figure (f).
4.4 TEST/TRAIN SPLITS
Rather than interpreting our theory as a description the average test loss during SGD in a one-pass setting, where data points sampled from the distribution each step of SGD, theory be suitably modified to accommodate multiple random passes over a finite training set. To accomplish this, one must first recognize that the training and test distributions are different.
Theorem 5. Let p(x) = M En 5(x - x!") be the empirical distribution on the M training data points and ler E = (xb(x))((x)) xx-p(x) Ev. Su be the feature correlation matrix on this training set. Let p(x) be the test distribution its corresponding feature correlation. Then we have
(Lrain,)) Tr [
(Lies,,) Tr [EC]]]
1/2 (x(())((x)) (x)(()(( (x) )x-((x) m
(I-)=(In
SC,S
(10)
We provide the proof of this theorem in Appendix The interpretation of this result is that it pro- vides the expected training and test loss if, at each step of SGD, points from the training set {x' x.. are sampled uniformly with replacement and used to calculate stochastic gradient. Note that while Σ can be full rank, the rank of Σ has rank upper bounded by M, the number of training samples. The recurrence for C, can again be more casily solved under a Gaussian approx- imation which we employ in Figure 4 Since learning will only occur along the M dimensional subspace spanned by the data, the test error will have an irreducible component at large time, as evidenced in Figure 4 While the training errors continue to go zero, the test errors saturate at M M-dependent final loss. This result can also allow one to predict errors on other test distributions.
5 COMPARING NEURAL NETWORK FEATURE MAPS
We can utilize our theory to compare how wide neural networks different depths generalize when trained with SGD real dataset. With certain parameterization. large width NNs approxi- mately linear in their parameters (Lee et al. 2020). To predict test loss dynamics with our theory, it therefore suffices to characterize the geometry of the gradient features xb(x) Vof(x, 0). In
8

---

Published as a conference paper at ICLR 2022
23-3
Degin Depti
bepo m3
Depes
1/4
15"
1844
(a) MNNIST NTK Spectra
(b) MNIST Task Spectra
(c) Test Loss Scaling Laws
CONY HIP
CONY
13
COW
35"
2-
13-
13-4
18-
3430-1
3433
(d) CIFAR- 10 NTK Spectra
(e) CIFAR-10 Task Spectra
(I) Test Loss Scalings
ReLU neural networks of depth D and width 500 are trained with SGD on full MNIST. (a)- (b) Feature and spectra are estimated by diagonalizing the infinite width NTK matrix on the training data. We fit simple power law each of the curves 2 -ー and 1 k (c) Experimental test loss during SGD (color) compared to theoretical power-law scalings t (dashed black). Deeper networks train faster due their slower decay in their feature eigenspectra Ag. though they have similar task spectra. (d)-(f) The spectra and test loss for convolutional and fully connected networks on CIFAR-10. 10. The CNN obtains a better convergence exponent due to its faster decaying task spectra. The predicted test loss scalings (dashed black) match experiments (color).
Figure 5:
Figure 5 we show the Neural Tangent Kernel (NTK) eigenspectra and task- lask-power spectra for fully connected neural networks of varying depth, calculated with the Neural Tangents API (Novak et al.) 2020). We compute the kernel on a subset of 10, 000 randomly sampled MNIST images and es- timate the power law exponents for the kernel and task spectra 0 and vf. Across architectures, the task spectra of highly similar, but that the kernel eigenvalues AS decay more slowly for deeper models, corresponding to a smaller exponent b. As a consequence, deeper neural network models train more quickly during stochastic gradient descent as we show in Figure 5 (c). After fitting power laws to the spectra A ~ k-b and the task power vž ~ k-a, we compared the true test loss dynamics (color) for width-500 neural network model with the predicted power-law scal- ings 3 = + from the fit exponents a, b. The predicted scalings from NTK regression accurately describe trained width-500 networks. On CIFAR- we compare the scalings of the CNN model and standard MLP and find that the CNN obtains better exponent due to its faster decaying tail sum Σm= 2 We stress that the exponents B were estimated from our one-pass theory, but were utilized experiments on finite training set. This approximate and convenient version of our theory quite accurate across these varying models, in line with recent conjectures about early training dynamics (Nakkiran et al. 2021
6 CONCLUSION
Studying a simple model of SGD, we were able to uncover how the feature geometry governs the dynamics of the test loss. We derived average learning curves (I.) for both Gaussian and general features and showed conditions under which the Gaussian approximation is accurate. The proposed model allowed us to explore the role of the data distribution and neural network architecture on the learning curves, and choice of hyperparameters on realistic learning problems. While our theory accurately describes networks in the lazy training regime, average case learning curves in the feature learning regime would be interesting future extension. Further extensions of this work could be used to calculate the expected loss throughout curriculum learning where the data distribution evolves over time well alternative optimization strategies such as SGD with momentum.

---

Published as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
The code to reproduce the experimental components of this paper can be found here https: // github , com/Pehlevan Group/sgo wttrggs_ssrctuur_efettttu
which contains jupyter notebook files which we ran in Google Colab, More details about the experiments can be found in Appendix Generally, detailed derivations of our theoretical results are provided in the Appendix.
ACKNOWLEDGEMENTS
We thank the Harvard Data Science Initiative and Harvard Dean's Competitive Fund for Promis- ing Scholarship for their support. We also thank Jacob Eavatone-Vett for useful discussions and comments on this manuscript.
REFERENCES
Andreas Anastasiou, Krishnakumar Balasubramanian, and Murat A. Erdogdu. Normal approxima- tion for stochastic gradient descent via non-asymptotic rates of martingale clt. In Alina Beygelz- imer and Daniel Hsu (eds.). Proceedings of the Thirty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp. 115-137, Phoenix, USA, 25-28 Jun 2019. PMLR. URL html
m press /99/ wtt..//rrreee///aaattttooooo anastasioul 9a.
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv:2102.06701, 2021.
Peter I. Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117(48):30633-3000 2020.
Carl Bender and Steven Orszag. Advanced Mathematical Methods for Scientists and Engineers:
Asymptotic Methods and Perturbation Theory, volume 1. 01 1999. ISBN 978-1-4419-3187-0. doi: 10.1007/978- 1141007/77811-475--333 1-4757-3009-2.
Raphaël Berthier, Francis Bach, and Pierre Gaillard. Tight nonparametric convergence rates for stochastic gradient descent under the noiseless linear model, 2020.
M Biehl and Riegler. On-line learning with a perceptron. Europhysics Letters (EPL), 28(7):525- 530, dec 1994. doi: 10112990295-505552777012 URL https:://aci.org/111 10299 0295 5075/50/777111111
Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. In International Conference Machine Learning, pp. 1024-1034. PMLR, 2020.
Abdulkadir Canatar, B. Bordelon, and C. Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature Communications, 12:1-12, 2020.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. Advances in Neural Information Processing Systems, 32: 12:10836-10846, 2019.
Yuan Cao, Quanquan Gu, and Misha Belkin. Risk bounds for over-parameterrized maximum margin classification on sub-gaussian mixtures. In Thirty-Fifil Conference on Neural Information Pro- cessing Systems, 2021. URL
Niladri Chatterji and Philip M Long. Finite-sample analysis of interpolating linear classifiers in the overparameterized regime. Journal of Machine Learning Research, 22(129);1-300 2021.
K. L. Chung. On Stochastic Approximation Method. The Annals of Mathematical Statiiss 25 (3):463 483, 1954. doi: 10. 1112144//mm/11777228 aoms 1
URL httttps://dc...og/11...
10

---

Published as a conference paper at ICLR 2022
Aymeric Dieuleveut and Francis Bach. Nonparametric stochastic approximation with large step- sizes. The Annals of Statistics, 44(4):1363 1399, 2016. doi: 10.1214/5-AAS1391, URL
https: /ttpss://oi... http:///.....111111//1111111111/ 12111111111111111111111111111111
Aymeric Dieuleveut, Nicolas Flammarion, and Francis Bach. Harder, better, faster, stronger conver- gence rates for least-squares regression. Journal of Machine Learning Research, 18, 02 2016.
John C. Duchi and Feng Ruan. Asymptotic optimality in stochastic optimization. The Annals of Statistics, 49(1):21 - 48, 2021. doi: 10.1214/19-A0S1831 URL httpt:ttt///ott. org/10.
1214/19 AOS131111 831
A. Engel and C. Van den Broeck. Statistical Mechanics of Learning. Cambridge University Press, 2001. doi:
Simon Fischer and Ingo Steinwart. Sobolev norm learning rates for regularized least-squares al- gorithms. Journal of Machine Learning Research, 21 2020. URL http://
Szc/paperssv22/11-- html
Nicolas Flammarion and Francis Bach. From averaging acceleration, there is only a step-size. In Conference on Learning Theory, pp. 658-695. PMLR, 2015.
Nicolas Flammarion and Francis Bach. Stochastic composite least-squares regression with con- vergence rate o(1/n). In Satyen Kale and Ohad Shamir (eds.), Proceedings of the 2017 Conference Learning Theory, volume 65 of Proceedings Machine Learning Rexearch, pp. 831-875. PMLR, 07-10 Jul 2017. URL https: hattps://proceedinng http://prooeeci.a............. www. mir.press/v65/
Flammarion html
Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dy- namics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019, URL htt hhtt:s::/pprocceedinng nww.rnng....euuueeeeee neur wtt.:ipp....pppeeeeeeee cc:/paperr2019/i1eee pdf
Sebastian Goldt, Marc Mézard, Florent Krzakala, and Lenka Zdeborová. Modeling the influence of data structure on learning in neural networks: The hidden manifold model. Phys. Rev. X, 10: 041044, Dec 2020. doi: 10. (003//hysRex.10.0441000 URL https:://inkk aps. org/ doii 1555/10.1103/Pyysttte 10.
Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Marc Mézard, and Lenka Zde- borová. The gaussian equivalence of generative models for learning with shallow neural networks. Proceedings of Machine Learning Research, 145:1-46. 2021.
Mert Gurbuzbalaban, Umut Simsekli, and Lingjiong Zhu. The heavy-tail phenomenon in sgd. In International Conference Machine Learning, pp.
PMLR, 2021.
Tom M. Heskes and Bert Kappen. Learning processes in neural networks. Phys. Rev. A, 44:2718 2726, Aug 1991, doi: 10. 10.1103//hysReev.44.27 URL httpp://ling http:::////kkk..s...... aps. wwt.ooggg/oooooo..ooo doi/ 10. 10.1103/PhysRevA, 44.2718
Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md. Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically, 2017.
Arthur Jacot, Berfin Simsek, Francesco Spadaro, Clément Hongler, and Franck Gabriel. Kernel alignment risk estimator: Risk prediction from training data. In NeurIPS, 2020.
cccccc/pae httpppaae//22//aait. b367e525a7e5 441771111 7el 9ac24b7635607 2o Abbstttct....ttt................ hhmml
URL
https: httpss://pccceededeee neuriis.............
Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerat- ing stochastic gradient descent for least squares regression. In Conference On Learning Theory, pp. 545-604. PMLR, 2018.
11

---

Published as a conference paper at ICLR 2022
Yann LeCun, Ido Kanter, and Sara A. Solla. Eigenvalues of covariance matrices: Application to neural-network learning. Phys. Rev. Let. 66:2396 2399, May 1991. doi: 10.1 10.1105/PhyRReveett 66.2396. htttt: aps. http://///....111/////aooooe////
66.
Jaehoon Lee, Lechao Xiao, Samuel S Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. Journal of Statistical Mechanics: Theory and Experiment, 2020(12): doi: 10. (0.1088/1742-4468/aaa
124002, Dec 2020. ISSN 1742-5468.
URL http://dx
152..org/1..188/// -177422-46658 5468/ 5468/abc62b abc62b
Bruno Loureiro, Cédric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Mare Mézard, and Lenka Zdeborová. Capturing the learning curves of generic features maps for realistic data sets with a teacher-student model, 2021.
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of sgd in modern over-parametrized learning. In ICMI., pp. 3331-3340, 2018. URL. http:://procedding ngs. press / 8a htm
C. Mace and A. Coolen. Statistical mechanical analysis of the dynamics of learning perceptrons. Statistics and Computing, 8:55-88, 1998.
Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and double descent curve, 2020.
Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynamical mean-field theory for stochastic gradient descent gaussian mixture classification, 2020.
Francesca Mignacco, Pierfrancesco Urbani, and Lenka Zdeborová. Stochasticity helps to navigate rough landscapes: comparing gradient-dese descent- based algorithms in the phase retrieval problem, 2021.
Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. In International Conference on Learning Represen- tations, 2021. URL/httts::/opperevrii
net feres//frrmmidddee qid-
Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference on Learning Representations, 2020. URL httph:t////tthhttt com/ poooll//neurl-ttee
Loucas Pillaud- Vivien, Alessandro Rudi, and Francis Bach. Statistical optimality of stochastic gra- dient descent on hard learning problems through multiple passes, 2018.
B. Polyak and A. Juditsky. Acceleration of stochastic approximation by averaging. Siam Journal on Control and Optimization, 30:838 855, 1992.
Phil Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning. In International Conference on Learning Repre- sentations, 2021. U /ttp::::////eeeeevieeee htt.:/iwwwweeeeeeeeee..eee.eeeee Fow.mmidd=..e... http://eenneeeeeeeeeeeeeeorerree
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In J. Platt, D. Koller, Y. Singer, and S. Roweis (eds.), Advances in Newral Information Processing Systems, volume 20, Curran Associates, Inc., 2008. URL https: http:///////ccecepsssoeoe wwwrips....... cc/
paper/2007/ http://000/111eee 000000033000000 3 3dbc5 92e www.see3eeeeeeeeeerraaaaaaaaaaaa -ttp::/iiis.faastsaae.eeeepeee pdf.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning). The MIT Press, 2005.
Herbert Robbins and Sutton Monro. A Stochastic Approximation Method. The Annals of Math- ematical Statistics, 22(3):400- 407, 1951. doi: 14.221//0m//1117779566 URL https:
aoms 202214///mmm/1772555552
David Ruppert. Efficient estimations from slowly convergent robbins-monro process. 02 1988.
12

---

Published as a conference paper at ICLR 2022
David Saad and Sara Solla. Dynamics of on-line gradient descent learning for multilayer neural networks. 04 1999.
Christopher Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Mcasuring the effects of data parallelism neural network training. arXiv preprint arXiv:1811.03600, 2018.
Alexander Shapiro. Asymptotic Propertics of Statistical Estimators in Stochastic Programming. The Annals of Statisties, 17(2):841 858, 1989. doi: 10.
URL https:
Stefano Spigler, Mario Geiger, and Matthieu Wyart. Asymptotic learning curves of kernel methods: empirical data versus teacher-studen. paradigm. Journal of Statistical Mechanics: Theory and Experiment. 12): Dee 2020. ISSN 1742- 5468. doi: 00.1088/1722-5466118
URL htllhttp://ale doi .
/bcbll
Alexander Tsigler and Peter L Bartlett. Benign overfitting in ridge regression. arXiv preprint [EII:2009.142866 2020.
Aditya Varre, Loucas Villaud-Vivie, and Nicolas Flammarion. Last iterate convergence of sgd for least-squares in the interpolation regime, 2021.
Maksim Velikanov and Dmitry Yarotsky. Universal scaling laws in the gradient descent training of neural networks, 2021.
Justin Werfel, Xiaohui Xie, and Seung. Learning curves for stochastic gradient descent in linear feedforward networks. In S. Thrun, I. Saul, and B. Schölkopf (eds.), Advances in Neural Information Processing Systems, volume 16. MIT Press, 2004.
http://ccee..ssssoppppeee//eoeoa URL https: Pacicce-ppprrppp.. EB032277000222560177222222888882
Yiming Ying and Massimiliano Pontil. Online gradient descent learning algorithms. Found. Comput. Math., 15):561-5 October 2008. ISSN 1615-3375.
Lu Yu, Krishnakumar Balasubramanian, Stanislav Volgushev, and Murat A. Erdogdu. An analysis of constant step size sgd in the non-convex regime: Asymptotic normality and bias, 2020.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems, 32:8 32:819688207, 2019.
Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham M Kakade. Benign over- fitting constant-stepsize sgd for linear regression. arXiv preprint arXiv:2103. 12692, 2021.
13

---

Published as a conference paper at ICLR 2022
A DECOMPOSITION OF THE FEATURES AND TARGET FUNCTION
Let xx) be square integrable target function with (u(x)²) < oc. Define the following integral
operator Tk for kernel K(x, x') <b(x) x5(x'):
Tx(o](x)= - (x)xxx)
(11)
We are interested in eigenfunctions of this operator, function ΦR for which Tx oo A For kernels with finite trace J [xx, x))(x) < oc, Mercer's ' theorem (Rasmussen & Williams/ 2005) guarantees the existence of a set of orthonormal cigenfunctions. Since (x(x) spans an I dimenstonal function space, only N of the kernel eigenfunctions will have non-zero eigenvalue. Since the basis of kernel eigenfunctions (including the zero eigenvalue functions) is complete over the space of square integrable functions. After ordering the eigenvalues A we obtain the expansion
X > AN with AN+e 0,
y(x) = Σ &
(x)) (x(x)= Σ Ux) +xbb(x)+U((x)) ASN
(x Σ KNN
(u(x)ox(x)) Φx(x)
(12)
Further, we can decompose the feature map in this basis +x(x Ex=1 Vx We
recognize through these decompositions the coefficients Uk can be computed uniquely as uh X-1/2 u (xxxx) This provides recipe for determining the necessary spectral quantities for our theory. We see that the feature map's 's decomposition above reveals that A₁ are also the eigenval- of the feature correlation matrix E since
Σ = (1b(x))(x))))
Σ
Vudeuuuu (0x (xx(x) Σ Auugu
(13)
A.1 FINITE SAMPLE SPACES
When we discuss experiments on MNIST or CIFAR, we use this technology for an atomic data distribution p(x) = Mf EM 6(x - x'). Plugging this into the integral operator gives
Txlol(x) M Σ x(xxx"))xx")
(14)
We see that, restricting the domain to the set of points { XMJ. this amounnts to solving 1 K Φb AudE where K IRMM the kernel gram matrix with matrix eigenvalue problem 3T
entries Knv = [xx", x") and Φ5 has entries ØR =x(x"..
B PROOF OF THEOREM
Let A, w, w" represent the difference between the current and optimal weights and define the correlation matrix for this difference
C, = (A,AT) ) D.
(15)
Using stochastic gradient descent, W,+1 Wt 7/1 with gradient vector g: 1 Σm b2,b,. A.. the C, matrix satisfies the recursion
C++1 = ((A, 7g:) A. - vg::) T), = C. - (g. 17 (Ag.)) ) + + (g.g.)
(16)
First, note that since 3 are all independently sampled at timestep t, we can break up the average into the fresh batch of m samples and an average over D,-1
(g Ag) = Σ (bba (AAT) P EC,
(17)
14

---

Published as a conference paper at ICLR 2022
The last term requires computation of fourth moments
(g/g))
(b
(A.AT) D.-1 +/, 01
(18)
m
2,
( Cebpybb) 2..0""
(19)
m
First, consider the case where i - j. Letting + - tb. we need to compute terms of the form
Σ
Cu.e (enhn))))))
(20)
k
For Gaussian random vectors, resort to Wick-Isserlis theorem for the fourth moment
(Wrbherr))) = (1ygbr) (teein) + (elyche) (Warta) + (Wyth) (Wethn)
(21)
giving
gig m + EC, -
(22)
m This correlation structure for g, implies that its covariance has the form
m
1 +Sc.I-LETC
(Cov (Bi))
(23)
Using the formula for (g:g.) ) arrive at the following recursion relation for C,
C++1 = - - n, n/C, + 2 +m++111 ΣC₁ + " Tr (EC))
(24)
Since we are ultimately interested in the generalization enror (Lx) (A) EA,) Ex du C,ug, it suffices to track the evolution of Ct.k ug C,u:
TrEC,
(1 1 21/Aa= + r)²
1 1 Σ Asot.
C,+1,
2 Ct
(25)
m
m
Vectorizing this equation for с generates the following solution
1 v/² diag(X²) + r/²
C₁ A' A'ca A I 2n diag(A)
XX
(26)
m
m
The coefficient Co,6 of (1, w") To get the generalization error, merely compute (Lg)
T a, A as desired.
C PROOF OF STABILITY CONDITIONS
We will first establish the following lower bound on the loss, where without loss of generality assumed the maximum correlation is 1:
L 2 (1 - 2
2 A
(27)
m
We will then this lower bound to provide necessary conditions the learning rate and batch size for stability of the loss evolution. First, note that the following inequality holds elementwise
7 AA (I (u-n 17 diag (A)) + diag(X ( +
A N (1 n) +
X 2
A
(28)
m
m
m
Repeating this inequality times gives A'X A A gives the desired inequality. Note that this inequality is very close to the true result for
[1 (1- + | x. Using the fact that L,
15

---

Published as a conference paper at ICLR 2022
isotropic features À 1 which gives L, (1 7)2 (1212 + +1)] For anisotropic features with small learning rate, this bound becomes less tight. For the loss to converge to zero at large time, the quantity brackets must necessarily be less than one. This implies the following necessary condition on the batchsize and learning rate
2m
7/A/2 2-1 1
+ |||2
m
Mimin
(29)
where Mymin is the minimal batch size for learning rate 1) and feature covariance eigenvalues A.
C.I HEURISTIC BATCH SIZE AND LEARNING RATE
We can derive heuristic optimal choices of the learning rate and batch size hyperparameters m at a fixed compute budget which optimize the lower bound derived above.
C.1.1 FIXED LEARNING RATE
We will first consider optimizing only the batch size at fixed learning rate before discussing the optimal m when 77 is chosen optimally. The loss at fixed compute budget C = tm is lower bounded by
(1-1) (1= 17² ||
C/m
Lo/m 2 A 0² (1
(30)
For the purposes of optimization, we introduce
1/ //m and consider optimizing
f(x) x In [A + Bz] A = (1 - 7) (01n))
= A
(31)
The first order optimality condition fx( 0 implies that In(A Bz) + A+Bz 0. Letting A Ba, this is equivalent to 2 In giving solutions z E (e-1,1). 1).
Br
A 0. This equation has solutions for all valid A (0, 1)
Letting z 2(n). represent the solution to 2 + In - (1 7)) = 0, the
optimal batchsize has the form
B
111
m (n) =(A) - A 2(n) (1 - 1)
(32)
We can gain intuition for this result by considering the limit of + 0 and 17 - 1. First, in the 2n| 2mmmnnn making contact with the stability bound derived in Appendix Section C Thus for small learning rates, this heuristic optimum suggests doubling the minimal stable batchsize for optimal convergence. At large learning rates, 1) ~ 1 with A ~ O, we find ~ e-1 so m m1)) ~ e Thus for small 17. we expect m to scale linearly with y while for large 7/ expect a scaling of the form 1/2 In either case, the optimal batchsize scales with feature eigenvalues through the sum of the squares A Σx X
?) -> limit, we find that A+1
2 so m
2-n
C.1.2 HEURISTIC OPTIMAL LEARNING RATE AND BATCH SIZE
We will now examine what happens when one first optimizes loss bounmd with respect the the learning rate at any value of m and then subsequently optimizes over the batch size m. We can easily find the 7y which minimizes the lower bound.
θ Oy
A
))
(33)
m
m
A
Note that this heuristic optimal learning rate is very close to the true optimum in the isotropic data m2 m+N Plugging this back into the loss bound, we find that at fixed
setting Thraue m N I
compute C = im, the loss scales like
C/m
Le/m To
(34)
m |
With the optimal choice of learning rate, the loss at fixed compute monotonically increases with batch size, giving an optimal batchsize of m = 1. This shows that if the learning rate chosen optimally then small batch sizes give the best performance per unit of computation. This corresponds
to the hyperparameter choices (1), m)
16

---

Published as a conference paper at ICLR 2022
D INCREASING m REDUCES THE LOSS AT FIXED t
We will show that for a fixed number of steps t, increasing the minibatch size m can only decrease the expected error. To do this, we will simply show that the derivative of the expected loss with respect to m,
d (Le) X
.A v².
(35)
Im
dm
is always non-positive.
The derivative of the 1-th power of can be identified with the chain rule
aA'
0A
0A
A
+ 1-1 0A
A
1-1
+
A
1-2
(36)
dm
dm
dm
3m
Note that the matrix
0A
[diag (A²) + AA ]
(37)
m
has all non non-positive entries. Thus we find that
0 (L)) 3m
0A A 1-
ΣXTA
(38)
dm
==0
Note that since all entries in Vž and A At-n-1 are non-negative, the vector 2, A A has non-negative entries. By the same argument, the vector qa = A"X is also non-negative in each
entry. Therefore, each of the terms in a(L+) above must be non-positive
Om
0 (Tx)
AA Zn
9n
o
ΣΣ
Zn.k [ōn, + dxdee] qn.t S
(39)
dm
0m
m
R.E
Thus we find
a(L,) 8m
< 0, implying that optimal (Lg) is always obtained (possibly non-unine non-uniquely) at
m
- 00.
E INCREASING m INCREASES THE LOss AT FIXED C' tm ON ISOTROPIC FEATURES
Unlike the previous section, which considered fixed and varying m, in this section we consider fixing the total number of samples (or gradient evaluations) which we call the compute budget C Im. For fixed compute budget C Im, and unstructured N dimensional Gaussian features and optimal learning rate 17" M+N+I+ we have
N N11
C/m
(Lc/m)
1|w*1|2
(40)
m + +1
Taking derivative with respect the batch size get dm
C
N +1
log (Lo/m)
log
Om m
m N +1
C
m
m++ N +1
C
log
m(m + N 1) > 0.
(41)
m
N+1
This exercise demonstrates that, for the isotropic features, smaller batch-sizes are preferred at fixed compute budget C. This result does not hold for arbitrary spectra Aş. In the general case, optimal minibatch sizes can exist as we show in Figure [ (c)-(f) for power law and MNIST spectra.
F PROOF OF THEOREM 2
Let Vee(.) denote a flattening of an N N matrix into a vector of length N2 and let Mat(-) represent a flattening of a 4D tensor into a 2 x N2 two-dimensional matrix. We will show that the expected loss (over D₂) is
(L.) Σ AECE,EK C.
(G+ Ma(ks) ) m
112
Vec(vr R
(42)
17

---

Published as a conference paper at ICLR 2022
where [G] = Ö5 (1 (1- 17(A, + X²) + y
)x2y) and [v]x = u w"
We rotate all of the feature vectors into the eigenbasis of the covariance, generating diagonalized features @x = ug s and introduce the following fourth moment tensor
Kight = (010))EE1))
(43)
We redefine C in the appropriate (rotated) basis by projecting onto the eigenvectors of the covari- ance
C. = U (A, U,
(44)
where
[U1, u2, U.]. With this definition, C's dynamics take the form
C++1 = C - AC, - C,A+ (m - 1) AC,A + (oo )
(45)
The elements of the matrix can be expressed with the fourth moment tensor
Σ
(d:050n0)) Cke = Σ
(46)
ke
We thus generate the following dynamics for C%2y
Ctti= (1''=(-n(A++))) (d)
1) 1 7m (m 2:dg) Cjs
n
m Σ Cis
(47)
m
Let c, Vec Ct), then we have
Go + " m
Matxx)) Ct [Golijke dzzdze 1 m7d; + 22)
17 1m \rdy1
1)
C++1
(48)
Solving these dynamics for c, recognizing that co = Vec(vv and computing (L,) = TrIC Ee CREAA gives the desired result.
=
F.1 PROOF OF THEOREM 3
Suppose that the features satisfy the regularity condition
(Won)" Gubab (a + + +
(49)
Recalling the recursion relation for C,
m m m²
EC, "
C.+1 = C. - -1EC - 7,C +
(ub² C
1 C N/C 7, 1/ m
m
ΣC, + "
((a +
(50)
m²
(I -I- 11 + anl
ΣC,E + ETREC]
(51)
m
Defining that C,,t u C₁ we note C++1 & (I- diag(A))² + n²a Idiag (10++AAT)) Using the fact that L: cj X, we find
2AT))
L TT (I-- diag (2))² + 11²a
[diag (X + AA ))²
(x2)+
(52)
which proves the desired bound.
G POWER LAW SCALINGS IN SMALL LEARNING RATE LIMIT
By either taking small learning rate or a large batch size, the test loss dynamics reduce to the test loss obtained from gradient descent the population loss. In this section, we consider the small learning rate limit 17 + O, where the average test loss follows
(L.)
A (1 - m7ax))
(53)
18

---

Published as a conference paper at ICLR 2022
14
2.8
0.04
02
3.4
0.4
04
50t
(a) Non-Gaussian Effects on MNIST
(b) Non-Gaussian Effects on CIFAR
Figure F.l: Non-Gaussian effects small on random feature models. (a)-(b) The first 20- dimensions of the summed fourth moment matrix Kjj u (abe) "bb4T) uy are plotted for the Gaussian approximation and the empirical fourth moment. Differences between the Gaussian ap- proximation and true fourth moment matrices on this example are visible, but are only on the order of ~ 5% of the size of the entries Ki4-
Under the assumption that the eigenvalue and target function power spectra both follow power laws A ~ - and of X a the loss can be approximated by an integral over all modes
(I4) = Σ k- k- npph by-
exp (2n In(1 - .nk - a In k.) dk
(54)
j, 0o
exp (-2rm) - a In k:) dl: -
(55)
We identify the function f(k 2nk In k and proceed with Laplace's s method (Bender & Orszag 1999). This consists of Taylor expanding f(k) around its minimum to second order and
computing Gaussian integral
2x exp(=tf(+))dk ~ fexp (-t5(()) - - km7) ~ exp( exp(-tf(("))
(k=)
(56)
We must identify the '- which minimizes f(k)). The interpretation of this value is that it indexes the mode which dominates the error at large time t. The first order condition gives
f(()) 2bmnk b-1
tk = Ks (2bnt/a)
(57)
The second derivative has the form
f"(k") 2b² 2b2ng bb-22 Thus we are left with a scaling of the form
the k- 262n (2bmt/a) - (b+-2//6
1
(2bmt/a) -2/b --1-2
2/h
(58)
(Ix)
1) - ~5
(59)
H PROOF OF THEOREM
Let (vi) 0² and (yL) 0, (y1 bb 0. The gradient descent updates take the following form
A++1 A 7 with
g:
2b
A + v1
(60)
Again, defining C, = (A,AT) we perform the average over each of the v/: vectors to obtain the following recursion relation
C+1==
= AA
A.g. 17 (g.
+ (g.g.)
C - 17 ECC C, m
EC,Σ + ET (C,E), + +
(61)
Again, analyzing Cr ig C,uu find
C+11k = 1 - 11-2 + 1 m +m+11
X C. + Σ 1220
(62)
m
19

---

Published as a conference paper at ICLR 2022
The vector o₁ follows the linear evolution
C++1 = Act + o
(63)
Let o2 X. Writing out the first few steps, we identify a pattern
C1 Ac +
C2 = Acı = A 2 Co + b
C3 Ac2 + = A co + + Ab+b
G A'co +
ΣA" A" b.
(64)
The geometric sum E Ar) can be computed exactly under the assumption that (I - A) is
invertible which holds provided all of A's eigenvalues are less than unity, which necessarily holds provided the system is stable. The geometric sum yields
-1-A) -
Σ n=0
Σ
A" -
(65)
Recalling the definition of b 021)² X and the definition of the average loss (I,) A Ct2 we have
(T+) = 0² + X A' + (I- (I- A
(66)
Recognizing Co 2 gives the desired result.
I PROOF OF THEOREM
We will now prove that if the training p(x) and test distributions p(x) are different and have feature correlation matrices E and Σ respectively, then the average training and test losses have the form
Tr [SC]] Liest Tr [EC.].
Limin
(67)
As before, will assume that there exist weights w" which satisfy y w w/. We start by noticing that the update rule for gradient descent
m De [wr w"]
W W. 7B.
g:
(68)
m
generates the following dynamics for the weight discrepancy correlation (ww, w")(w, - >>)")>;
C.
ECC n² (1-
C++1= = - C, -
(vτ T C,pb)T
(69)
m
m
This formula can be obtained through the simple averaging procedure shown in sian approximation, we can obtain a simplification
B Under the Gaus-
C++1 = (1- (I- --111-1111 -))
ΣC,Σ + ETΣEC,
(70)
We can solve for the evolution of the diagonal and off-diagonal entries in this matrix giving
ug Caug- [A'o2] ug C.u. (1 (1-1A,, 17 ne + 172 (1 1 11 Sede) vwve m
(71)
20

---

Published as a conference paper at ICLR 2022
To calculate the training and test error, have
(Lre)) = (((b(x) We (x) -w=""" =²)²) Xxx((x) = ((W, - (oW--w))>" (w, - -w")) = TrEC, (Lurin) = ((1b(xx) W. - (x)) w² =²))²) Xxx())) = ((w,, w w'j (we - w TTSC, (72)
Note that in the training error formula, since Σ has eigenvectors uy only the diagonal terms u C.u. enter into the formula for Lerains but off-diagonal components ug C.u do enter into the formula for Liest
J EXPERIMENTAL DETAILS
For Figures we use the last two classes of MNIST and CIFAR-I0. We encode the target values as binary E [+1,-1) For Figure we use 6000 random training points drawn from entire MNIST and 10 datasets to calculate the spectrum of the Fisher information matrix. We train with SGD on these training data, using one-hot label vectors for each training example and plot the error on the test set. We train our models all experimental results the supplementary materials. To match our theory, we use fixed learning rate SGD, Both evaluation of the infinite width kernels and training were performed with the Neural Tangents API (Novak et al. 2020).
Google Colab GPU and include code to reproduce
21


Reviews:
[{"review_id": "wchgphMhsdS", "reviewer": "Reviewer_My4V", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I find the following pros and cons of the paper.\n\nPros: \n\n- The paper is written clearly. The problem setup is described rigorously and cleanly, and the content of the paper is well-organized.\n\n- The results of the paper are solid. The theoretical results are correct, and the experiments are also reasonable.\n\nCons:\n\n- 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.\n\n- 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.\n\n- 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.\n\n\n[1] Cao and Gu. \"Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks.\" NeurIPS 2019\n\n[2] Bartlett et al. \"Benign overfitting in linear regression.\", PNAS, 2020\n\n[3] Tsigler and Bartlett. \"Benign overfitting in ridge regression.\" arXiv:2009.14286, 2020\n\n[4] Chatterji and Long. \"Finite-sample analysis of interpolating linear classifiers in the overparameterized regime.\" JMLR, 2021\n\n[5] Zou et al. \"Benign Overfitting of Constant-Stepsize SGD for Linear Regression\", COLT 2021\n\n[6] Cao et al. \"Risk bounds for over-parameterized maximum margin classification on sub-gaussian mixtures.\" NeurIPS 2021\n", "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.", "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.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "oYjX-Fd7yT0", "reviewer": "Reviewer_Ack9", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nI 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.\nMissing references:\n\n(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;\n\n(ii) Ref. [2] extends the above work to include target distributions generated from a pre-trained teacher with a static feature map.\n\nThese 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).\n\nThe 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.\n\nMinor 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.\n\nFigure 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? \n\nIn 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.\n\nThe 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…).\n\nMinor 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.\n\nSome typos are present both in the text and in the equations. For instance:\n-  The left-hand side of eq. (4) is missing $y_{\\perp}(x)$;\n-  Right below eq. (6) the average brackets are missing: $L_t$ should be $\\langle L_t \\rangle$;\n-  The left-hand side of eq. (67), Appendix I should be ${\\bf w}_{t+1}$;\n-  Appendix I ends with an incomplete sentence.\n\nReferences:\n\n[1] arXiv:2006.14709\n\n[2] arXiv:2102.08127\n\n[3] arXiv:1811.03600\n\n[4] arXiv:2006.06098\n\n[5] arXiv:2103.04902\n\n[6] arXiv:1711.04623", "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.", "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.\n\n*** Update after reading the authors response***\n\nThe 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.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "_1B7LVaPIQe", "reviewer": "Reviewer_AoHd", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The paper analyzes the influence of data structure (e.g., the spectrum of covariance) on the loss dynamics, which is very interesting.\n\nWeaknesses:\n- 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.\n- 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.\n- 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.\n\n\n[1] Which algorithmic choices matter at which batch sizes?   \n[2] Accelerating Stochastic Gradient Descent For Least Squares Regression\n\n", "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.\n", "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. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OpC541ySNLz", "reviewer": "Reviewer_FVgD", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nThe paper is generally well written, although exposition in the main text or in the appendix could be made a bit more pedagogical. \nThe 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. \n\n\nSome additional comments:\n\n- Use \\lVert \\lambda \\rVert  instead of |\\lambda| at the top of pag. 4 for consistency. \n\n- 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 \nreport.\n\n- Theorem 3.2: make explicit t depence  in c_k\n- Theorem 3.4: this is assuming Gaussian features as in Theorem 3.1, right? If so, should be clearly stated. \n\n- Theorem 3.5: Why these expressions don't involve 4-th moments? Is some Gaussianity assumption made also here?\n- Theorem 3.5: make t dependence explicit in Ltrain and Ltest\n- In the MNIST and CIFAR tasks, the regression targets are 0s and 1s?\n - Related Work: \"Their predicted exponent\nin 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\nsaddle point approximations.\n\n- Appendix A: \"Since the basis of kernel eigenfunctions\n(including the zero eigenvalue functions) is complete over the space of square integrable functions.\" should end with a comma instead?\n- 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. \n", "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.\nThe 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.\nFor 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 \nGaussian case seems to match very well the experiments on real world datasets.\nThe formalism covers both the online setting and the setting of multiple passes over a fixed training set. ", "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.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

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)