A Geometric Analysis of Neural Collapse with Unconstrained Features

1 Introduction . In the past decade , the revival of deep neural networks ( DNN ) has led to dramatic success in numerous applications ranging from computer vision , to natural language processing , to scientific discovery and beyond [ 2–5 ] . Nevertheless , the practice of deep networks has been shrouded with mystery as our theoretical understanding for the success of deep learning remains elusive . There are many intriguing phenomena , such as implicit algorithmic bias in training [ 6–10 ] , and good generalization of highly-overparameterized networks [ 7,11–15 ] , that seem often contradictory to , or can not be explained by , classical optimization and learning theory . ∗The first two authors contributed to this work equally . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . … … … … … Input Hidden layer 1 OutputHidden layer L-1 Figure 2 : Illustration of the unconstrained feature model , where the gray box is peeled off so that the representation h is modeled by a simple decision variable for every training sample . Towards demystifying DNN , recent seminal work [ 1 , 16 ] empirically discovered an intriguing phenomenon that persists across a range of canonical classification problems during the terminal phase of training . As illustrated in Figure 1 , it has been widely observed that last-layer features and classifiers of a trained DNN exhibit simple but elegant mathematical structures : • Variability Collapse : cross-example within-class variability of last-layer features collapses to zero , as the individual features of each class themselves concentrate to their isolated class-means . • Convergence to Simplex ETF : the class-means centered at their global mean are not only linearly separable , but are actually maximally distant and located on a sphere centered at the origin up to scaling ( i.e. , they form a Simplex Equiangular Tight Frame ( ETF ) – or Simplex ETF , which is formally defined in Definition C.1 in the Appendix ) . • Convergence to Self-duality : the last-layer linear classifiers , living in the dual vector space to that of the class-means , are perfectly matched with their class-means . • Simple Decision Rule : the last-layer classifier is behaviorally equivalent to a Nearest Class-Center decision rule . These results suggest that deep networks are learning maximally separable features between classes , and a max-margin classifier in the last layer upon these learned features , touching the ceiling in terms of the performance . This phenomenon is referred to as Neural Collapse ( NC ) [ 1 ] , and it persists across a range of canonical classification problems , on different neural network architectures ( e.g. , VGG [ 17 ] , ResNet [ 18 ] , and DenseNet [ 19 ] ) and on a variety of standard datasets ( such as MNIST [ 20 ] , CIFAR [ 21 ] , and ImageNet [ 22 ] ) . Fully demystifying the NC phenomenon in theory can be very challenging . Perhaps the most difficult hurdle lies in the nonconvexity of the optimization problem for training neural networks , which , loosely speaking , stems from the nonlinear interaction across many different layers of neural networks . Towards this goal , a recent line of work [ 23–29 ] studied the properties of last-layer classifiers and features based on the assumption of the so-called unconstrained feature model [ 23 ] or layer-peeled model [ 26 ] . At a high level , the unconstrained feature model takes a top-down approach to the analysis of deep neural networks [ 23–26 , 29–31 ] , wherein the last-layer features are modeled as free optimization variables ( hence we call them unconstrained features ) along with the last-layer classifiers ( see Figure 2 for an illustration ) ; this is in contrast to the conventional bottomup approach that studies the problem starting from the input [ 32–42 ] .2 The underlying reasoning is that modern deep networks are often highly overparameterized with the capacity of learning any representations [ 43–46 ] , so that the last-layer features can approximate , or interpolate , any point in the feature space . In this way , the model simplifies the study of last-layer features , enabling us to analyze the interaction between them and the last-layer classifiers . Nonetheless , the simplified unconstrained feature model still leaves us a highly nonconvex training loss to be dealt with . Despite the nonconvexity , recent work [ 23–28 ] studied the global minimizers , proving that Simplex ETFs ( i.e. , NC ) are indeed global solutions to the nonconvex loss . In particular , the work [ 23 , 47 ] studied the training problem with the least-squared loss , proving that the gradient flow converges to NC solutions with extra assumptions . Another line of work [ 24–27 ] considered the commonly used cross-entropy loss for classification , showing that the only global minimizers of the loss function are Simplex ETFs with different constraints on the weights and features.3 However , 2Here , top-down means that we study the network starting from the last-layer down to the input layer , whereas bottom-up refers to an approach from the input up to the last layer . 3The constraints on the features are mainly used to prevent it from approaching infinity since the objective function , with the cross-entropy loss , is not coercive . Note that we still refer to this model as an unconstrained feature model even if they include norm constraints or regularization . these results still suffer from several limitations : ( i ) due to the nonconvex nature , only characterizing optimality conditions is not enough to explain the empirical convergence of iterative algorithms to NC , such as stochastic gradient descent ( SGD ) ; ( ii ) the problem formulations differ from those typically used in practice , which deploy norm regularization ( i.e. , weight decay ) on the weights , rather than enforcing constraints , for the ease of optimization.4 Contributions of This Work . Inspired by these pioneering results [ 1 , 23–26 , 29 ] , in this work we take a step further by characterizing the global optimization landscape of the network training loss based on the unconstrained feature model . Our contributions are summarized as follows . • Benign Global Landscape . For the unconstrained feature model , we provide the first result show- ing that a commonly used , regularized cross-entropy loss is a strict saddle function [ 49–51 ] . In other words , every critical point is either a global solution ( corresponding to Simplex ETFs ) or a strict saddle point5 with negative curvature , so that there is no spurious local minimizer on the optimization landscape . As summarized in Table 1 , this is in contrast to previous work [ 23–26 ] that only characterizes global minimizers . • Efficient , Algorithmic Independent , Global Optimization . The benign global landscape implies that any method that can escape strict saddle points ( e.g . stochastic gradient descent ) converges to a global solution [ 52 ] that exhibits NC . This result supports our empirical observation , as shown in Section 4.1 , that practical overparameterized networks always converge to Simplex ETF solutions with a diverse choice of optimization algorithms . • Cost Reduction for Practical Network Training . Moreover , the universality of NC implies that there is no need of training the last-layer classifiers since the weights can be simply fixed as a Simplex ETF throughout the training process . On the other hand , since NC happens whenever d ≥ K , this implies that we can choose the feature dimension d comparable to the number of classes K , reducing the feature dimension for further computational benefits . In Section 4.3 , our experiments demonstrate that such a strategy achieves on par performance with classical training methods , leading to substantial cost reductions on both memory and computation . Our results shed new light on the question raised in the recent paper [ 53 ] on the role of the optimization strategy ( e.g. , stochastic gradient descent ) for achievingNC in training practical deep networks . This question also relates to the recent highly influential work [ 7 ] on the implicit algorithmic bias . For multi-class classification problems with linearly separable data , this work [ 7 ] showed that linear predictors optimized by gradient descent converge to the max-margin classifiers even without adding any explicit regularization on the cross-entropy loss . Based on this result , a sequence of works [ 54–61 ] laid great emphasis on the notion of “ inductive bias ” of particular optimization algorithms as a reason for the surprising success in training deep learning models.6 In contrast , both our theoretical result on the global landscape for the unconstrained feature model and the empirical evidence on practical deep models demonstrate that NC in network training is facilitated not only by the choice of the optimization methods , but more importantly , by the choice of loss functions and the power of overparameterization in the network architecture . 2 The Problem Setup . A deep neural network is essentially a nonlinear mapping ψ ( · ) : RD 7→ RK , which can be modeled by a composition of simple maps : ψ ( x ) = ψL ◦ · · · ◦ ψ2 ◦ ψ1 ( x ) for x ∈ RD , where ψ ` ( · ) ( 1 ≤ ` ≤ L ) are called “ layers ” . Each layer is composed of an affine transform , represented by 4A recent concurrent work [ 48 ] studied the gradient flow as well as landscape of the cross-entropy loss under the unconstrained feature model , but without using any constraints or regularizers on the weights . 5Throughout the paper , for a minimization problem , we will not distinguish between local maxima and saddle points . We call a critical point strict saddle if the Hessian at this point has at least one negative eigenvalue . 6While ( stochastic ) gradient descent and generic steepest descent methods can converge to max-margin classifiers , it should be noted that other commonly used optimization algorithms in deep learning , such as AdaGrad [ 62 ] and Adam [ 63 ] , do not have max-margin properties in general and their solutions depend upon initialization , step-size and other algorithm hyper-parameters [ 64–66 ] . some weight matrix W ` , and bias b ` , followed by a simple nonlinear7 activation function σ ( · ) . More precisely , a vanilla L-layer neural network can be written as ψΘ ( x ) = WL σ ( WL−1 · · ·σ ( W1x+ b1 ) + bL−1 ) ︸ ︷︷ ︸ φθ ( x ) +bL . ( 1 ) For convenience , we use Θ = { Wk , bk } Lk=1 to denote all the network parameters , and use θ = { Wk , bk } L−1k=1 to denote the network parameters up to the last layer . The output of the penultimate layer , denoted by φθ ( x ) , is usually referred to as the representation or feature of the input x learned from the network . In this way , the function implemented by a neural network classifier can also be expressed as a linear classifier acting upon φθ ( x ) . The goal of deep learning is to fit the parameters Θ so that the output of the model on an input samples x approximates the corresponding output y , i.e . so that ψΘ ( x ) ≈ y , in expectation over a distribution of input-output pairs , D. This can be achieved by optimizing an appropriate loss function L ( ψΘ ( x ) , y ) which quantifies this approximation . In this work , we focus on multi-class classification tasks ( say , with K classes ) , where the class label of a sample x is given by a one-hot vector y ∈ RK representing its membership to one of the K classes . In this setting , cross-entropy is one of the most popular choices for the loss function . Naturally , the distribution D is unknown , but we have access to training samples that are drawn i.i.d . from D. In this way , one can minimize the empirical risk over these samples by optimizing the following problem min Θ K∑ k=1 nk∑ i=1 LCE ( ψΘ ( xk , i ) , yk ) + λ 2 ‖Θ‖2F , ( 2 ) where yk ∈ RK is a one-hot vector with only the kth entry equal to unity ( 1 ≤ k ≤ K ) , { nk } Kk=1 are the numbers of training samples in each class , and λ > 0 is the regularization parameter ( or weight decay penalty ) , and LCE ( · , · ) is the cross-entropy loss . As introduced in Section 1 , recent work [ 1 ] showed that the features learned by minimizing the above objective showcase theNC phenomenon : their within-class variability vanishes , and the features converge to a Simplex ETF .

Summary: The paper studies collapse phenomena of representations from neural classification models trained beyond zero training error with cross entropy. The theoretical contribution of the paper is an analysis of such phenomena in the so called unconstrained feature model. In this setting, the authors derive the global loss minimizers, and show that these can be obtained via gradient based optimization, as the critical points of the loss function are benign. These theoretical results are underpinned by experiments, which investigate the collapse phenomena for a ResNet-18 architecture when trained via different optimization algorithms or with randomly labeled data. The theoretical results also indicate, that in practice, the last linear layer of a neural network can be fixed a priory and of low dimension.