[Paper Review] Generalization Guarantees for Neural Networks via Harnessing the Low-rank Structure of the Jacobian

Samet Oymak, Zalan Fabian+2 more|arXiv (Cornell University)|Jun 12, 2019

Model Reduction and Neural NetworksPhysics and Astronomy53 references40 citations

TL;DR

The paper develops a Jacobian-based, data-dependent theory showing how neural networks generalize by separating learning dynamics into an information space (fast, well-aligned with labels) and a nuisance space (slow, potentially overfitting), and demonstrates that even constant-width nets can generalize on nicely structured data.

ABSTRACT

Modern neural network architectures often generalize well despite containing many more parameters than the size of the training dataset. This paper explores the generalization capabilities of neural networks trained via gradient descent. We develop a data-dependent optimization and generalization theory which leverages the low-rank structure of the Jacobian matrix associated with the network. Our results help demystify why training and generalization is easier on clean and structured datasets and harder on noisy and unstructured datasets as well as how the network size affects the evolution of the train and test errors during training. Specifically, we use a control knob to split the Jacobian spectum into "information" and "nuisance" spaces associated with the large and small singular values. We show that over the information space learning is fast and one can quickly train a model with zero training loss that can also generalize well. Over the nuisance space training is slower and early stopping can help with generalization at the expense of some bias. We also show that the overall generalization capability of the network is controlled by how well the label vector is aligned with the information space. A key feature of our results is that even constant width neural nets can provably generalize for sufficiently nice datasets. We conduct various numerical experiments on deep networks that corroborate our theoretical findings and demonstrate that: (i) the Jacobian of typical neural networks exhibit low-rank structure with a few large singular values and many small ones leading to a low-dimensional information space, (ii) over the information space learning is fast and most of the label vector falls on this space, and (iii) label noise falls on the nuisance space and impedes optimization/generalization.

Motivation & Objective

• Motivate and quantify how gradient-descent trained neural networks generalize despite overparameterization.
• Introduce a data-dependent decomposition of learning dynamics via the Jacobian spectrum into information and nuisance spaces.
• Show how alignment of labels with the information space and low-rank Jacobian enable strong generalization, even with modest width.
• Analyze bias–variance tradeoffs and network size effects on training and test performance.
• Incorporate arbitrary initializations (including pretrained models) into the generalization framework.

Proposed method

• Define information and nuisance spaces using the singular value decomposition of the Jacobian of the network.
• Decompose training dynamics and generalization error into information-space and nuisance-space contributions.
• Use a bias–variance framework where bias arises from misalignment with the information space and variance from model movement away from initialization.
• Provide finite-sample, data-dependent guarantees (Theorems 3.2 and 3.3) for random and arbitrary initializations, via the Multiclass Neural Tangent Kernel (M-NTK).
• Show that, under low-rank Jacobian structure, width can be modest (e.g., logarithmic in data size) and still achieve good generalization.

Experimental results

Research questions

• RQ1 Can gradient descent generalize for overparameterized networks by exploiting a low-rank Jacobian structure?
• RQ2 How does the alignment of labels with the Jacobian's information space affect generalization performance?
• RQ3 What is the role of network width in generalization when the Jacobian is effectively low-rank?
• RQ4 Do pre-trained or arbitrarily initialized models admit similar generalization guarantees under Jacobian-based analysis?
• RQ5 How do bias and variance components separate in the context of information and nuisance spaces?

Key findings

• The Jacobian of typical neural networks exhibits low-rank structure with a few large singular values and many small ones, defining a low-dimensional information space.
• Learning is fast over the information space and most of the label vector lies in this space, enabling fast reduction of train error.
• Learning over the nuisance space is slower and early stopping helps generalization at the cost of some bias.
• Generalization improves when the label vector is well-aligned with the information space; width can be constant or modest for sufficiently structured data.
• The framework yields data-dependent guarantees that do not require extremely wide networks, and results extend to arbitrary initializations including pretrained models.
• Numerical experiments corroborate the theoretical claims and illustrate fast convergence on information directions and slower, bias-prone learning on nuisance directions.