Source-linked AI summary
Generalization Guarantees for Neural Networks via Harnessing the Low-rank Structure of the Jacobian
Samet Oymak, Zalan Fabian, Mingchen Li, Mahdi Soltanolkotabi
TL;DR
The paper asks why overparameterized neural networks trained by gradient descent can generalize despite being able to fit noise. It develops a Jacobian-based theory that separates information and nuisance directions, and establishes multiclass generalization bounds while describing distinct learning dynamics across those spaces.
Problem
The paper investigates why overparameterized neural networks can fit arbitrary training data yet generalize to unseen data when trained by gradient descent.
Source: §1.1 · p. 1Method
The paper uses the network Jacobian and its singular directions to characterize optimization and generalization, with bounds for randomly initialized and arbitrary-initialization networks.
Source: §1.1 · p. 1, §3 · p. 6Results
The paper establishes multiclass generalization bounds and describes rapid learning over the information space followed by slower nuisance-space learning and a slight test-error increase.
Source: §2.1 · p. 4, §3 · p. 6Takeaways & Limitations
The linear example illustrates how information-space error can fall quickly while nuisance-space error rises slowly, producing a dip in total test error.
Source: §2.1 · p. 4The proposed adaptation of the Jacobian to the dataset is observed numerically but is not formally proven in this paper.
Source: §3.2 · p. 9
Abstract
from arXiv · showhide
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.
1 Introduction
The paper studies why overparameterized neural networks trained by gradient descent can generalize despite their capacity to fit noise. It develops a Jacobian-based theory and specifies a one-hidden-layer model and gradient-descent training setup.
1.1 Motivation and contributions
- The paper uses the network Jacobian to characterize directions where learning is fast and generalizable versus slow and prone to overfitting.Its stated aim is to help explain how networks can fit noise yet generalize when trained on real datasets.Source: §1.1 · p. 1
- The theory links easier training and generalization to semantically linked features and labels, and separates learning into information and nuisance spaces with different bias–variance behavior.Gradient descent nearly interpolates over the information space with small bias, while optimization over the nuisance space is slow and overfits due to higher variance.Source: §1.1 · p. 2
- Larger networks have smaller prediction bias, while sufficiently structured datasets can let small networks generalize despite typically higher bias.The paper contrasts this with cited guarantees that require widths growing inversely with input distances, class margins, or related quantities [4] [21] [2] [15] [40] [1] [14].Source: §1.1 · p. 2
1.2 Model and training
- The theoretical analysis focuses on a one-hidden-layer network with d inputs, k hidden neurons, K outputs, and input-to-hidden weights W.The model maps inputs in R^d to outputs in R^K; the training setup fixes the output layer and optimizes W.Source: §1.2 · p. 2
- The paper studies the generalization of models produced by iterative gradient-descent updates from an initialization W0.The stated analysis does not require random initialization and also applies to arbitrary initialization, including pretrained models.Source: §1.1 · p. 2, §1.2 · p. 3
2 Components of a Jacobian-based theory of generalization
A linear regression example illustrates how total test error can fall and then rise as gradient descent moves from fitting informative directions to fitting nuisance directions. The paper extends this intuition to neural networks by splitting Jacobian directions according to their singular values and relating that split to training and generalization dynamics.
2.1 Prelude: fitting a linear model
- The example models overparameterized linear regression with low-rank uncorrupted features and Gaussian corruption in the features and labels.The setup assumes n ≤ d and defines the clean labels through the low-rank feature component.Source: §2.1 · p. 3
- In the linear example, test error falls quickly at first, then slowly rises as error on nuisance directions grows despite decreasing information-space error.The information space is the span of U, while the nuisance space is its complement.Source: §2.1 · p. 4
- The decomposition explains the test-error dip: prominent singular directions align with information and are learned quickly, while less prominent directions accumulate residual energy with further iterations.The total test error combines the information and nuisance components.Source: §2.1 · pp. 4–5
2.2 Information and nuisance spaces of the Jacobian
- The Jacobian spectrum governs gradient dynamics: residual alignment with top singular directions enables substantial training-error reduction, and label alignment with prominent directions is associated with low test error.The paper describes a low-rank Jacobian with a few large and many small singular values, and reports fast early learning followed by slower nuisance-space learning and a slight test-error increase.Source: §2.2 · pp. 5–6
- For neural networks, the paper defines information and nuisance spaces from a reference Jacobian by thresholding its singular values.The Jacobian may be defined using its expected value at random initialization or its value at an iterate.Source: §2.2 · p. 6
3 Main results
This section develops Jacobian-based generalization bounds for neural networks trained by gradient descent, using information and nuisance spaces to analyze generalization dynamics under random and arbitrary initializations. It covers capacity-dependent interpolation, information/nuisance-space tradeoffs, multiclass classification, and pre-trained networks.
3.1 Results for random initialization
- With an approximately low-rank kernel and labels mostly in the information space, moderate-width networks can generalize; under stated parameter choices, width logarithmic in n suffices.The theorem’s conditions connect generalization to both low-dimensional kernel structure and label alignment.Source: §3.1 · pp. 7–8
- The nuisance-space label component creates bias, while distance from initialization contributes variance; larger networks can reduce bias, and low-rank cases can converge in a constant number of iterations.The iteration bound is inversely proportional to the smallest singular value over the information space.Source: §3.1 · pp. 7–8
3.2 Generalization guarantees with arbitrary initialization
- For arbitrary initialization, the generalization guarantee depends on the initial residual’s nuisance-space projection and applies to pretrained models and intermediate iterates without requiring random initialization.A Jacobian that gives a better low-rank representation and aligns the residual with its information space yields sharper guarantees.Source: §3.2 · p. 9
3.3 Case Study: Gaussian mixture model
- In the zero-noise Gaussian mixture, the Jacobian rank is at most K^2C, labels lie entirely in the information space, and information-space singular values scale with n/KC.The rank bound follows from KC distinct input vectors and K output nodes.Source: §3.3 · pp. 10–11
- In the small-noise regime, good generalization is achievable when n ≳ K^2C, the effective rank of the M-NTK matrix; required width is independent of n and k ≳ K^8C^4.Source: §3.3 · p. 11
3.4 Prior Art
- The prior-art discussion introduces classical statistical learning theory work dating to the 1990s and renewed study of neural-network capacity following the success of deep networks.It also notes studies of capacity under different norm constraints and architectures.Source: §3.4 · p. 11
- The paper analyzes generalization dynamics across gradient-descent iterations, unlike bounds for empirical-risk minimizers and mean-field analyses focused on asymptotic behavior.Prior work also studies gradient-descent generalization empirically, linearly separable logistic regression, and kernel methods at interpolation.Source: §3.4 · p. 12
- Unlike prior NTK-based analyses that typically require very wide networks, the work allows capacity-dependent interpolation with small or constant width and no restrictive input assumptions.The paper also reports faster convergence rates.Source: §3.4 · p. 12
- The work distinguishes its generalization bounds through information/nuisance-space bias-variance tradeoffs, no margin, distance, or minimum-eigenvalue assumptions, and coverage of multiclass and pre-trained networks.The multiclass and pre-trained network results are stated in Theorem 3.3.Source: §3.4 · p. 12
4 Numerical experiments
The experiments find a low-rank Jacobian spectrum and faster learning along its information space, with label corruption shifting energy toward the nuisance space and worsening test error. On uncorrupted data, the Jacobian becomes more aligned with labels during training; under corruption, test error first falls and then rises as nuisance-space residual grows.
4 Numerical experiments
- The experiments use a modified ResNet20 with batch normalization and softmax removed, least-squares loss, and SGD for the reported SGD runs.The full-Jacobian calculation motivates restricting the CIFAR-10 experiments to three classes and subsampling training data; broader spectrum experiments are planned using distributed computation.Source: §4 · p. 12
- The Jacobian has tens of large singular values and a majority of small ones, motivating low-dimensional information and complementary nuisance spaces.The experiments calculate the spectrum on a modified ResNet20 using three CIFAR-10 classes and the top 50 singular vectors as the information space.Source: §4 · pp. 12–13
- After training on uncorrupted data, more label and initial-residual energy lies in the information space of the final Jacobian than of the initial Jacobian.The reported alignment is associated with less nuisance-space energy and a relatively small ∥J†Iy∥ℓ2.Source: §4 · p. 13
- On uncorrupted data with SGD, information-space residual energy drops rapidly, nuisance-space residual declines slowly, and both track training and test misclassification error.The decrease in total residual energy is described as overwhelmingly governed by the information space.Source: §4 · p. 14
- With Adam on uncorrupted data, training reaches a perfect fit in fewer iterations than with SGD, and information-space residual falls below 4 ⋅10−4 after 100 epochs.The authors also report that Adam's Jacobian adapts to labels and initial residual faster than SGD's on this dataset.Source: §4 · p. 14
- With 50% label corruption, test error reaches its minimum at 100 epochs and then rises as nuisance-space residual increases while information-space residual falls.Training error steadily decreases during this behavior.Source: §4 · p. 15
- After 800 epochs, across 25%, 75%, and 100% label corruption, test error increases with corruption, while corrupted labels become less correlated with the information space and more label energy falls on the nuisance space.This is consistent with the authors’ theory, which predicts worse generalization in this case.Source: §4 · p. 17
5 Technical approach and General Theory
The section develops a gradient-descent analysis based on a reference Jacobian and its information and nuisance subspaces. A meta-theorem relates nonlinear and linearized training dynamics under bounded-spectrum and bounded-perturbation assumptions.
5 Technical approach and General Theory
- The framework models fitting a nonlinear multi-output function to feature-label data with nonlinear least-squares loss and gradient-descent updates.It uses the Jacobian of the model output with respect to parameters to express the gradient.Source: §5 · p. 17
- A reference Jacobian close to the initialization Jacobian supplies the spectrum used to define information and nuisance subspaces through a cutoff.The information space is associated with prominent singular directions, and the truncated Jacobian acts on that space.Source: §5 · pp. 17–18
- The theory treats information and nuisance subspaces as directions where learning is respectively fast and generalizable versus slow and prone to overfitting.The analysis assumes a bounded Jacobian spectrum and bounded perturbation around the initialization point.Source: §5 · pp. 18–19
- Under the stated conditions, the meta-theorem shows that nonlinear and linearized gradient-descent iterates and their residuals closely track through the specified time horizon.The comparison concerns a reference Jacobian associated with the Jacobian at initialization and assumes the displayed inequalities in the theorem.Source: §5 · p. 19
6 Proofs
This section introduces notation for reshaping matrices and vectors and for derivatives used in the proofs.
6 Proofs
- vect(W) denotes the vector formed by concatenating the rows of matrix W.The rows are ordered w1, w2, …, wk.Source: §6 · p. 19
- mat(w) denotes the matrix obtained by reshaping vector w across its rows.For w in R^kd, the resulting matrix is k × d.Source: §6 · p. 19
- For a differentiable scalar function φ, φ′ and φ′′ denote its first and second derivatives.Source: §6 · p. 19
6.1 Proofs for General Theory (Proof of Theorem 5.3)
The proof proceeds from preliminary definitions and lemmas to bounds on linearized residuals and their deviation from nonlinear residuals, then completes the meta-theorem by induction.
6.1.1 Preliminary definitions and lemmas
- The preliminary section defines the early-stopping value and distance and establishes bounds on the stopping value, alongside matrix perturbation lemmas.The early-stopping distance is intended to characterize distance from initialization at an appropriate stopping time.Source: §6.1.1 · pp. 20–21
6.1.2 Key lemmas for general nonlinearities
- For linearized gradient descent, enough iterations fit the residual over the information space while leaving it unchanged in the worst case over the nuisance space.The proof also develops bounds for tracking the difference between nonlinear and linearized residuals.Source: §6.1.2 · p. 24
- The key lemmas bound residual evolution and perturbation between nonlinear and linearized problems under the stated assumptions and iteration horizon.These bounds provide the ingredients used to complete Theorem 5.3.Source: §6.1.2 · p. 25, §6.1.3 · p. 26
6.1.3 Completing the proof of Theorem 5.3
- The proof completes Theorem 5.3 by induction, using the lemmas to bound residual tracking, parameter distance from initialization, and the early-stopped residual.The argument applies within the theorem's stated neighborhood and time-horizon conditions.Source: §6.1.3 · pp. 26, 28, 29
6.2 Key lemmas and identities for neural networks
This appendix establishes identities and bounds for one-hidden-layer network Jacobians, their behavior at random initialization, and the initial residual under stated assumptions.
6.2.1 Preliminary identities and calculations
- The one-hidden-layer network Jacobian is represented class by class and can also be written using Khatri-Rao products.These identities provide the representations used in the later Jacobian analyses.Source: §6.2.1 · p. 30
6.2.2 Fundamental properties of the Jacobian of the neural network
- With bounded activation derivatives, the single-output and multiclass Jacobian mappings have spectral-norm bounds and are Lipschitz under the stated smoothness and unit-input assumptions.The section extends the single-output properties to the multiclass setting.Source: §6.2.2 · pp. 30–31
6.2.3 Properties of the Jacobian at random initialization
- At Gaussian input-weight initialization and zero-mean output weights, the expected Jacobian covariance is block diagonal and related to the multiclass kernel matrix.The section also proves concentration of the initialized Jacobian around its expected covariance under the stated sampling and weight assumptions.Source: §6.2.3 · pp. 32–33
6.2.4 Upper bound on initial residual
- For the stated one-hidden-layer model with unit-norm inputs and bounded activation values and derivatives, Lemma 6.14 provides a high-probability upper bound on the initial residual.The displayed bound’s value is not present in the supplied passage.Source: §6.2.4 · p. 34
6.3 Rademacher complexity and generalization bounds
This section bounds the complexity of a constrained neural-network function class and connects that bound to generalization for Lipschitz losses and classification error.
Generalization bounds
- The section uses standard Rademacher generalization bounds and vector contraction to relate empirical complexity to test loss for bounded, 1-Lipschitz losses.The stated results assume i.i.d. training data and unit-norm inputs.Source: §6.3 · p. 36
Rademacher complexity
- Lemma 6.17 bounds empirical Rademacher complexity for a neural-network class constrained around reference weights.The class and bound depend on weight constraints, reference-model energy, and sample size.Source: §6.3 · p. 37
Generalization bounds
- Lemma 6.18 connects test error for a Lipschitz loss to quadratic training loss and also gives a classification-error guarantee for one-hot labels.The classification result is stated for the specified function class and sampling assumptions.Source: §6.3 · pp. 39–40
6.4 Proofs for neural nets with arbitrary initialization (Proof of Theorem 3.3)
The arbitrary-initialization proofs bound gradient-descent residuals through information and nuisance components, then use those optimization bounds to obtain a generalization guarantee.
6.4.1 Preliminary Optimization Result
- The preliminary optimization bound makes the information-space component of the initial residual decay exponentially, while the nuisance component remains in the residual bound.A further term depends on the reference-Jacobian approximation tolerance.Source: §6.4.1 · p. 41
6.4.2 Main Optimization Result
- The main optimization result applies the preliminary bound under an additional condition that the initial residual has sufficient information-space energy.The supplied passages state the condition and theorem setup but omit the full resulting bound.Source: §6.4.2 · p. 43
6.4.3 Main generalization result (completing the proof of Theorem 3.3)
- After the stated number of gradient-descent iterations, the generalization error is bounded using the nuisance component of the initial residual and an early-stopping-distance term.The theorem assumes unit-norm inputs, bounded activation derivatives, bounded initial residual, and the specified learning-rate and width conditions.Source: §6.4.3 · p. 44
6.5 Proofs for neural network with random initialization (proof of Theorem 3.2)
The random-initialization proofs establish optimization and generalization bounds using information and nuisance spaces defined from the initialized Jacobian, with stated data, network, and label-alignment assumptions.
6.5.1 Optimization result
- Under the theorem’s assumptions, the random-initialization optimization result bounds residual and weight movement after gradient descent using the information and nuisance components of the labels.The setup uses Gaussian initial input weights, random output weights, and a cutoff defining the two subspaces.Source: §6.5.1 · p. 46
6.5.2 Generalization result (completing the proof of Theorem 3.2)
- The generalization theorem gives a high-probability upper bound on classification error after the specified gradient-descent iterations.The theorem assumes unit-norm inputs, unit-norm concatenated labels, and an information-space component of the labels bounded below; the supplied passages omit the bound’s full expression.Source: §6.5.2 · pp. 49–50
- The training-residual bound combines the nuisance-label component, an exponentially weighted information-label component, and a precision term.This residual estimate is used in the generalization proof.Source: §6.5.2 · p. 50
A The Jacobian of the Mixture Model is low-rank (Proofs for Section 3.3)
For the noiseless mixture model, the Jacobian’s information space is low-rank and contains the concatenated multiclass labels. The section also states a generalization result for Gaussian-mixture data with input noise under specified conditions.
A The Jacobian of the Mixture Model is low-rank (Proofs for Section 3.3)
- The noiseless mixture model’s Jacobian has a K^2C-dimensional information space containing the concatenated multiclass label vectors.Each class-specific label vector lies in its corresponding single-output information space.Source: A The Jacobian of the Mixture Model is… · pp. 54–55
- For a single output, the information space has dimension KC, and the proof uses this separability to reduce the analysis to one output.The full multiclass space is formed from the per-output spaces.Source: A The Jacobian of the Mixture Model is… · pp. 54–55
- The proof relates the dataset Jacobian to the cluster-center Jacobian by mapping repeated class/cluster samples into a KC-dimensional subspace.The samples belonging to the same class/cluster pair have equal coordinates in that subspace.Source: A The Jacobian of the Mixture Model is… · p. 55
- The mixture-model generalization theorem considers iid Gaussian-mixture data with input noise and gradient descent under its stated conditions.The theorem specifies a training schedule and learning-rate condition.Source: A The Jacobian of the Mixture Model is… · pp. 56–57
B Joint input-output optimization
The section constructs a joint Jacobian for training input and output weights and states its properties, including Lipschitzness.
B Joint input-output optimization
- The joint Jacobian combines derivatives with respect to both the input and output weights.The section denotes this combined Jacobian by J(V,W).Source: B Joint input-output optimization · p. 57
- The output-weight Jacobian is block diagonal, with each block equal to φ(XW^T).Its blocks correspond to the output-layer vectors.Source: B Joint input-output optimization · pp. 57–58
- Theorem B.1 states that the combined input/output Jacobian is Lipschitz, with the proof bounding changes in its two components separately.The output-weight component is linear in the output layer.Source: B Joint input-output optimization · p. 58