Learning Curves for SGD on Structured Features

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 mean square loss 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.

arXiv commentCamera Ready for ICLR 2022: https://openreview.net/forum?id=WPI2vbkAl3Q

References in corpus (14)

• The generalization error of random features regression: Precise asymptotics and double descent curve
• Explaining Neural Scaling Laws
• The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning
• From Averaging to Acceleration, There is Only a Step-size
• Accelerating Stochastic Gradient Descent For Least Squares Regression
• The Heavy-Tail Phenomenon in SGD
• Statistical Optimality of Stochastic Gradient Descent on Hard Learning Problems through Multiple Passes
• Harder, Better, Faster, Stronger Convergence Rates for Least-Squares Regression
• The Intrinsic Dimension of Images and Its Impact on Learning
• Kernel Alignment Risk Estimator: Risk Prediction from Training Data
• Tight Nonparametric Convergence Rates for Stochastic Gradient Descent under the Noiseless Linear Model
• An Analysis of Constant Step Size SGD in the Non-convex Regime: Asymptotic Normality and Bias
• Normal Approximation for Stochastic Gradient Descent via Non-Asymptotic Rates of Martingale CLT
• Universal scaling laws in the gradient descent training of neural networks