summaries: 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. 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.  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.  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.
abstractText: The generalization performance of a machine learning algorithm such as a neural network depends in an intricate 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) 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 we 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 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.
references.authors: Andreas Anastasiou Krishnakumar Balasubramanian Murat A. Erdogdu
references.title: Normal approximation for stochastic gradient descent via non-asymptotic rates of martingale clt
references.venue: Proceedings of the Thirty-Second Conference on Learning Theory,
references.authors: Yasaman Bahri Ethan Dyer Jared Kaplan Jaehoon Lee Utkarsh Sharma
references.title: Explaining neural scaling laws
references.venue: arXiv preprint arXiv:2102.06701,
references.authors: Peter L Bartlett Philip M Long Gábor Lugosi Alexander Tsigler
references.title: Benign overfitting in linear regression
references.venue: Proceedings of the National Academy of Sciences,
references.authors: Carl Bender Steven Orszag
references.title: Advanced Mathematical Methods for Scientists and Engineers: Asymptotic Methods and Perturbation Theory, volume 1
references.authors: Raphaël Berthier Francis Bach Pierre Gaillard
references.title: Tight nonparametric convergence rates for stochastic gradient descent under the noiseless linear model, 2020
references.authors: M Biehl P Riegler
references.title: On-line learning with a perceptron
references.venue: Europhysics Letters (EPL),
references.authors: Blake Bordelon Abdulkadir Canatar Cengiz Pehlevan
references.title: Spectrum dependent learning curves in kernel regression and wide neural networks
references.venue: In International Conference on Machine Learning,
references.authors: Abdulkadir Canatar B. Bordelon C. Pehlevan
references.title: Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks
references.venue: Nature Communications,
references.authors: Yuan Cao Quanquan Gu
references.title: Generalization bounds of stochastic gradient descent for wide and deep neural networks
references.venue: Advances in Neural Information Processing Systems,
references.authors: Yuan Cao Quanquan Gu Misha Belkin
references.title: Risk bounds for over-parameterized maximum margin classification on sub-gaussian mixtures
references.venue: In Thirty-Fifth Conference on Neural Information Processing Systems,
references.authors: Niladri S Chatterji Philip M Long
references.title: Finite-sample analysis of interpolating linear classifiers in the overparameterized regime
references.venue: Journal of Machine Learning Research,
references.authors: K.L. Chung
references.title: On a Stochastic Approximation Method
references.venue: The Annals of Mathematical Statistics,
references.authors: Aymeric Dieuleveut Francis Bach
references.title: Nonparametric stochastic approximation with large stepsizes
references.venue: The Annals of Statistics,
references.authors: Aymeric Dieuleveut Nicolas Flammarion Francis Bach. Harder
references.title: better, faster, stronger convergence rates for least-squares regression
references.venue: Journal of Machine Learning Research,
references.authors: John C. Duchi Feng Ruan
references.title: Asymptotic optimality in stochastic optimization
references.venue: The Annals of Statistics,
references.authors: A. Engel C. Van den Broeck
references.title: Statistical Mechanics of Learning
references.authors: Simon Fischer Ingo Steinwart
references.title: Sobolev norm learning rates for regularized least-squares algorithms
references.venue: Journal of Machine Learning Research,
references.authors: Nicolas Flammarion Francis Bach
references.title: From averaging to acceleration, there is only a step-size
references.venue: In Conference on Learning Theory,
references.authors: Nicolas Flammarion Francis Bach
references.title: Stochastic composite least-squares regression with convergence rate o(1/n)
references.venue: Proceedings of the 2017 Conference on Learning Theory,
references.authors: Sebastian Goldt Madhu Advani Andrew M Saxe Florent Krzakala Lenka Zdeborová
references.title: Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup
references.venue: Advances in Neural Information Processing Systems,
references.authors: Sebastian Goldt Marc Mézard Florent Krzakala Lenka Zdeborová
references.title: Modeling the influence of data structure on learning in neural networks: The hidden manifold model
references.venue: Phys. Rev. X,
references.authors: Sebastian Goldt Bruno Loureiro Galen Reeves Florent Krzakala Marc Mézard Lenka Zdeborová
references.title: The gaussian equivalence of generative models for learning with shallow neural networks
references.venue: Proceedings of Machine Learning Research,
references.authors: Mert Gurbuzbalaban Umut Simsekli Lingjiong Zhu
references.title: The heavy-tail phenomenon in sgd
references.venue: In International Conference on Machine Learning,
references.authors: Tom M. Heskes Bert Kappen
references.title: Learning processes in neural networks
references.venue: Phys. Rev. A,
references.authors: Joel Hestness Sharan Narang Newsha Ardalani Gregory Diamos Heewoo Jun Hassan Kianinejad Md. Mostofa Ali Patwary Yang Yang Yanqi Zhou
references.title: Deep learning scaling is predictable
references.authors: Arthur Jacot Berfin Simsek Francesco Spadaro Clément Hongler Franck Gabriel
references.title: Kernel alignment risk estimator: Risk prediction from training data
references.venue: NeurIPS,
references.authors: Prateek Jain Sham M Kakade Rahul Kidambi Praneeth Netrapalli Aaron Sidford
references.title: Accelerating stochastic gradient descent for least squares regression
references.venue: In Conference On Learning Theory,
references.authors: Yann LeCun Ido Kanter Sara A. Solla
references.title: Eigenvalues of covariance matrices: Application to neuralnetwork learning
references.venue: Phys. Rev. Lett.,
references.authors: Jaehoon Lee Lechao Xiao Samuel S Schoenholz Yasaman Bahri Roman Novak Jascha SohlDickstein Jeffrey Pennington
references.title: Wide neural networks of any depth evolve as linear models under gradient descent
references.venue: Journal of Statistical Mechanics: Theory and Experiment,
references.authors: Bruno Loureiro Cédric Gerbelot Hugo Cui Sebastian Goldt Florent Krzakala Marc Mézard Lenka Zdeborová
references.title: Capturing the learning curves of generic features maps for realistic data sets with a teacher-student model, 2021
references.authors: Siyuan Ma Raef Bassily Mikhail Belkin
references.title: The power of interpolation: Understanding the effectiveness of sgd in modern over-parametrized learning
references.venue: In ICML,
references.authors: C. Mace A. Coolen
references.title: Statistical mechanical analysis of the dynamics of learning in perceptrons
references.venue: Statistics and Computing,
references.authors: Song Mei Andrea Montanari
references.title: The generalization error of random features regression: Precise asymptotics and double descent curve, 2020
references.authors: Francesca Mignacco Florent Krzakala Pierfrancesco Urbani Lenka Zdeborová
references.title: Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification, 2020
references.authors: Francesca Mignacco Pierfrancesco Urbani Lenka Zdeborová
references.title: Stochasticity helps to navigate rough landscapes: comparing gradient-descent-based algorithms in the phase retrieval
references.authors: Preetum Nakkiran Behnam Neyshabur Hanie Sedghi
references.title: The deep bootstrap framework: Good online learners are good offline generalizers
references.venue: In International Conference on Learning Representations,
references.authors: Roman Novak Lechao Xiao Jiri Hron Jaehoon Lee Alexander A. Alemi Jascha Sohl-Dickstein Samuel S. Schoenholz
references.title: Neural tangents: Fast and easy infinite neural networks in python
references.venue: In International Conference on Learning Representations,
references.authors: Loucas Pillaud-Vivien Alessandro Rudi Francis Bach
references.title: Statistical optimality of stochastic gradient descent on hard learning problems through multiple passes, 2018
references.authors: B. Polyak A. Juditsky
references.title: Acceleration of stochastic approximation by averaging
references.venue: Siam Journal on Control and Optimization,
references.authors: Phil Pope Chen Zhu Ahmed Abdelkader Micah Goldblum Tom Goldstein
references.title: The intrinsic dimension of images and its impact on learning
references.venue: In International Conference on Learning Representations,
references.authors: Ali Rahimi Benjamin Recht
references.title: Random features for large-scale kernel machines
references.venue: Advances in Neural Information Processing Systems,
references.authors: Carl Edward Rasmussen Christopher K.I. Williams
references.title: Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning)
references.authors: Herbert Robbins Sutton Monro
references.title: A Stochastic Approximation Method
references.venue: The Annals of Mathematical Statistics,
references.authors: David Ruppert
references.title: Efficient estimations from a slowly convergent robbins-monro process
references.authors: David Saad Sara Solla
references.title: Dynamics of on-line gradient descent learning for multilayer neural networks
references.authors: Christopher J Shallue Jaehoon Lee Joseph Antognini Jascha Sohl-Dickstein Roy Frostig George E Dahl
references.title: Measuring the effects of data parallelism on neural network training
references.venue: arXiv preprint arXiv:1811.03600,
references.authors: Alexander Shapiro
references.title: Asymptotic Properties of Statistical Estimators in Stochastic Programming
references.venue: The Annals of Statistics,
references.authors: Stefano Spigler Mario Geiger Matthieu Wyart
references.title: Asymptotic learning curves of kernel methods: empirical data versus teacher–student paradigm
references.venue: Journal of Statistical Mechanics: Theory and Experiment,
references.authors: Alexander Tsigler Peter L Bartlett
references.title: Benign overfitting in ridge regression
references.venue: arXiv preprint arXiv:2009.14286,
references.authors: Aditya Varre Loucas Pillaud-Vivien Nicolas
references.title: Flammarion. Last iterate convergence of sgd for least-squares in the interpolation regime, 2021
references.authors: Maksim Velikanov Dmitry Yarotsky
references.title: Universal scaling laws in the gradient descent training of neural networks, 2021
references.authors: Justin Werfel Xiaohui Xie H. Seung
references.title: Learning curves for stochastic gradient descent in linear feedforward networks
references.venue: Advances in Neural Information Processing Systems,
references.authors: Yiming Ying Massimiliano Pontil
references.title: Online gradient descent learning algorithms
references.venue: Found. Comput. Math.,
references.authors: Lu Yu Krishnakumar Balasubramanian Stanislav Volgushev Murat A. Erdogdu
references.title: An analysis of constant step size sgd in the non-convex regime: Asymptotic normality and bias, 2020
references.authors: Guodong Zhang Lala Li Zachary Nado James Martens Sushant Sachdeva George Dahl Chris Shallue Roger B Grosse
references.title: Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model
references.venue: Advances in neural information processing systems,
references.authors: Difan Zou Jingfeng Wu Vladimir Braverman Quanquan Gu Sham M Kakade
references.title: Benign overfitting of constant-stepsize sgd for linear regression
references.venue: arXiv preprint arXiv:2103.12692,
references.authors: ∫ K(x x)p(x)dx
references.title: Mercer’s theorem (Rasmussen & Williams, 2005) guarantees the existence of a set of orthonormal eigenfunctions. Since ψ(x) spans an N dimensional 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
sections.text: The generalization performance of a machine learning algorithm such as a neural network depends in an intricate 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) 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 we 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 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.
sections.heading: 1 INTRODUCTION
sections.text: Understanding the dynamics of SGD on realistic learning problems is fundamental to learning theory. Due to the challenge of modeling the structure of realistic data, theoretical studies of generalization often attempt to derive data-agnostic generalization bounds or study the typical performance of the algorithm on high-dimensional, factorized data distributions (Engel & Van den Broeck, 2001). Realistic datasets, however, often lie on low dimensional structures embedded in high dimensional ambient spaces (Pope et al., 2021). For example, MNIST and CIFAR-10 lie on surfaces with intrinsic 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 choice of 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 ·ψ(x) for nonlinear feature map ψ trained with MSE loss. We express test loss dynamics in terms of the induced second and fourth moments of ψ. Under a regularity condition on the fourth moments, we show that the test error can be accurately predicted in terms of second moments alone. We demonstrate the accuracy of our theory on random feature models and wide neural networks trained on MNIST and CIFAR-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→∞). To explore the computational advantages SGD compared to standard full batch gradient descent we analyze the loss achieved at a fixed compute budget C = tm for different minibatch size m and number of steps t, 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 a predictive account of training dynamics in wide neural networks.
sections.heading: 1.1 OUR CONTRIBUTIONS
sections.text: The novel contributions of this work are described below.
  • We calculate the exact expected test error for SGD on MSE loss for arbitrary feature structure in terms of second and fourth moments as we discuss in Section 4.2. We identify how the structured gradient noise induced by sampling alters the loss curve compared to vanilla GD.
  • For Gaussian features (or those with regular fourth moments), we compute the test error in Section 4.1. This theory is 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.
