An Effective Gram Matrix Characterizes Generalization in Deep Networks

TL;DR

The paper tackles the problem of explaining generalization in deep networks trained by gradient methods by formulating a dynamical systems framework for the averaged generalization gap. It develops a differential equation for the gap, governed by a time-varying contraction factor $\bar{c}_n(t)$ and a perturbation factor $\bar{\epsilon}_n(t)$, and introduces an effective Gram matrix $K_n$ that expresses the asymptotic gap as $\vec{r}_n(0)^\top K_n \vec{r}_n(0)$, with $\vec{r}_n(0)$ the initial residual. The results show how the residual dynamics couple to the gradient covariance to drive generalization and demonstrate, across multiple datasets and architectures, that this quadratic form accurately predicts test loss while revealing a benign training regime in which the residual concentrates in the low-eigenvalue subspace of $K_n$. The framework links data- and architecture-dependent complexity to generalization, connects to kernel-method intuitions, and suggests extensions to other optimizers and interpolation settings with practical implications for understanding and predicting neural-network generalization.

Abstract

We derive a differential equation that governs the evolution of the generalization gap when a deep network is trained by gradient descent. This differential equation is controlled by two quantities, a contraction factor that brings together trajectories corresponding to slightly different datasets, and a perturbation factor that accounts for them training on different datasets. We analyze this differential equation to compute an ``effective Gram matrix'' that characterizes the generalization gap in terms of the alignment between this Gram matrix and a certain initial ``residual''. Empirical evaluations on image classification datasets indicate that this analysis can predict the test loss accurately. Further, during training, the residual predominantly lies in the subspace of the effective Gram matrix with the smallest eigenvalues. This indicates that the generalization gap accumulates slowly along the direction of training, charactering a benign training process. We provide novel perspectives for explaining the generalization ability of neural network training with different datasets and architectures through the alignment pattern of the ``residual" and the ``effective Gram matrix".

Figures (10)

• Figure 1: The contraction factor calculated through its analytical expression in \ref{['eq: c']} (orange) compared to its approximation using \ref{['eq: approx c']} (blue) for FC trained on MNIST with two selected classes, $n=1000$, $m=100$.
• Figure 2: Left: FC trained on MNIST with all 10 classes, with $n=1100$ samples and statistics computed over datasets perturbed by $m=100$ samples. Right: LeNet-5 trained on CIFAR-10 with 2 selected classes, $n=1100$, $m=100$. To be consistent with the calculations, all our experiments are conducted with gradient descent, not stochastic gradient descent. Practically, this means that in order to get the network to fit the training data well enough, we need to use small sample sizes $n$.
• Figure 3: Statistics of the residual $\vec{r}_n$ and effective Gram matrix $K_n$ for two different tasks. Benign task: FC trained on MNIST with all 10 classes, $n=1000$, $m=100$ for which the network generalizes well. Random task: FC trained on MNIST with 10 randomly assigned classes, $n=50$, $m=5$ where the network does not generalize. Left: Eigenspectrum of the Gram matrix $\sigma(K_n)$ and the normalized projection of initial residual $P(\vec{r}_n(0), E(K_n))$ for benign and random tasks. Right: Explained magnitude of the initial residual $M(\vec{r}_n(0),E(K_n))$ for benign and random tasks. We use relative index ranges from 0 to 1 for better comparison of configurations with different number of samples (see \ref{['rmk: normalization']} for illustration).
• Figure S.1: Explained magnitude of the residual $M(\vec{r}_n(t_0),E(K_n(t_0)))$ (Y-axis) as a function of the explained magnitude of the effective Gram matrix $M(K_n(t_0))$ (X-axis) for FC trained on MNIST with all 10 classes, with $n=1100$ and $m=100$, but computed for different times $t_0$. We see that as the number of training iterations increases, the explained magnitude of the residuals in the subspace of the effective Gram matrix with a small explained magnitude, i.e., the non-principal subspace, decreases. Residuals at later training times project more and more predominantly in the principal subspace of the effective Gram matrix.
• Figure S.2: Evaluation on synthetic datasets

Theorems & Definitions (15)

• Lemma 1
• Lemma 2
• Remark 3: Classical contraction theory with uniform bounds on contraction and perturbation
• Lemma 4
• Theorem 5
• Lemma 6: Existence of the effective Gram matrix
• Remark 7
• Remark 8: Data and architecture dependent generalization bound
• Theorem 9: Theorem 2.1 from TsukamotoContraction2021
• Theorem 10: Theorem 2.3 from TsukamotoContraction2021