WHAT DEEP REPRESENTATIONS SHOULD WE LEARN? -A NEURAL COLLAPSE PERSPECTIVE

Abstract

For classification tasks, when sufficiently large networks are trained until convergence, an intriguing phenomenon has recently been discovered in the last-layer classifiers, and features termed neural collapse (NC): (i) the within-class variability of the features collapses to zero, and (ii) the between-class feature means are maximally and equally separated. Despite of recent endeavors to understand why NC happens, a fundamental question remains: whether NC is a blessing or a curse for deep learning? In this work, we investigate the problem under the setting of transfer learning that we pretrain a model on a large dataset and transfer it to downstream tasks. Through various experiments, our findings on NC are twofold: (i) when pre-training models, preventing intra-class variability collapse (to a certain extent) better preserve the structures of data, and leads to better model transferability; (ii) when fine-tuning models on downstream tasks, obtaining features with more NC on downstream data results in better test accuracy on the given task. Our findings based upon NC not only explain many widely used heuristics in model pretraining (e.g., data augmentation, projection head, self-supervised learning), but also leads to more efficient and principled transfer learning method on downstream tasks.

1. INTRODUCTION

Recently, an intriguing phenomenon has been discovered in terms of learned deep representations, in which the last-layer features and classifiers collapse to simple but elegant mathematical structures on the training data: (i) for each class, the intra-class variability of last-layer features collapses to zero, and (ii) the between-class class means and the last-layer classifiers all collapse to the vertices of a Simplex Equiangular Tight Frame (ETF) up to scaling. This phenomenon, termed Neural Collapse (N C) (Papyan et al., 2020; Han et al., 2022) , has been empirically demonstrated to persist across a variety of network architectures and datasets. Theoretically, more recent works (Fang et al., 2021; Zhu et al., 2021; Zhou et al., 2022; Tirer & Bruna, 2022) justified the prevalence of N C under simplified unconstrained feature models across a variety of training losses and problem formulations. Despite of recent endeavors of demystifying such an interesting phenomenon, a fundamental question lingers: is N C a blessing or a curse for deep representation learning? Understanding such a question could address many important but mysterious aspects of deep representation learning. For example, quite a few recent works (Papyan et al., 2020; Galanti et al., 2022; Hui et al., 2022) studied the connection between N C and generalization of overparameterized deep networks. In this work, we aim to understand transfer learning by studying the relationship between N C and the transferability of pretrained deep models. Transfer learning has become an increasingly popular approach in computer vision, medical imaging, and natural language processing (Zhuang et al., 2020) . With domain similarity, a pretrained large model on upstream datasets is reused as a starting point for fine-tuning a new model on a much smaller downstream task (Zhuang et al., 2020) . The pretrained model reuse with fine-tuning significantly reduces the computational cost, and achieves superior performances on problems with limited training datasets. However, without principled guidance, the underlying mechanism of transfer learning is not very well understood. First, when we are pretraining deep models on the upstream dataset, we lack good metrics for measuring the quality of the learned model or representation. In the past, people tended to rely empirically on controversial metrics for predicting the transferred test performance, such as the validation accuracy on the pretrained data (e.g., validation accuracy on ImageNet (Kornblith et al., 2019) ). For example, some popular approaches (e.g., label smoothing (Szegedy et al., 2016) 

