Understanding Gradient Descent through the Training Jacobian
This provides insights into training dynamics for machine learning researchers, though it is incremental as it builds on existing analysis methods.
The paper analyzes neural network training geometry using the Jacobian of trained parameters with respect to initial values, revealing low-dimensional structure dependent on input data but independent of labels, with singular value spectra showing chaotic, bulk, and stable regions where perturbations in bulk directions have minimal in-distribution effect but affect out-of-distribution performance.
We examine the geometry of neural network training using the Jacobian of trained network parameters with respect to their initial values. Our analysis reveals low-dimensional structure in the training process which is dependent on the input data but largely independent of the labels. We find that the singular value spectrum of the Jacobian matrix consists of three distinctive regions: a "chaotic" region of values orders of magnitude greater than one, a large "bulk" region of values extremely close to one, and a "stable" region of values less than one. Along each bulk direction, the left and right singular vectors are nearly identical, indicating that perturbations to the initialization are carried through training almost unchanged. These perturbations have virtually no effect on the network's output in-distribution, yet do have an effect far out-of-distribution. While the Jacobian applies only locally around a single initialization, we find substantial overlap in bulk subspaces for different random seeds. Our code is available at https://github.com/EleutherAI/training-jacobian