sections.heading: 2 RELATED WORK
sections.text: The analysis of 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 a noisy problem. Many more works have examined a similar setting, identifying how averaged and accelerated versions of SGD perform asymptotically when the target function is noisy (not a deterministic function of the input) (Flammarion & Bach, 2015; 2017; Shapiro, 1989; Robbins & Monro, 1951; Chung, 1954; Duchi & Ruan, 2021; Yu et al., 2020; Anastasiou et al., 2019; Gurbuzbalaban et al., 2021).
  Recent studies have also analyzed the asymptotics of noise-free MSE problems with arbitrary feature structure to see what stochasticity arises from sampling. Prior works have found exponential 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 overparameterized (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 problem (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., 2021). These works provide bounds of the formO(t−β) 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; Biehl & 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 a Gaussian covariates model (Goldt et al., 2020; 2021). In the continuous time limit considered in these works, SGD converges to gradient flow on the population loss, where fluctuations due to 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; Tsigler & Bartlett, 2020; Bordelon et al., 2020; Canatar et al., 2020), hinge loss (Chatterji & Long, 2021; Cao et al., 2021; Cao & Gu, 2019) and general convex loss functions (Loureiro et al., 2021) 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 m. Ma et al. (2018) study the tradeoff between taking many steps of SGD at small m and taking a small number of steps at large m. After a 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 a noisy quadratic model. Since they stipulate constant gradient noise 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.
