- The paper analyzes neural network training geometry by examining the singular value spectrum of the training Jacobian, which links initial and final parameters under gradient descent.
- It identifies three spectral regions
having distinct effects on parameter perturbations during training: chaotic (singular values > 1), bulk (singular values ~ 1), and stable (singular values < 1).
- Experiments show a prevalent and robust bulk region, suggesting that SGD operates within a low-dimensional subspace largely determined by data structure, not initialization or labels.
Understanding Gradient Descent Through the Training Jacobian
The paper provides an analytical perspective on the geometry of neural network training, focusing on the singular value spectrum of the Jacobian matrix. This Jacobian, termed the training Jacobian, describes the relationship between initial and final parameters of a trained neural network, particularly under gradient descent optimization.
The research identifies three distinct spectral regions within the singular value spectrum of the training Jacobian: the chaotic, bulk, and stable regions. The chaotic region is characterized by singular values significantly greater than one, indicating non-convexity in the optimization landscape. The bulk region, where singular values are extremely close to one, suggests that perturbations to initial parameters are carried through training unchanged, revealing an intrinsic geometry largely invariant to the labels and initializations but sensitive to data distribution. Lastly, the stable region has singular values less than one, consistent with traditional convex optimization settings where the optimizer dampens parameter perturbations.
Experimentally, the authors employ small neural networks to compute the training Jacobian using JAX's forward-mode automatic differentiation. Their analysis, conducted on an MLP trained on the UCI digits dataset, reveals that approximately two-thirds of the parameter space falls into the bulk, marginalizing the effect of perturbations on the training data. Further, they demonstrate that these bulk directions have a muted influence on network output when evaluated on in-distribution data compared to far out-of-distribution scenarios, such as inputs comprising random noise.
Results indicate that the bulk subspace is surprisingly robust across different training initializations and even when training with shuffled labels. This aligns with the hypothesis that the bulk is a result of data structure, suggesting that SGD navigates a low-dimensional subspace to optimize network performance. Conversely, training constrained to the chaotic or stable subspaces still achieves successful optimization, supporting the conjecture that these regions are more critical to the training dynamics than the bulk.
The findings are nuanced by extending spectral analysis to larger models like LeNet-5 on MNIST. The effort to compute the singular value decomposition highlights the computational challenges of this approach, finding a similar bulk region as observed in smaller models, thus reinforcing the scalability of their observations.
In conclusion, the paper furnishes new insights into neural network training by elucidating the roles of chaotic, bulk, and stable regions of the training Jacobian. It paves the way for investigating the structural aspects of parameter space, providing a deeper understanding of SGD's behavior across various data types. Future research could focus on efficiently computing the training Jacobian in larger-scale models, potentially utilizing novel randomized linear algebra techniques to further disentangle the intrinsic dimensionality that dictates the trainability and generalization capabilities of neural networks.