Learning Curves for SGD on Structured Features

TL;DR

This paper develops an exact theoretical model to analyze how data structure influences the test loss dynamics of SGD in neural networks, revealing the benefits of small batch sizes and accurately predicting performance on real datasets.

Contribution

It introduces an exactly solvable model for SGD on structured features, connecting theory with real neural network training and data covariance structures.

Findings

Gaussian feature model accurately predicts test loss for neural networks.

Small batch sizes are often optimal for fixed compute budgets.

The theory extends to real datasets, predicting training and test error accurately.

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.