Learning Curves for SGD on Structured Features

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 . 1 INTRODUCTION . 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 . 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 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 . 2 RELATED WORK . 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 . 3 PROBLEM DEFINITION AND SETUP . 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 ) .

Summary: 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.