sections.heading: 3 PROBLEM DEFINITION AND SETUP
sections.text: We study stochastic gradient descent on a linear model with parameters w and feature map ψ(x) ∈ RN (with N possibly infinite). Some interesting examples of linear models are random feature models, where ψ(x) = φ(Gx) for random matrixG and point-wise nonlinearity φ (Rahimi & Recht, 2008; Mei & Montanari, 2020). Another interesting linearized setting is wide neural networks with neural tangent kernel (NTK) parameterization (Jacot et al., 2020; Lee et al., 2020). Here the features are parameter gradients of the neural network function ψ(x) = ∇θf(x,θ)|θ0 at initialization. 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) = 〈 (w ·ψ(x)− y(x))2 〉 x∼p(x) , (1)
  where x are input data vectors associated with a probability distribution p(x), ψ(x) is a nonlinear feature map and y(x) is a target function which we can evaluate on training samples. We assume that the target function is square integrable 〈 y(x)2 〉 x <∞ 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 eigendecomposition of the feature correlation matrix and the target function
  Σ = 〈 ψ(x)ψ(x)> 〉 x = N∑ k=1 λkuku > k , y(x) = ∑ k vku > k ψ(x) + y⊥(x), (2)
  where 〈y⊥(x)ψ(x)〉 = 0. We justify this decomposition of y(x) in the Appendix A using an eigendecomposition and show that it is general for target functions and features with finite 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∗ ·ψ + y⊥. At each time step t, the weights are updated by taking a stochastic gradient step on a fresh mini-batch of m examples
  wt+1 = wt − η
  m m∑ µ=1 ψt,µ (wt ·ψt,µ − yt,µ) , (3)
  where each of the vectors ψt,µ are sampled independently and yt,µ = w∗ ·ψt,µ. The learning rate η controls the gradient descent step size while the batch size m gives a empirical estimate of the gradient at timestep t. At each timestep, the test-loss, or generalization error, has the form
  Lt = 〈 (wt ·ψ(x)−w∗ ·ψ(x)− y⊥(x))2 〉 x = (wt −w∗)>Σ(wt −w∗) + 〈 y⊥(x) 2 〉 , (4) which quantifies exactly the test error of the vector wt. Note, however, that Lt is a random variable since wt depends on the precise history of sampled feature vectors Dt = {ψt,µ}. Our theory, which generalizes the recursive method of (Werfel et al., 2004) allows us to compute the expected test loss by averaging over all possible sequences to obtain 〈Lt〉Dt . 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 a separate test set which we address in Section 4.4.
  In summary, we will develop a theory that predicts the expected test loss 〈Lt〉Dt averaged over training sample sequences Dt in terms of the quantities {λk, vk, 〈 y⊥(x) 2 〉 x }. 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 ψ, and target functions y(x).
sections.heading: 4 ANALYTIC FORMULAE FOR LEARNING CURVES
sections.heading: 4.1 LEARNABLE AND NOISE FREE PROBLEMS
sections.text: 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 is a reasonable assumption, especially when applying our theory to real datasets by fitting an atomic measure on P points 1P ∑ µ δ(x− xµ). We will further assume that the induced feature distribution is Gaussian so that all moments of ψ can be written in terms of the covariance Σ. We will remove these assumptions in later sections.
  Theorem 4.1. Suppose the features ψ follow a Gaussian distribution ψ ∼ N (0,Σ) and the target function is learnable in these features y = w∗ ·ψ. After t steps of SGD with minibatch size m and learning rate η, the expected (over possible sample sequences Dt) test loss 〈Lt〉Dt has the form
  〈Lt〉Dt = λ >Atv2 , A = (I− η diag(λ))2 + η
  2
  m diag
  ( λ2 ) + η2
  m λλ> (5)
  where λ is a vector containing the eigenvalues of Σ and v2 is a vector containing elements (v2)k = v2k = (uk ·w∗)2 for eigenvectors uk of Σ. The function diag(·) constructs a diagonal matrix with the argument vector placed along the diagonal.
  Proof. See Appendix B for the full derivation. We will provide a brief sketch of the proof here. The strategy of the proof relies on the fact that 〈Lt〉 = Tr Σ Ct where Ct = 〈 (wt −w∗) (wt −w∗)> 〉 Dt
  . We derive the following recursion relation for this error matrix
  Ct+1 = (I− ηΣ)Ct(I− ηΣ) + η2
  m [ΣCtΣ + ΣTr (ΣCt)] (6)
  The loss only depends on ck,t = u>k Ctuk. Solving the recurrence, ct = A tv2 and using 〈Lt〉 =∑
  k λku > k Ctuk = ∑ k ck,tλk = λ >Atv2, we obtain the desired result.
  Below we provide some immediate interpretations of this result.
  • The matrix A contains two components; a matrix (I− η diag(λ))2 which represents the timeevolution of the loss under average gradient updates. The remaining matrix η 2
  m
  ( diag(λ2) + λλ> ) arises due to fluctuations in the gradients, a consequence of the stochastic sampling process.
  • The test loss obtained when training directly on the population loss can be obtained by taking the minibatch size m→∞. In this case, A→ (I− η diag(λ))2 and one obtains the population loss Lpopt = ∑ k v 2 kλk(1 − ηλk)2t. This population loss can also be obtained by considering small
  learning rates, i.e. the η → 0 limit, where A = (I− η diag(λ))2 +O(η2). • For general λ and η2/m > 0, A is non-diagonal, indicating that the components {u1, ...,uk} are
  not learned independently as t increases like for Lpopt , but rather interact during learning due to non-trivial coupling across eigenmodes at large η2/m. This is unlike offline theory for learning in feature spaces (kernel regression), (Bordelon et al., 2020; Canatar et al., 2020), this observation of mixing across covariance eigenspaces agrees with a recent analysis of SGD, which introduced recursively defined “mixing terms” that couple each mode’s evolution (Varre et al., 2021).
  • Though increasing m always improves generalization at fixed time t (proof given in Appendix D), learning with a fixed compute budget (number of gradient evaluations) C = tm, can favor smaller batch sizes. We provide an example of this in the next sections and Figure 1 (d)-(f).
  • The lower bound 〈Lt〉 ≥ λ>v2 [ (1− η)2 + η 2 m |λ| 2 ]t can be used to find necessary stability
  conditions on m, η. This bound implies that 〈Lt〉 will diverge if m < η2−η |λ| 2. The learning rate must be sufficiently small and the batch size sufficiently large to guarantee convergence. This stability condition depends on the features through |λ|2 = ∑ k λ 2 k. One can derive heuristic optimal batch sizes and optimal learning rates through this lower bound. See Figure 2 and Appendix C.
