How training shapes neural network accuracy over time with weight decay
A Theoretical Analysis of Generalization Dynamics in Neural Networks under Gradient Descent with Weight Decay
Machine Learning
Summary
Understanding why neural networks get better at guessing new data is complicated and depends on the data, the network design, and how the training is done. The authors studied neural networks trained with a common technique called gradient descent combined with weight decay, a method to keep the network simple. They broke down the errors into parts related to the data, the training process, and how much the model’s predictions wiggle around. From this, they found conditions that explain why some layers in a network learn better than others and why improvement can sometimes appear late during training, a phenomenon known as grokking.
generalizationneural networksgradient descentweight decayempirical losspopulation erroroptimization errorprediction variationgrokkingtraining dynamics
Authors
Yuqing Wang, Ioannis G. Kevrekidis, Mikhail Belkin
Abstract
Understanding generalization remains a central challenge in machine learning because it requires jointly considering data, architecture, and training dynamics. In this paper, we develop a theoretical framework that characterizes how these factors jointly shape generalization performance throughout training. More precisely, we study a broad class of neural networks trained under the $\ell^2$ loss by gradient descent (GD) with weight decay, and prove the convergence of GD to a neighbourhood of the global minimizers of the empirical loss. By partitioning the space based on the input data, we then decompose the population error into data error, optimization error, and prediction variation error, and bound them separately. In particular, for the prediction variation error, which measures the oscillations of the learned function, we propose (local) approximate homogeneity and derive explicit cellwise and layerwise bounds for its evolution along the training trajectory. These bounds yield two important implications: a necessary condition of improved generalization explains differences in layerwise generalization behavior; a sufficient condition describes delayed generalization and provides a theoretical characterization of grokking.