On the Generalization Mystery in Deep Learning

TL;DR

The paper tackles why over-parameterized neural networks generalize well despite memorizing random data. It proposes Coherent Gradients, a dataset-dependent stability mechanism, centered on the alignment (coherence) of per-example gradients and measured by the alpha metric. A data-dependent generalization bound ties the gap between training and test loss to coherence, and extensive experiments show real data induces higher gradient coherence than random data, with depth amplifying coherent signals. The work also introduces practical regularization methods (winsorization, median-of-means) that suppress weak gradient directions and improve generalization, offering a unifying, causal explanation that spans simple linear settings to large-scale vision models and points to concrete directions for future theory and algorithm design.

Abstract

The generalization mystery in deep learning is the following: Why do over-parameterized neural networks trained with gradient descent (GD) generalize well on real datasets even though they are capable of fitting random datasets of comparable size? Furthermore, from among all solutions that fit the training data, how does GD find one that generalizes well (when such a well-generalizing solution exists)? We argue that the answer to both questions lies in the interaction of the gradients of different examples during training. Intuitively, if the per-example gradients are well-aligned, that is, if they are coherent, then one may expect GD to be (algorithmically) stable, and hence generalize well. We formalize this argument with an easy to compute and interpretable metric for coherence, and show that the metric takes on very different values on real and random datasets for several common vision networks. The theory also explains a number of other phenomena in deep learning, such as why some examples are reliably learned earlier than others, why early stopping works, and why it is possible to learn from noisy labels. Moreover, since the theory provides a causal explanation of how GD finds a well-generalizing solution when one exists, it motivates a class of simple modifications to GD that attenuate memorization and improve generalization. Generalization in deep learning is an extremely broad phenomenon, and therefore, it requires an equally general explanation. We conclude with a survey of alternative lines of attack on this problem, and argue that the proposed approach is the most viable one on this basis.

Figures (29)

• Figure 1: An experiment in the spirit of Zhang17 illustrating the generalization mystery in deep learning. We train two ResNet-50 models, one on ImageNet with original labels ("real", top row), and another on ImageNet with images replaced by Gaussian noise ("random", bottom row) using vanilla SGD and no explicit regularization. As the loss and accuracy curves (first two columns) show, the network has sufficient capacity to memorize the training data, yet, it generalizes in one case and not the other. We believe that the reason for this difference in behavior can be found by analyzing the similarity between per-example gradients during training, that is, coherence. Using the $\alpha / \alpha_{m}^{\perp}$ metric for coherence (last two columns), we see that in the case of real data, the per-example gradients are much more similar, and each example helps reduce the loss on many other examples, as compared to the random case.
• Figure 2: Unlike the situation with ResNet-50 (\ref{['fig:chap1:ResNet-50imagenet']}), with AlexNet we find that the peak coherence for random data (second row) as measured by $\mathop{\mathrm{\mathnormal{\alpha_m / \alpha^{\perp}_{m}}}}\nolimits$ can be surprisingly high, even though it happens much later in training, and is lower than that of real (first row). Although this appears to be a contradiction to the theory, it is not; it is a limitation of the metric. $\mathop{\mathrm{\mathnormal{\alpha_m / \alpha^{\perp}_{m}}}}\nolimits$ in this plot is a measure of coherence over the entire network (that is, over entire per-example gradients), and is therefore an average quantity. A closer look at the layer-by-layer values of $\mathop{\mathrm{\mathnormal{\alpha_m / \alpha^{\perp}_{m}}}}\nolimits$ as shown in \ref{['fig:chap1:alexnetimagenet_layer']} reveals, once again, a significant difference between real and random data.
• Figure 3: A layer-by-layer breakdown of $\mathop{\mathrm{\mathnormal{\alpha_m / \alpha^{\perp}_{m}}}}\nolimits$ for AlexNet from \ref{['fig:chap1:alexnetimagenet_overall']} shows that on random data (second row), $\mathop{\mathrm{\mathnormal{\alpha_m / \alpha^{\perp}_{m}}}}\nolimits$ is indeed close to 1 and much lower than that of real data (first row) for the first few layers. For the higher (dense) layers, coherence is comparable between real and random, though note the difference in scale of $\mathop{\mathrm{\mathnormal{\alpha_m / \alpha^{\perp}_{m}}}}\nolimits$ between the convolutional and dense layer plots.
• Figure 4: Generalization improves when weak directions in the average gradient are suppressed during gradient descent. Weak directions are suppressed by winsorization, that is, by clipping extreme per-examples gradients (independently for each coordinate of the gradient). The parameter $c$ controls the level of winsorization and $c = 0$ corresponds to using the (usual) average gradient. We train a fully connected network on mnist with varying amounts of label noise (see \ref{['fig:overview:winsorization_gaussian']} for a similar experiment with random pixels).
• Figure 5: Taking the coordinate-wise median of 3 micro-batches ( M3) suppresses weak gradient directions more than taking their coordinate-wise average ( A3) (the same as ordinary SGD), leading to better generalization. Here, we train a ResNet-50 on 3 datasets derived from ImageNet by replacing some fraction of the training images with Gaussian noise. In the case of random data (100% noise), M3 prevents memorization, so the training and test curves both lie on the x-axis. In contrast, as expected, A3 (SGD) memorizes the training set. When only half the training images are replaced by noise (50% noise), M3 reaches a training accuracy near 50% suggesting that only the real images are learned. This is confirmed in \ref{['fig:overview:resnet_imagenet_a3_vs_m3_50pct_breakdown']}.

Theorems & Definitions (42)

• Example 1: "Real" and "Random"
• Example 2: Perfect similarity v/s pairwise orthogonal
• Example 3: Orthogonal Sample and Orthogonal Limit
• Example 4: "Real" and "Random"
• Theorem 1: Generalization Theorem
• proof
• Example 5: WGD applied to \ref{['ex:two-datasets']}
• Example 6: M3 applied to \ref{['ex:two-datasets']}
• Example 7: Adding Depth to Linear Regression
• Example 8: A Tale of Two Neurons