sections.heading: 4.1.1 SPECIAL CASE 1: UNSTRUCTURED ISOTROPIC FEATURES
sections.text: This special case was previously analyzed by Werfel et al. (2004) which takes Σ = I ∈ RN×N and m = 1. We extend their result for arbitrary m, giving the following learning curve
  〈Lt〉Dt = ( (1− η)2 + 1 +N m η2 )t ||w∗||2 , 〈L∗t 〉Dt = ( 1− m m+N + 1 )t ||w∗||2, (7)
  where the second expression has optimal η. First, we note the strong dependence on the ambient dimension N : as N m, learning happens at a rate 〈Lt〉 ∼ e−tm/N . 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 at any learning rate η is much slower than the convergence of the m → ∞ limit, Lpop = (1− η)2t||w∗||2 which does not suffer from a dimensionality dependence due to gradient noise. Lastly, for a fixed compute budget C = tm, the optimal batch size is m∗ = 1; see Figure 1 (d). This can be shown by differentiating 〈 LC/m 〉 with respect to m (see Appendix E). In Figure 1 (a) we show theoretical and simulated learning curves for this model for varying values of N at the optimal learning rate and in Figure 1 (d), we show the loss as a function of minibatch size for a fixed compute budget C = tm = 100. While fixed C represents fixed sample complexity, we stress that it may not represent wall-clock run time when data parallelism is available (Shallue et al., 2018).
sections.heading: 4.1.2 SPECIAL CASE 2: POWER LAWS AND EFFECTIVE DIMENSIONALITY
sections.text: Realistic datasets such as natural images or audio tend to exhibit nontrivial correlation structure, which often results in power-law spectra when the data is projected into a feature space, such as a randomly intialized neural network (Spigler et al., 2020; Canatar et al., 2020; Bahri et al., 2021). In the η 2
  m 1 limit, if the feature spectra and task specra follow power laws, λk ∼ k −b and
  λkv 2 k ∼ k−a with a, b > 1, then Theorem 4.1 implies that generalization error also falls with a power law: 〈Lt〉 ∼ Ct−β , β = a−1b where C is a constant. See Appendix G for a derivation with 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 feature correlation structure (Berthier et al., 2020; Dieuleveut et al., 2016; Velikanov & Yarotsky, 2021; Varre et 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 Section 5.
  (1−η)2 |λ| 2 (red). (c) The loss at fixed
  compute C = 150, a = 2, b = 0.85, optimal batchsize m for each η shown in dashed black. For sufficiently small η, the optimal batchsize is m = 1. For large η, it is better to trade off update steps for denoised gradients resulting in m∗ > 1.
  We show an example of such a power law scaling with synthetic features in Figure 1 (b). Since the total variance approaches a finite value as N →∞, the learning curves are relatively insensitive to N , and are rather sensitive to the eigenspectrum through terms like |λ|2 and 1>λ, etc. In Figure 1 (c), we see that the scaling of the loss is more similar to the power law setting than the isotropic features setting in a 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 = tm is fixed (Figure 1 (e) and (f)). In Appendix C.1, we heuristically argue that the optimal batch size for power law features should scale as, m∗ ≈ 1(2b−1) . Figure 2 tests this result.
  We provide further evidence of the existence of power law structure on realistic data in Figure 3 (a)-(c), where we provide spectra and test loss learning curves for MNIST and CIFAR-10 on ReLU random features. The eigenvalues λk ∼ k−b and the task power tail sums ∑∞ n=k λnv 2 n ∼ k−a+1 both follow power laws, generating power law test loss curves. These learning curves are contrasted with isotropically distributed data in R784 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).
