Limitations of Neural Collapse for Understanding Generalization in Deep Learning

TL;DR

This paper scrutinizes Neural Collapse by separating train-collapse (an optimization property) from test-collapse (a potential generalization property) and introducing weak/strong variants for tests. It shows train-collapse reliably emerges in practice while test-collapse does not, across diverse datasets and architectures, and that stronger train-collapse can harm generalization and downstream transfer. The work also reveals a negative link between collapse and representational quality in transfer tasks, challenging claims that Neural Collapse universally benefits representation learning. Additionally, it provides preliminary evidence of cascading collapse across layers and calls for a tempered view of NC's role, emphasizing its optimization rather than generalization relevance and outlining directions for future theory and experiments.

Abstract

The recent work of Papyan, Han, & Donoho (2020) presented an intriguing "Neural Collapse" phenomenon, showing a structural property of interpolating classifiers in the late stage of training. This opened a rich area of exploration studying this phenomenon. Our motivation is to study the upper limits of this research program: How far will understanding Neural Collapse take us in understanding deep learning? First, we investigate its role in generalization. We refine the Neural Collapse conjecture into two separate conjectures: collapse on the train set (an optimization property) and collapse on the test distribution (a generalization property). We find that while Neural Collapse often occurs on the train set, it does not occur on the test set. We thus conclude that Neural Collapse is primarily an optimization phenomenon, with as-yet-unclear connections to generalization. Second, we investigate the role of Neural Collapse in feature learning. We show simple, realistic experiments where training longer leads to worse last-layer features, as measured by transfer-performance on a downstream task. This suggests that neural collapse is not always desirable for representation learning, as previously claimed. Finally, we give preliminary evidence of a "cascading collapse" phenomenon, wherein some form of Neural Collapse occurs not only for the last layer, but in earlier layers as well. We hope our work encourages the community to continue the rich line of Neural Collapse research, while also considering its inherent limitations.

Figures (6)

• Figure 1: Failure of Test Collapse. Neural Collapse for ResNet18 on CIFAR-10. Collapse occurs on the train set, but not on test.
• Figure 2: Failure of Test Collapse. Training and test variance vs. SGD iterations, for various dataset and architecture combinations. We train all models to 0 training error and continue training to achieve close to 0 training loss. All test sets (black line) do not collapse to negligible variance, and have much less collapse than the train sets (purple line).
• Figure 3: Neural Collapse on CIFAR-10. Collapse occurs on the train set, but not on the test set (neither Strong nor Weak). Weak test collapse has smaller variance than Strong test collapse.
• Figure 4: Train vs. Test Anti-Correlation. We vary the size of the train set ($N$), and observe that train and test collapse are anti-correlated. Top: ResNet18 trained on subsets of CIFAR-10. Bottom: VGG11 trained on subsets of FashionMNIST.
• Figure 5: Collapsed Features Transfer Worse. We save different checkpoints during pre-training, and use them to initialize the downstream models. The $x$-axis shows the TrainVariance of those checkpoints on the pre-training train set, and $y$-axis shows the test accuracy after fine-tuning on downstream tasks. We find that stronger train collapse (i.e. lower variance) is correlated with lower downstream test accuracy. Left: MNIST with a 3 hidden layer fully-connected network. Right: CIFAR-10 with a standard Resnet18.

Theorems & Definitions (4)

• Definition 1: Train-Collapse
• Conjecture 1: Train-Collapse Conjecture, informal
• Definition 2: Weak Test-Collapse
• Definition 3: Strong Test-Collapse