A Theoretical Analysis of Generalization Dynamics in Neural Networks under Gradient Descent with Weight Decay
Theoretical framework bounds generalization for gradient descent with weight decay, deriving conditions that explain delayed generalization and grokking.
The paper proves convergence of gradient descent with weight decay to a neighborhood of global minimizers of the empirical l2 loss for a broad class of neural networks. It decomposes population error into data, optimization, and prediction variation errors, deriving cellwise and layerwise approximate-homogeneity bounds on prediction variation along the training trajectory. The resulting necessary and sufficient conditions explain layerwise generalization differences and provide a theoretical characterization of grokking.
- Proves gradient descent with weight decay converges near empirical loss minimizers
- Decomposes population error into data, optimization, and prediction variation terms
- Derives sufficient condition characterizing delayed generalization and grokking
Full article170 words · extracted from arxiv.org · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.07755