sections.heading: 4.2 ARBITRARY INDUCED FEATURE DISTRIBUTIONS: THE GENERAL SOLUTION
sections.text: The result in the previous section was proven exactly in the case of Gaussian vectors (see Appendix B). For arbitrary (possibly non-Gaussian) distributions, we obtain a slightly more involved result (see Appendix F). Theorem 4.2. Let ψ(x) ∈ RN be an arbitrary feature map with covariance matrix Σ =∑ k λkuku > k . After diagonalizing the features φk(x) = u > k ψ(x), introduce the fourth moment
  tensor κ4ijkl = 〈φiφjφkφl〉. The expected loss is exactly 〈Lt〉 = ∑ k λkck(λ,κ,v, t).
  We provide an exact formula for ck in the Appendix F We see that the test loss dynamics depends only on the second and fourth moments of the features through quantities λk and κijk` respectively. We recover the Gaussian result as a special case when κijkl is a simple weighted sum of these three products of Kronecker tensors κGaussijkl = λiλjδikδjl + λiλkδijδkl + λiλjδilδjk. As an alternative to the above closed form expression for 〈Lt〉, a 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 a regularity condition, similar to those assumed in other recent works (Jain et al., 2018; Berthier et al., 2020; Varre et al., 2021), on the fourth moment structure of the features allows derivation of an upper bound which is qualitatively similar to the Gaussian theory.
  Theorem 4.3. If the fourth moments satisfy 〈 ψψ>Gψψ> 〉 (α+ 1)ΣGΣ + αΣTrΣG for any positive-semidefiniteG, then
  Lt ≤ λ>Atv2 , A = (I− η diag(λ))2 + αη2
  m
  [ diag(λ2) + λλ> ] . (8)
  We provide this proof in Appendix F.1. We note that the assumed bound on the fourth moments is tight for Gaussian features with α = 1, recovering our previous theory. Thus, if this condition on the fourth moments is 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/α.
  The question remains whether the Gaussian approximation will provide an accurate model on realistic data. We do not provide a 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.1, we show that the fourth moment matrix for a ReLU random feature model and its projection along the eigenbasis of the feature covariance is accurately approximated by the equivalent Gaussian model.
sections.heading: 4.3 UNLEARNABLE OR NOISE CORRUPTED PROBLEMS
sections.text: In general, the target function y(x) may depend on features which cannot be expressed as linear combinations of featuresψ(x), y(x) = w∗ ·ψ(x)+y⊥(x). Let 〈 y⊥(x) 2 〉 x
  = σ2. Note that y⊥ does not have to be a deterministic function of x, but can also be a stochastic process which is uncorrelated with ψ(x).
  Theorem 4.4. For a target function with unlearnable variance 〈 y2⊥ 〉
  = σ2 trained on Gaussian ψ, the expected test loss has the form
  〈Lt〉 = λ>Atv2 + 1
  m η2σ2λ>(I−A)−1(I−At)λ (9)
  which has an asymptotic, irreducible error 〈L∞〉 = 1mη 2σ2λ>(I−A)−1λ as t→∞.
  See Appendix H for the proof. The convergence to the asymptotic error takes the form 〈Lt − L∞〉 = λ>At ( v2 − 1mη 2σ2(I−A)−1λ ) . We note that this quantity is not necessarily monotonic in t and can exhibit local maxima for sufficiently large σ2, as in Figure 3 (f).
sections.heading: 4.4 TEST/TRAIN SPLITS
sections.text: Rather than interpreting our theory as a description of the average test loss during SGD in a one-pass setting, where data points are sampled from the a distribution at each step of SGD, our theory can 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 4.5. Let p̂(x) = 1M ∑ µ δ(x− xµ) be the empirical distribution on the M training data
  points and let Σ̂ = 〈 ψ(x)ψ(x)> 〉 x∼p̂(x) = ∑ k λ̂kuku > k be the feature correlation matrix on this training set. Let p(x) be the test distribution Σ its corresponding feature correlation. Then we have
  〈Ltrain,t〉 = Tr [ Σ̂Ct ] , 〈Ltest,t〉 = Tr [ΣCt]
  Ct+1 = (I− ηΣ̂)Ct(I− ηΣ̂) + η2
  m
  [〈 ψ(x)ψ(x)>Ctψ(x)ψ(x) >〉 x∼p̂(x) − Σ̂CtΣ̂ ] (10)
  We provide the proof of this theorem in Appendix I. The interpretation of this result is that it provides the expected training and test loss if, at each step of SGD,m points from the training set {x1, ...,xM} are sampled uniformly with replacement and used to calculate a 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 Ct can again be more easily solved under a Gaussian approximation 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 to zero, the test errors saturate at a M -dependent final loss. This result can also allow one to predict errors on other test distributions.
sections.heading: 5 COMPARING NEURAL NETWORK FEATURE MAPS
sections.text: We can utilize our theory to compare how wide neural networks of different depths generalize when trained with SGD on a real dataset. With a certain parameterization, large width NNs are approximately 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 ψ(x) = ∇θf(x,θ).
  In Figure 5, we show the Neural Tangent Kernel (NTK) eigenspectra and task-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 estimate the power law exponents for the kernel and task spectra λk and v2k. Across architectures, the task spectra v2k are highly similar, but that the kernel eigenvalues λk 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 λk ∼ k−b and the task power v2k ∼ k−a, we compared the true test loss dynamics (color) for a width-500 neural network model with the predicted power-law scalings β = a−1b from the fit exponents a, b. The predicted scalings from NTK regression accurately describe trained width-500 networks. On CIFAR-10, we compare the scalings of the CNN model and a standard MLP and find that the CNN obtains a better exponent due to its faster decaying tail sum ∑∞ n=k λnv 2 n. We stress that the exponents β were estimated from our one-pass theory, but were utilized experiments on a finite training set. This approximate and convenient version of our theory is quite accurate across these varying models, in line with recent conjectures about early training dynamics (Nakkiran et al., 2021).
sections.heading: 6 CONCLUSION
sections.text: 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 〈Lt〉 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 as well as alternative optimization strategies such as SGD with momentum.
sections.heading: A DECOMPOSITION OF THE FEATURES AND TARGET FUNCTION
sections.text: Let y(x) be a square integrable target function with 〈 y(x)2 〉 < ∞. Define the following integral operator TK for kernel K(x,x′) = ψ(x) ·ψ(x′):
  TK [φ](x ′) = ∫ p(x)K(x,x′)φ(x)dx (11)
  We are interested in eigenfunctions of this operator, function φk for which TK [φk] = λkφk. For kernels with finite trace ∫ K(x,x)p(x)dx <∞, Mercer’s theorem (Rasmussen & Williams, 2005) guarantees the existence of a set of orthonormal eigenfunctions. Since ψ(x) spans an N dimensional 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 λ1 > λ2 > ... > λN with λN+` = 0, we obtain the expansion
  y(x) = ∑ k 〈y(x)φk(x)〉x φk(x) = ∑ k≤N vkφk(x) + y⊥(x) , y⊥(x) = ∑ k>N 〈y(x)φk(x)〉φk(x)
  (12)
  Further, we can decompose the feature map in this basis ψ(x) = ∑N k=1 √ λkukφk(x). We recognize through these decompositions the coefficients vk can be computed uniquely as vk = λ −1/2 k u > k 〈ψ(x)y(x)〉. This provides a recipe for determining the necessary spectral quantities for our theory. We see that the feature map’s decomposition above reveals that λk are also the eigenvalues of the feature correlation matrix Σ since
  Σ = 〈 ψ(x)ψ(x)> 〉 = ∑ k` √ λkλ`uku > ` 〈φk(x)φ`(x)〉 = ∑ k λkuku > k . (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) = 1M ∑M µ=1 δ(x− x′). Plugging this into the integral operator gives
  TK [φ](x) = 1
  M ∑ µ K(x,xµ)φ(xµ) (14)
  We see that, restricting the domain to the set of points {x1, ...,xM}, this amounnts to solving a matrix eigenvalue problem 1MKφk = λkφk where K ∈ R
  M×M is the kernel gram matrix with entries Kµν = K(x µ,xν) and φk has entries φk,µ = φk(xµ).
