An Unconstrained Layer-Peeled Perspective on Neural Collapse

1 INTRODUCTION . Deep learning has achieved state-of-the-art performance in various applications ( LeCun et al. , 2015 ) , such as computer vision ( Krizhevsky et al. , 2012 ) , natural language processing ( Brown et al. , 2020 ) , and scientific discovery ( Long et al. , 2018 ; Zhang et al. , 2018 ) . Despite the empirical success of deep learning , how gradient descent or its variants lead deep neural networks to be biased towards solutions with good generalization performance on the test set is still a major open question . To develop a theoretical foundation for deep learning , many studies have investigated the implicit bias of gradient descent in different settings ( Li et al. , 2018 ; Amari et al. , 2020 ; Vaswani et al. , 2020 ; Soudry et al. , 2018 ; Lyu & Li , 2019 ; Arora et al. , 2019 ) . It is well acknowledged that well-trained end-to-end deep architectures can effectively extract features relevant to a given label . Although theoretical analysis of deep learning has been successful in recent years ( Arora et al . ; Goldblum et al. , 2019 ) , most of the studies that aim to analyze the properties of the final output function fail to understand the features learned by neural networks . Recently , in Papyan et al . ( 2020 ) , the authors observed that the features in the same class will collapse to their mean and the mean will converge to an equiangular tight frame ( ETF ) during the terminal phase of training , that is , the stage after achieving zero training error . This phenomenon , namely , neural collapse ( Papyan et al. , 2020 ) , provides a clear view of how the last-layer features in the neural network evolve after interpolation and enables us to understand the benefit of training after achieving zero training error to achieve better performance in terms of generalization and robustness . To theoretically analyze the neural collapse phenomenon , Fang et al . ( 2021 ) proposed the layer-peeled model ( LPM ) as a simple surrogate for neural networks , where the last-layer features are modeled as free optimization variables . In particular , in a balanced K-class classification problem using a neural network with d neurons in the last hidden layer , the LPM takes the following form : min W , H 1 n n∑ i=1 L ( Whi , yi ) , s.t . 1 2 ||W ||2F ≤ C1 , 1 2 ||H||2F ≤ C2 , ( 1.1 ) where C1 , C2 are positive constants . Here , W = [ w1 , w2 , · · · , wK ] > ∈ RK×d is the weight of the final linear classifier , H = [ h1 , h2 , · · · , hn ] ∈ Rd×n is the feature of the last layer and yi is the corresponding label . The intuition behind the LPM is that modern deep networks are often highly over-parameterized , with the capacity to learn any representations of the input data . It has been shown that an equiangular tight frame ( ETF ) , i.e. , feature with neural collapse , is the only global optimum of the LPM objective ( 1.1 ) ( Fang et al. , 2021 ; Lu & Steinerberger , 2020 ; Wojtowytsch & E , 2020 ; Zhu et al. , 2021 ) . However , feature constraints in LPMs are not equivalent to weight decay used in practice . In this study , we directly deal with the unconstrained model and show that gradient flow can find those neural collapse solutions without the help of explicit constraints and regularization . To do this , we build a connection between the neural collapse and recent theories on max-margin implicit regularization ( Lyu & Li , 2019 ; Wei et al. , 2018 ) , and use it to provide a convergence result to the first-order stationary point of a minimum-norm separation problem . Furthermore , we illustrate that the crossentropy loss enjoys a benign global landscape where all the critical points are strict saddles in the tangent space , except for the only global minimizers that exhibit the neural collapse phenomenon . Finally , we verify our insights via empirical experiments . In contrast to previous theoretical works on neural collapse , our analysis does not incorporate any explicit regularization or constraint on features . A comparison with other results can be found in Table 1 and we defer a detailed discussion to Section 5.2 . The reasons we investigate the unregularized objective are summarized as follows : 1 . Feature regularization or constrain is still not equivalent to weight decay used in practice . However , previous studies have justified that neural networks continue to perform well without any regularization or constraint ( Zhang et al. , 2021 ) . Moreover , it is proved that SGD with exponential learning rate on unconstrained objective is equivalent to SGD with weight decay . ( Li & Arora , 2019 ) . 2 . As shown in this study , neural collapse exists even under an unconstrained setting , which implies the emergence of neural collapse should be attributed to gradient descent and cross-entropy loss rather than explicit regularization . 3 . Regularization or constraint feature constraint can be barriers for existing theories of neural networks ( Jacot et al. , 2018 ; Lyu & Li , 2019 ) . By allowing features to be totally free , we hope our results can inspire further analysis to plug in a realistic neural network . 1.1 CONTRIBUTIONS . The contributions of the present study can be summarized as follows . • We build a relationship between the max-margin analysis ( Soudry et al. , 2018 ; Nacson et al. , 2019b ; Lyu & Li , 2019 ) and the neural collapse and provide the implicit bias analysis to the feature rather than the output function . Although both parameters and features diverge to infinity , we prove that the convergent direction is along the direction of the minimum-norm separation problem . • Previous works ( Lyu & Li , 2019 ; Ji et al. , 2020 ) only prove that gradient flow on homogeneous neural networks will converge to the KKT point of the corresponding minimum-norm separation problem . However , the minimum-norm separation problem remains a highly non-convex problem and a local KKT point may not be the neural collapse solution . In this study , we perform a more detailed characterization of the convergence direction via landscape analysis . • Previous analysis about neural collapse relies on the explicit regularization or constraint . In this study , we show that the implicit regularization effect of gradient flow is sufficient to lead to a neural collapse solution . The emergence of neural collapse should be attributed to gradient descent and loss function , rather than explicit regularization or constraint . We put detailed discussion in Section 5.2 . 1.2 RELATED WORKS . Implicit Bias of Gradient Descent : To understand how gradient descent or its variants helps deep learning to find solutions with good generalization performance on the test set , a recent line of research have studied the implicit bias of gradient descent in different settings . For example , gradient descent is biased toward solutions with smaller weights under ` 2 loss ( Li et al. , 2018 ; Amari et al. , 2020 ; Vaswani et al. , 2020 ) and will converge to large margin solution while using logistic loss ( Soudry et al. , 2018 ; Nacson et al. , 2019b ; Lyu & Li , 2019 ; Chizat & Bach , 2020 ; Ji et al. , 2020 ) . For linear networks , Arora et al . ( 2019 ) ; Razin & Cohen ( 2020 ) ; Gidel et al . ( 2019 ) have shown that gradient descent determines a low-rank approximation . Loss Landscape Analysis : Although the practical optimization problems encountered in machine learning are often nonconvex , recent works have shown the landscape can enjoy benign properties which allow further analysis . In particular , these landscapes do not exhibit spurious local minimizers or flat saddles and can be easily optimized via gradient-based methods ( Ge et al. , 2015 ) . Examples include phase retrieval ( Sun et al. , 2018 ) , low-rank matrix recovery ( Ge et al. , 2016 ; 2015 ) , dictionary learning ( Sun et al. , 2016 ; Qu et al. , 2019 ; Laurent & Brecht , 2018 ) and blind deconvolution ( Lau et al. , 2019 ) . 2 PRELIMINARIES AND PROBLEM SETUP . In this paper , || · ||F denotes the Frobenius norm , ‖ · ‖2 denotes the matrix spectral norm , ‖ · ‖∗ denotes the nuclear norm , ‖ · ‖ denotes the vector l2 norm and tr ( · ) is the trace of matrices . We use [ K ] : = { 1 , 2 , · · · , K } to denote the set of indices up to K . 2.1 PRELIMINARIES . We consider a balanced dataset with K classes ⋃K k=1 { xk , i } ni=1 . A standard fully connected neural network can be represented as : f ( x ; Wfull ) = bL + WLσ ( bL−1 + WL−1σ ( · · ·σ ( b1 + W1x ) ) ) . ( 2.1 ) Here Wfull = ( W1 , W2 , · · · , WL ) denote the weight matrices in each layer , ( b1 , b2 , · · · , bL ) denote the bias terms , and σ ( · ) denotes the nonlinear activation function , for example , ReLU or sigmoid . Let hk , i = σ ( bL−1 + WL−1σ ( · · ·σ ( b1 + W1xk , i ) ) ) ∈ Rd denote the last layer feature for data xk , i and h̄k = 1n ∑n i=1 hk , i denotes the feature mean within the k-th class . To provide a formal definition of neural collapse , we first introduce the concept of a simplex equiangular tight frame ( ETF ) : Definition 2.1 ( Simplex ETF ) . A symmetric matrix M ∈ RK×K is said to be a simplex equiangular tight frame ( ETF ) if M = √ K K − 1 Q ( IK − 1 K 1K1 > K ) . ( 2.2 ) Where Q ∈ Rd×K is a matrix with orthogonal columns . Let W ∈ RK×d = WL = [ w1 , w2 , · · · , wK ] > be the weight of the final layer classifier , the four criteria of neural collapse can be formulated precisely as : • ( NC1 ) Variability collapse : As training progresses , the within-class variation of the activation becomes negligible as these activation collapse to their class mean h̄k = 1n ∑n i=1 hk , i . ||hk , i − h̄k|| = 0 , 1 ≤ k ≤ K. • ( NC2 ) Convergence to Simplex ETF : The vectors of the class-means ( after centering by their global-mean ) converge to having equal length , forming equal-size angles between any given pair , and being the maximally pairwise-distanced configuration constrained to the previous two properties . cos ( h̄k , h̄j ) = − 1 K − 1 , ||h̄k|| = ||h̄j || , k 6= j . • ( NC3 ) Convergence to self-duality : The linear classifiers and class-means will converge to align with each other , up to appropriate rescaling , that is , there exist a universal constant C > 0 such that wk = Ch̄k , k ∈ [ K ] . • ( NC4 ) Simplification to Nearest Class-Center For a given deepnet activation : h = σ ( bL−1 + WL−1σ ( · · ·σ ( b1 + W1x ) ) ) ∈ Rd , the network classifier converges to choose whichever class has the nearest train class-mean arg min k 〈wk , h〉 → arg min k ∥∥h− h̄k∥∥ , In this paper , we say that a point W ∈ RK×d , H ∈ Rd×nK satisfies the neural collapse conditions or is a neural collapse solution if the above four criteria are all satisfied for ( W , H ) .

Summary: This paper shows that for the cross-entropy loss, when updating both the weight vector and the features by gradient flow, the final solution exhibits the neural collapse behavior, where the neural collapse refers to an observation pointed out by (Papyan 2020), which is that (NC1) the within-class features collapse to their class mean, (NC2) the average of the features of each class converges to having the equal length and all pairs of the class-means have equal-size angles, (NC3) the weight and the class-mean are aligned, (NC4) the weight converges to the nearest class-center of the features. The optimization problems (or the simplified model) in  the previous works that theoretically show the neural collapse require some constraints or regularization, which are seldom used in practice. The optimization model (ULPM) in this paper does not require these, and so is much closer to the realistic (of training deep networks).