On the Role of Neural Collapse in Transfer Learning (2112.15121v2)

Abstract: We study the ability of foundation models to learn representations for classification that are transferable to new, unseen classes. Recent results in the literature show that representations learned by a single classifier over many classes are competitive on few-shot learning problems with representations learned by special-purpose algorithms designed for such problems. In this paper we provide an explanation for this behavior based on the recently observed phenomenon that the features learned by overparameterized classification networks show an interesting clustering property, called neural collapse. We demonstrate both theoretically and empirically that neural collapse generalizes to new samples from the training classes, and -- more importantly -- to new classes as well, allowing foundation models to provide feature maps that work well in transfer learning and, specifically, in the few-shot setting.

• The paper demonstrates that neural collapse, quantified by low CDNV, generalizes from source to target classes in few-shot learning scenarios.
• It establishes theoretical bounds and validates empirically on datasets like Mini-ImageNet and CIFAR-FS to link feature clustering to improved classification.
• The findings suggest that diverse pretraining induces transferable feature spaces, enabling simple classifiers to perform effectively with minimal data.

This paper investigates why foundation models, typically trained on large-scale classification tasks with many classes, are highly effective for transfer learning, particularly in few-shot learning scenarios where data for new tasks is scarce. The authors propose that the phenomenon of Neural Collapse (NC) plays a crucial role in this effectiveness.

Neural Collapse is an observation made in the training of deep neural networks for classification, where the features (outputs of the penultimate layer) for samples belonging to the same class cluster tightly around their class mean. Furthermore, these class means exhibit specific geometric structures (approaching a simplex equiangular tight frame), and the last-layer classifier behaves like a nearest-class-mean rule. The paper defines Neural Collapse using Class-Distance Normalized Variance (CDNV), which is the ratio of the within-class variance to the squared distance between class means. A low CDNV indicates strong neural collapse.

The core contribution of the paper is demonstrating both theoretically and empirically that this neural collapse property generalizes beyond the training data and source classes to new, unseen classes.

The paper's theoretical analysis proceeds in three steps:

1. Generalization to New Samples from Source Classes: Proposition~\ref{prop:generalizationInner} shows that if NC (measured by CDNV) is observed on the training data from a pair of source classes, it is also expected to be small for unseen test samples from those same classes. This generalization is bounded by terms related to the standard generalization gap (differences between expected values and empirical averages), which typically decrease as the number of training samples per class increases.
2. Generalization to New Classes: Proposition~\ref{prop:newclasses} extends this idea to unseen target classes. It argues that if neural collapse is present across a large set of source classes, it is likely to emerge for new pairs of target classes drawn from the same underlying distribution of classes. The bound on the expected CDNV for target classes depends on the average CDNV over the source classes and terms that decrease as the number of source classes increases. This suggests that training on a wide variety of source classes encourages feature representations where new classes will also exhibit the collapsed structure.
3. CDNV and Classification Error: Section~\ref{sec:error} and Proposition~\ref{prop:cdnv_error} connect the presence of neural collapse (small CDNV) in the feature space of target classes to good classification performance in a few-shot setting. The paper shows that the classification error of a simple nearest-class-mean classifier on the learned features is bounded by terms related to the CDNV on the target classes. Specifically, if the CDNV is small, the error is expected to be low, even with few samples. For spherical Gaussian feature distributions, the error bound can be exponentially small in the feature dimension relative to the CDNV.

Practical Implications and Implementation:

The findings suggest a practical strategy for building transferable representations:

• Train a deep neural network for standard classification on a large dataset covering many classes (acting as the source task). This training process naturally encourages neural collapse in the penultimate layer's feature space.
• Once the feature extractor (the network excluding the final classification layer) is trained, freeze its weights.
• For a new target classification task, use the frozen feature extractor to compute features for the available few-shot training samples.
• Train a simple linear classifier or use a nearest-class-mean rule directly on these extracted features. The paper demonstrates that ridge regression is a suitable and competitive choice for this step. The weight matrix for ridge regression on the feature vectors $f(X)$ and labels $Y$ is $W = (f(X)^\top f(X)+\lambda_n I)^{-1} f(X)^\top Y$.

This approach simplifies the adaptation to new tasks because the "hard part" (learning a general, well-structured feature space via the deep network) is done once on the source data. The beneficial geometric properties induced by neural collapse mean that the target task adaptation becomes a much simpler problem in the feature space, requiring only a few examples to determine the class locations (means or hyperplane boundaries).

Experimental Validation:

The authors validate their theory with experiments on Mini-ImageNet, CIFAR-FS, FC-100, and EMNIST datasets using Wide ResNets and standard convolutional networks.

• They show that the CDNV on source test data decreases during training, generalizing NC from training samples to unseen samples of the same classes (Figure~\ref{fig:var_nc_main}b).
• Crucially, they demonstrate that the CDNV on target classes also decreases as the number of source classes used for pretraining increases (Figure~\ref{fig:var_nc_main}c). This supports the theoretical finding that training on a broader source task improves generalization of the collapsed structure to new classes.
• The decrease in target CDNV is strongly correlated with improved few-shot classification accuracy on target tasks using the learned features (Figure~\ref{fig:var_nc_main}d).
• Their simple ridge regression method on frozen features achieves competitive results compared to specialized few-shot learning algorithms (Table~\ref{tab:main_results}).
• Additional experiments confirm that the phenomenon is consistent across different learning rates (Figures~\ref{fig:var_nc_cifar_fs_lr}, \ref{fig:var_nc_emnist_lr}), that the minimum distance between class means in the feature space tends to increase (Figure~\ref{fig:dist_fs}), and that similar (though less pronounced) collapse can be observed in lower layers of the network (Figure~\ref{fig:mid_layer}). They also show consistency when using the original Class-Covariance Normalized Variance (CCNV) metric (Figures~\ref{fig:ccnv_cifar_fs}, \ref{fig:ccnv_emnist}).

Implementation Considerations:

• The quality and diversity (number of classes) of the source dataset are critical for inducing generalizable neural collapse.
• Standard classification training objectives (like cross-entropy minimization) appear sufficient to achieve this.
• The target task adaptation is computationally inexpensive, involving only a simple linear model trained on the extracted features of the few-shot support set. This makes deployment efficient.
• The paper's theoretical bounds rely on assumptions like bounded feature norms and specific properties of the class distributions, which may not perfectly hold in all real-world scenarios. However, the empirical results suggest the phenomenon is broadly applicable.

In summary, the paper provides a compelling explanation rooted in Neural Collapse for the effectiveness of pretraining foundation models for transfer learning. It shows that standard classification training on a large source task creates a feature space where examples from new classes are well-clustered and separable, greatly simplifying the task of learning a classifier from few examples. This suggests that inducing strong neural collapse during pretraining is a key goal for developing transferable representations.