sections.heading: B PROOF OF THEOREM 4.1
sections.text: Let ∆t = wt −w∗ represent the difference between the current and optimal weights and define the correlation matrix for this difference
  Ct = 〈 ∆t∆ > t 〉 Dt−1 . (15)
  Using stochastic gradient descent, wt+1 = wt − ηgt with gradient vector gt = 1m ∑m i=1ψiψ > i ∆t, the Ct matrix satisfies the recursion
  Ct+1 = 〈 (∆t − ηgt)(∆t − ηgt)> 〉 Dt = Ct − η 〈 gt∆ > t 〉 − η 〈 ∆tg > t 〉 + η2 〈 gtg > t 〉 . (16)
  First, note that since ψi are all independently sampled at timestep t, we can break up the average into the fresh batch of m samples and an average over Dt−1
  〈gt∆t〉Dt = 1
  m m∑ i=1 〈 ψiψ > i 〉 ψi 〈 ∆∆>t 〉 Dt−1 = ΣCt. (17)
  The last term requires computation of fourth moments〈 gtg > t 〉 = 1
  m2 ∑ i,j 〈 ψiψ > i 〈 ∆t∆ > t 〉 Dt−1 ψjψ > j 〉 ψi,ψj
  (18)
  = 1
  m2 ∑ i,j 〈 ψiψ > i Ctψjψ > j 〉 ψi,ψj . (19)
  First, consider the case where i = j. Letting ψ = ψi, we need to compute terms of the form∑ k,` Ck,` 〈ψjψkψ`ψn〉 . (20)
  For Gaussian random vectors, we resort to Wick-Isserlis theorem for the fourth moment
  〈ψjψkψlψn〉 = 〈ψjψk〉 〈ψ`ψn〉+ 〈ψjψ`〉 〈ψkψn〉+ 〈ψjψn〉 〈ψ`ψk〉 (21)
  giving 〈 gtg > t 〉 = m+ 1
  m ΣCtΣ +
  1 m Σ Tr (ΣCt) . (22)
  This correlation structure for gt implies that its covariance has the form
  〈Covψ(gt)〉Dt = 1
  m ΣCtΣ +
  1 m ΣTr(ΣCt). (23)
  Using the formula for 〈 gtg > t 〉 , we arrive at the following recursion relation for Ct
  Ct+1 = Ct − ηCtΣ− ηΣCt + η2 m+ 1
  m ΣCtΣ +
  1 m η2Σ Tr (ΣCt) . (24)
  Since we are ultimately interested in the generalization error 〈Lt〉 = 〈 ∆>t Σ∆t 〉 = TrΣCt =∑
  k λku > k Ctuk, it suffices to track the evolution of ct,k = u > k Ctuk
  ct+1,k = ( 1− 2ηλk + η2 m+ 1
  m λ2k
  ) ct,k + 1
  m η2λk ∑ j λjct,j . (25)
  Vectorizing this equation for c generates the following solution
  ct = A tc0 , A = I− 2η diag(λ) +
  m+ 1
  m η2 diag(λ2) +
  η2 m λλ>. (26)
  The coefficient c0,k = v2k = ( u>k w ∗)2. To get the generalization error, we merely compute 〈Lt〉 = λ>at = λ>Atv2 as desired.
sections.heading: C PROOF OF STABILITY CONDITIONS
sections.text: We will first establish the following lower bound on the loss, where without loss of generality we assumed the maximum correlation is 1:
  Lt ≥ λ>v2 [ (1− η)2 + η 2
  m |λ|2
  ]t . (27)
  We will then use this lower bound to provide necessary conditions on the learning rate and batch size for stability of the loss evolution. First, note that the following inequality holds elementwise
  Aλ = [ (I − η diag (λ))2 + η 2
  m diag(λ2) +
  η2 m λ
  ] λ ≥ [ (1− η)2 + η 2
  m |λ|2
  ] λ (28)
  Repeating this inequality t times gives Atλ ≥ [ (1− η)2 + η 2 m |λ| 2 ]t λ. Using the fact that Lt = λ>Atv2 gives the desired inequality. Note that this inequality is very close to the true result for
  isotropic features λ = 1 which gives Lt ∝ [ (1− η)2 + η 2 m (|λ| 2 + 1) ]t . For anisotropic features with small learning rate, this bound becomes less tight. For the loss to converge to zero at large time, the quantity in brackets must necessarily be less than one. This implies the following necessary condition on the batchsize and learning rate
  η < 2m
  m+ |λ|2 ⇐⇒ m > mmin =
  η|λ|2 2− η (29)
  where mmin is the minimal batch size for learning rate η and feature covariance eigenvalues λ.
  C.1 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 a fixed learning rate η before discussing the optimal m when η is chosen optimally. The loss at a fixed compute budget C = tm is lower bounded by
  LC/m ≥ λ>v2 [ (1− η)2 + η 2
  m |λ|2
  ]C/m (30)
  For the purposes of optimization, we introduce x = 1/m and consider optimizing
  f(x) = x ln [A+ x] , A = (1− η)2
  η2|λ|2 (31)
  The first order optimality condition f ′(x) = 0 implies that (A+ x) ln(A+ x) = −x. This equation has solutions for A ∈ (0, 1) giving solutions x ∈ (0, e−1). Thus the predicted heuristic optimal batch size is always greater than Euler’s constant m∗ = 1x∗ > e. As A increases the optimal batch-size m ∗ increases monotonically. We can get an intuition for this optimality criterion by noticing that x ≈ A gives a decent approximation to the solution. Thus we expect optimal batch size to scale roughly as
  m∗ ≈ η 2|λ|2
  (1− η)2 (32)
  which is roughly 2/η the size of the stability bound mmin when η is small.
  C.1.2 HEURISTIC OPTIMAL LEARNING RATE AND BATCH SIZE
  We will now examine what happens when one first optimizes loss bounmd with respect to the the learning rate at any value of m and then subsequently optimizes over the batch size m. We can easily find the η which minimizes the lower bound.
  ∂
  ∂η
  [ (1− η)2 + η 2|λ|2
  m
  ] = 0 =⇒ η∗ = m
  m+ |λ|2 (33)
  Note that this heuristic optimal learning rate is very close to the true optimum in the isotropic data setting ηtrue = mm+N+1 ≈ m m+N . Plugging this back into the loss bound, we find that at fixed compute C = tm, the loss scales like
  LC/m ≥ λ>v2 [ |λ|2
  m+ |λ|2
  ]C/m (34)
  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 is chosen optimally then small batch sizes give the best performance per unit of computation. This corresponds to the hyperparameter choices (η,m) = ( 1 1+|λ|2 , 1 ) .
  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,
  ∂ 〈Lt〉 ∂m = λ> ∂At ∂m v2, (35)
  is always non-positive. The derivative of the t-th power of A can be identified with the chain rule ∂At
  ∂m = ∂A ∂m At−1 + A ∂A ∂m At−2 + A2 ∂A ∂m At−3 + ...+ At−1 ∂A ∂m . (36)
  Note that the matrix ∂A
  ∂m = − η
  2 m2 [ diag(λ2) + λλ> ] (37)
  has all non-positive entries. Thus we find that
  ∂ 〈Lt〉 ∂m = t−1∑ n=0 λ>An ∂A ∂m At−n−1v2. (38)
  Note that since all entries in v2k and A t−n−1 are non-negative, the vector zn = At−n−1v2 has non-negative entries. By the same argument, the vector qn = Anλ is also non-negative in each entry. Therefore, each of the terms in ∂〈Lt〉∂m above must be non-positive
  ∂ 〈Lt〉 ∂m = t−1∑ n=0 z>n ∂A ∂m qn = − η2 m ∑ n ∑ k,` zn,k [ δk,`λ 2 k + λkλ` ] qn,` ≤ 0. (39)
  Thus we find ∂〈Lt〉∂m ≤ 0, implying that optimal 〈Lt〉 is always obtained (possibly non-uniquely) at m→∞.
  E INCREASING m INCREASES THE LOSS AT FIXED C = tm ON ISOTROPIC FEATURES
  Unlike the previous section, which considered fixed t and varyingm, in this section we consider fixing the total number of samples (or gradient evaluations) which we call the compute budget C = tm. For a fixed compute budget C = tm, and unstructured N dimensional Gaussian features and optimal learning rate η∗ = mm+N+1 , we have〈
  LC/m 〉 =
  ( N + 1
  m+N + 1
  )C/m ||w∗||2. (40)
  Taking a derivative with respect to the batch size we get ∂ ∂m log 〈 LC/m 〉 = ∂ ∂m C m log
  ( N + 1
  m+N + 1 ) = C
  m2 log
  ( m+N + 1
  N + 1
  ) +
  C
  m(m+N + 1) > 0. (41)
  This exercise demonstrates that, for the isotropic features, smaller batch-sizes are preferred at a fixed compute budget C. This result does not hold for arbitrary spectra λk. In the general case, optimal minibatch sizes can exist as we show in Figure 1 (e)-(f) for power law and MNIST spectra.
