I mentioned in an earlier post that my research project into progressive sharpening and the edge of stability began with the hypothesis that, to understand neural network generalization, we need to understand their training process. A new paper has shaken my confidence in that hypothesis.
The alternative is called the volume hypothesis, and I first saw it advanced in Chiang et al. 2023. We know from Zhang et al. 2016 that there do exist sets of neural network parameters that memorize their training data but have very poor performance on test data, but we don't land on those parameters when we actually train networks. My interpretation of that has been that there is something about the training process that leads us to "good" parameters instead of bad ones. Chiang et al. 2023 argues instead that, while such bad parameters exist, they are very rare. Specifically, Chiang et al. 2023 argues that the volume of the ste of parameters with low training loss and high test loss is much smaller than the volume of parameters with low training loss and low test loss. When we train a neural network, we land at random somewhere inside of the set of parameters with low training loss, and since the subset of parameters with high test loss is very small, we are vanishingly unlikely to land there - so unlikely that it effectively never happens. If this is true, the training process doesn't really matter: as long as we land somewhere in the set with low training loss, we'll almost certainly have good test loss too.
The figure below illustrates the difference. In all cases, our networks start out initialized at random points in the set of all parameters. In the conventional view, the trajectory of training bends to ensure that the final network lands somewhere inside the set of parameters with both low training and test losses. Under the volume hypothesis, the trajectory just gets us somewhere inside the set of parameters with low training loss, but almost all of that set also has low test loss.
To test this idea, Chiang et al. 2023 sampled neural network parameters randomly, looking for parameters that had good performance on the training set by sheer chance, and then measured the performance on the test set - an approach they named guess and check (G&C). As you might expect, this process is extremely computationally expensive - it scales exponentially with the dataset size, making my own exponential Euler solver look cheap. As a result, they were only able to do this for binary classification problems with training sets of less than 32 examples, taken from CIFAR-10 and MNIST. However, the parameters they found this way did, in fact, consistently have test performance about as good as parameters obtained using gradient descent, consistent with the hypothesis:
The above graph shows the mean test accuracy of networks trained on binary classification of CIFAR-10 classes (e.g., cat vs. dog, rather than classifying among all ten classes). The blue curve is networks trained using the conventional stochastic gradient descent (SGD) method, while the purple curve is networks "trained" using G&C. The horizontal axis is the number of "training samples". The vertical axis is the test accuracy, measured over the entire CIFAR-10 test set of the two classes. So at first glance, these look very similar, with G&C actually outperforming SGD. However, these training sets are so small that we can't know if this observation will generalize to meaningful problems.
However, Chiang et al. 2023 used a non-standard scaling for the initialization of their networks. For G&C, that doesn't matter, because scaling the parameters has no effect on the prediction. But for training with SGD, the scaling does matter, a lot. Until the development of modern initialization techniques, neural networks were notoriously difficult to train, because the scaling of the parameters caused gradients to explode or die out during backpropagation. Proper scaling like the Glorot initialization and the Kaiming initialization keep the magnitude of the gradients roughly stable, allowing successful training even for relatively deep networks. So using a non-standard scaling is likely to have a significant impact on the results. And ineed, Peleg and Hein 2024 repeated Chiang et al. 2023's experiment using the Kaiming initialization, and found SGD noticeably outperformed G&C, by as much as 82.6% vs. 96.3%.
That said, 82.6% is not that far from 96.3%. The G&C networks are still doing much better than I would have expected. We could imagine that perhaps most of the test performance is due to volume effects, while SGD just increases the test accuracy by a small amount. We also need to ask whether rejection sampling from a normal distribution is the right probability distribution to be considering here. The statistical properties of trained neural networks are noticeably different from the properties of networks initialized with the Kaiming initialization -- for example, the expected spectrum of the Jacobian singular values is very different. Possibly if we were using the right probability distribution in G&C, its test accuracy would match the test accuracy of SGD.
On the other hand, perhaps the relatively good performance of G&C is an artifact of the small size of the dataset. There's some indicators that that may be what's going on in Peleg and Hein 2024. In particular, they find that the performance of SGD-trained networks improves as the network width increases, but not with G&C, which suggests we're dealing with two different phenomena.
I think the only way we're going to be able to solve this question is by pushing the G&C method far enough to extract scaling laws. We know that, as the network capacity and the training dataset size increase, the test loss of a network trained with SGD improves in a predictable way, following a power law. If we could stretch the curves in these figures out a bit, we could derive a similar scaling law for G&C, letting us estimate how it would perform on datasets of meaningful size. If we can show that G&C scales significantly worse than SGD, than that would imply that the volume hypothesis is wrong. But doing that will require some significant improvements in computational efficiency, given the exponential scaling of the cost with the dataset size. For now, the best option is to reserve judgement. That said, this is one of the most interesting results in deep learning theory that I've seen recently, and it's given me a great deal of food for thought.