sections.heading: F PROOF OF THEOREM 4.2
sections.text: Let Vec(·) 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 N2 ×N2 two-dimensional matrix. We will show that the expected loss (over Dt) is
  〈Lt〉 = ∑ k λkct,kk , ct = ( G + η2 m Mat(κ) )t Vec(vv>) ∈ RN 2
  (42)
  where [G]ij,kl = δikδjl ( 1− η(λi + λj) + η 2(m−1) m λiλj ) and [v]k = uk ·w∗.
  We rotate all of the feature vectors into the eigenbasis of the covariance, generating diagonalized features φk = u>k ψ and introduce the following fourth moment tensor
  κijkl = 〈φiφjφkφ`〉 . (43) We redefine Ct in the appropriate (rotated) basis by projecting onto the eigenvectors of the covariance
  Ct = U > 〈∆t∆>t 〉U, (44)
  where U = [u1,u2, ...,uN ]. With this definition, C’s dynamics take the form
  Ct+1 = Ct −ΛCt −CtΛ + η2(m− 1)
  m ΛCtΛ +
  〈 φφ>Ctφφ >〉 . (45) The elements of the matrix can be expressed with the fourth moment tensor∑
  k` 〈φiφjφkφ`〉Ck` = ∑ k` κijklCk,`. (46)
  We thus generate the following dynamics for Ctij
  Ct+1ij = ( 1− η(λi + λj) + η2(m− 1) m λiλj ) Ctij + η2 m ∑ kl κijklC t kl. (47)
  Let ct = Vec(Ct), then we have
  ct+1 = ( G0 + η2
  m Mat(κ)
  ) ct , [G0]ij,k` = δikδj` [ 1− η(λi + λj) +
  η2(m− 1) m λiλj
  ] . (48)
  Solving these dynamics for c, recognizing that c0 = Vec(vv>), and computing 〈Lt〉 = TrΣCt =∑ k Ckkλk gives the desired result.
  F.1 PROOF OF THEOREM 4.3
  Suppose that the features satisfy the regularity condition〈 ψψ>Gψψ> 〉 (α+ 1)ΣGΣ + αΣTr (ΣG) (49)
  Recalling the recursion relation for Ct
  Ct+1 = Ct − ηΣCt − ηCtΣ + η2 m2 −m m2 ΣCtΣ + η2 m
  〈 ψψ>Ctψψ >〉 Ct − ηΣCt − ηCtΣ + η2
  m2 −m m2 ΣCtΣ + η2 m [(α+ 1)ΣCtΣ + αΣTrCtΣ] (50)
  = (I− ηΣ)Ct(I− ηΣ) + αη2
  m [ΣCtΣ + ΣTrΣCt] (51) Defining that ck,t = u>k Ctuk, we note ct+1 ≤ ( (I − η diag(λ))2 + η 2α m [ diag(λ2) + λλ> ]) ct. Using the fact that Lt = c>t λ, we find
  Lt ≤ λ> ( (I − η diag(λ))2 + η 2α
  m
  [ diag(λ2) + λλ> ])t v2 (52)
  which proves the desired bound.
sections.heading: G POWER LAW SCALINGS IN SMALL LEARNING RATE LIMIT
sections.text: By either taking a small learning rate η or a large batch size, the test loss dynamics reduce to the test loss obtained from gradient descent on the population loss. In this section, we consider the small learning rate limit η → 0, where the average test loss follows
  〈Lt〉 ∼ ∞∑ k=1 λkv 2 k(1− ηλk)2t. (53)
  Gauss
  0.0 0.2 0.4 0.6 0.8 1.0 4
  0.0 0.2 0.4 0.6 0.8 1.0 | 4 Gauss| 0.00 0.02 0.04
  (a) Non-Gaussian Effects on MNIST
  Gauss
  0.0 0.2 0.4 0.6 0.8 1.0 4
  0.0 0.2 0.4 0.6 0.8 1.0 | 4 Gauss| 0.01 0.02 0.03 0.04
  (b) Non-Gaussian Effects on CIFAR
  Figure F.1: Non-Gaussian effects are small on random feature models. (a)-(b) The first 20-dimensions of the summed fourth moment matrix κ4ij = u > i 〈 ψψ>ψψ> 〉 uj are plotted for the Gaussian approximation and the empirical fourth moment. Differences between the Gaussian approximation and true fourth moment matrices on this example are visible, but are only on the order of ∼ 5% of the size of the entries in κ4.
  Under the assumption that the eigenvalue and target function power spectra both follow power laws λk ∼ k−b and v2kλk ∼ k−a, the loss can be approximated by an integral over all modes k
  〈Lt〉 = ∑ k k−a(1− ηk−b)2t ∼ ∫ ∞ 1 exp ( 2η ln(1− ηk−b)t− a ln k ) dk (54)
  ∼ ∫ ∞ 1 exp ( −2ηηk−bt− a ln k ) dk , η → 0 (55)
  We identify the function f(k) = 2ηk−b + 1t ln k and proceed with Laplace’s method (Bender & Orszag, 1999). This consists of Taylor expanding f(k) around its minimum to second order and computing a Gaussian integral∫
  exp(−tf(k))dk ∼ ∫ exp ( −tf(k∗)− t
  2 f ′′(k∗)(k − k∗)2
  ) ∼ exp(−tf(k∗)) √ 2π
  tf ′′(k∗) .
  (56) We must identify the k∗ which minimizes f(k). The interpretation of this value is that it indexes the mode which dominates the error at a large time t. The first order condition gives
  f ′(k) = −2bηk−b−1 + a tk = 0 =⇒ k∗ = (2bηt/a)1/b . (57)
  The second derivative has the form
  f ′′(k∗) = 2b2ηk−b−2 − 1 tk2 |k∗ = 2b2η (2bηt/a)−(b+2)/b − 1 t (2bηt/a) −2/b ∼ t−1−2/b. (58)
  Thus we are left with a scaling of the form
  〈Lt〉 ∼ exp(−a/b ln t)t1/b ∼ t− a−1 b . (59)
sections.heading: H PROOF OF THEOREM 4.4
sections.text: Let 〈 y2⊥ 〉
  = σ2 and 〈y⊥〉 = 0, 〈y⊥ψ〉 = 0. The gradient descent updates take the following form ∆t+1 = ∆t − ηgt with
  gt = 1
  m m∑ i=1 ψi [ ψ>i ∆t + y⊥,i ] . (60)
  Again, defining Ct = 〈 ∆t∆ > t 〉 we perform the average over each of the ψi vectors to obtain the following recursion relation
  Ct+1 = 〈 ∆t∆ > t 〉 − η 〈 ∆tg > t 〉 − η 〈 gt∆ > t 〉 + η2 〈 gtg > t 〉 = Ct − ηΣCt − ηCtΣ + m+ 1
  m ΣCtΣ +
  1 m ΣTr (CtΣ) + η2σ2Σ. (61)
  Again, analyzing ct,k = u>k Ctuk we find
  ct+1,k = ( 1− 2ηλk + η2 m+ 1
  m λ2k
  ) ct,k + 1
  m ∑ ` λ`ct,` + η 2σ2λk. (62)
  The vector ct follows the linear evolution
  ct+1 = Act + η 2σ2λ. (63)
  Let b = η2σ2λ. Writing out the first few steps, we identify a pattern
  c1 = Ac0 + b
  c2 = Ac1 + b = A 2c0 + Ab+ b
  c3 = Ac2 + b = A 3c0 + A 2b+ Ab+ b
  ...
  ct = A tc0 + ( t−1∑ n=0 An ) b. (64)
  The geometric sum (∑t−1 n=0 A n )
  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(
  t−1∑ n=0 An
  ) = (I−A)−1 ( I−At ) . (65)
  Recalling the definition of b = σ2η2λ and the definition of the average loss 〈Lt〉 = λ>ct, we have 〈Lt〉 = λ>Atc0 + η2σ2λ>(I−A)−1 ( I−At ) λ. (66)
  Recognizing c0 = v2 gives the desired result.
sections.heading: I PROOF OF THEOREM 4.5
sections.text: We will now prove that if the training p̂(x) and test distributions p(x) are different and have feature correlation matrices Σ̂ and Σ respectively, then the average training and test losses have the form
  Ltrain = Tr [ Σ̂Ct ] Ltest = Tr [ΣCt] . (67)
  As before, we will assume that there exist weights w∗ which satisfy y = w∗ ·ψ. We start by noticing that the update rule for gradient descent
  wt = wt − ηgt , gt = 1
  m m∑ µ=1 ψt,µψ > t,µ [wt −w∗] (68)
  generates the following dynamics for the weight discrepancy correlation Ct =〈 (wt −w∗)(wt −w∗)> 〉 Dt .
  Ct+1 = Ct − ηΣ̂Ct − ηCtΣ̂ + η2 (
  1− 1 m
  ) Σ̂CtΣ̂ + η2
  m
  〈 ψψ>Ctψψ >〉 (69) This formula can be obtained through the simple averaging procedure shown in B. Under the Gaussian approximation, we can obtain a simplification
  Ct+1 = (I− ηΣ)Ct(I− ηΣ) + η2
  m [ΣCtΣ + ΣTrΣCt] (70)
  We can solve for the evolution of the diagonal and off-diagonal entries in this matrix giving
  u>k Ctuk = [ Atv2 ] k , u>k Ctu` = ( 1− ηλ̂k − ηλ̂` + η2 ( 1 + 1
  m
  ) λ̂kλ̂` )t vkv` (71)
  To calculate the training and test error, we have 〈Ltest〉 = 〈 (ψ(x) ·wt −ψ(x) ·w∗)2 〉 x∼p(x),Dt = 〈(wt −w∗)Σ (wt −w∗)〉 = TrΣCt.
  〈Ltrain〉 = 〈 (ψ(x) ·wt −ψ(x) ·w∗)2 〉 x∼p̂(x),Dt = 〈 (wt −w∗)Σ̂ (wt −w∗) 〉 = TrΣ̂Ct. (72)
  Note that in the training error formula, since Σ̂ has eigenvectors uk only the diagonal terms u>k Ctuk enter into the formula for Ltrain, but off-diagonal components u>k Ctu` do enter into the formula for Ltest
sections.heading: J EXPERIMENTAL DETAILS
sections.text: For Figures 3, we use the last two classes of MNIST and CIFAR-10. We encode the target values as binary y ∈ {+1,−1}. For Figure 5, we use 6000 random training points drawn from entire MNIST and CIFAR-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 on a Google Colab GPU and include code to reproduce all experimental results in 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).