Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path

1 INTRODUCTION . Modern deep learning includes paradigmatic procedures that are commonly adopted , but not entirely understood . Some examples are multi-layered architectures , stochastic gradient descent , batch normalization , cross-entropy ( CE ) loss , and training past zero error towards zero loss . Analyzing the properties of these practices is an important research task . In this work , we theoretically investigate behavior of last-layer features in classification deep nets . In particular , we consider the training of deep neural networks on datasets containing images from C different classes with N examples in each class . After passing the i-th example in the c-th class through all layers except the last-layer of the network , the network outputs some last-layer features hi , c ∈ RP . The last-layer of the network—which , for each class c , possesses a classifier wc ∈ RP and bias bc ∈ R—then predicts a label for the example using the rule argmaxc′ ( ⟨wc′ , hi , c⟩+ bc′ ) . The network ’ s performance is evaluated by calculating the error defined by Error = Ave i , c 1 { c ̸= argmax c′ ( ⟨wc′ , hi , c⟩+ bc′ ) } , while the weights , biases , and other parameters of the network ( that determine the behavior of the layers before the last layer ) in the network are updated by minimizing the CE loss defined by CE = −Ave i , c log exp { ⟨wc , hi , c⟩+ bc } ∑C c′=1 exp { ⟨wc′ , hi , c⟩+ bc′ } , ∗Equal contribution . Listed alphabetically . where Ave is the operator that averages over its subscript indices . Prior works such as Zhang et al . ( 2016 ) ; Belkin et al . ( 2019 ) have shown that overparameterized classifiers ( such as deep nets ) can “ memorize ” their training set without harming performance on unseen test data . Moreover , works such as Soudry et al . ( 2018 ) have further shown that continuing to train networks past memorization can still lead to performance improvements1,2 . Papyan , Han , and Donoho ( 2020 ) recently examined this setting , referring to the phase during which one trains past zero-error towards zero-CE-loss as the Terminal Phase of Training ( TPT ) . During TPT , they exposed a phenomenon called Neural Collapse ( NC ) . 1.1 NEURAL COLLAPSE . NC is defined relative to the feature global mean , µG = Ave i , c hi , c , the feature class-means , µc = Ave i hi , c , c = 1 , . . . , C , the feature within-class covariance , ΣW = Ave i , c ( hi , c − µc ) ( hi , c − µc ) ⊤ , ( 1 ) and the feature between-class covariance , ΣB = Ave c ( µc − µG ) ( µc − µG ) ⊤ . It is characterized by the following four limiting behaviors where limits take place with increasing training epoch t : ( NC1 ) Within-class variability collapse3 : Σ†BΣW → 0 , where † denotes the Moore-Penrose pseudoinverse . ( NC2 ) Convergence to Simplex ETF : ⟨µc − µG , µc′ − µG⟩ ∥µc − µG∥2∥µc′ − µG∥2 → { 1 , c = c′ −1 C−1 , c ̸= c ′ ∥µc − µG∥2 − ∥µc′ − µG∥2 → 0 ∀c ̸= c′ ( NC3 ) Convergence to self-duality : wc ∥wc∥2 − µc − µG ∥µc − µG∥2 → 0 ( NC4 ) : Simplification to nearest class center : argmax c′ ⟨wc′ , h⟩+ bc′ → argmin c′ ∥h− µc′∥2 1A similar phenomenon had been previously observed in boosting ( Bartlett et al. , 1998 ) . 2An alternative line of work on “ early stopping ” ( Prechelt , 1998 ; Li et al. , 2020 ; Rice et al. , 2020 ) advocates for terminating the training process early and shows that it could be beneficial when training on noisy data or small datasets . Such datasets are out of the scope of this paper . 3Our characterization of ( NC1 ) is more precise than that given by Papyan , Han , and Donoho ( 2020 ) , which only states ΣW→0 for rhetorical simplicity . The convergence of the trace-norm of Σ†BΣW to zero is the actual quantity measured by Figure 6 of Papyan , Han , and Donoho ( 2020 ) when demonstrating ( NC1 ) . Σ†BΣW The ( NC2 ) property captures convergence to a simple geometric structure called an Equiangular Tight Frame ( ETF ) . An ETF is a collection of vectors { vc } Cc=1 having equal lengths and equal , maximally separated pair-wise angles . In classification deep nets , last-layer features are of higher dimension than the number of classes i.e . P > C . In this setting , the maximal angles are given by ⟨vc , vc′⟩ ∥vc∥2∥vc′∥2 = { 1 , for c = c′ − 1C−1 , for c ̸= c ′ , and the ETF is called a Simplex ETF4 . Observe that as C increases , the ETF approaches a ( partial ) orthogonal matrix . Thus , when C is large , this translates to the intuitive notion that classifiers and class-means tend with training to ( near ) orthogonality . 1.2 DEEP NET CLASSIFICATION WITH MSE LOSS . While classification deep nets are typically trained with CE loss , Demirkaya et al . ( 2020 ) and Hui & Belkin ( 2020 ) recently reported deep nets trained with squared error ( MSE ) loss , L ( W , b , H ) =1 2 Ave i , c ∥Whi , c + b− yi , c∥22 + λ 2 ( ∥W ∥2F + ∥b∥22 ) ( 2 ) = 1 2CN ∥WH + b1⊤CN − Y ∥2F + λ 2 ( ∥W ∥2F + ∥b∥22 ) , achieve comparable test performance as those trained with CE loss . Above , H ∈ RP×CN and Y ∈ RC×CN are matrices resulting from stacking5 together the feature vectors hi , c and one-hot vectors yi , c as their respective columns ; W ∈ RC×P is the matrix resulting from the stacking of classifiers wc as rows ; b ∈ RC is the vector resulting from concatenating the scalars { bc } Cc=1 ; and 1CN is the length-CN vector of ones . Table 1 in Appendix A shows supplementary measurements affirming the findings of Hui & Belkin ( 2020 ) ; Demirkaya et al . ( 2020 ) , i.e . the measurements affirm that MSE-trained networks indeed achieve accuracies on testing data comparable to those of CE-trained networks ( cf . the analogous Table 1 in Papyan , Han , and Donoho ( 2020 ) ) . The analytically-tractable MSE loss offers more mathematical opportunities than the hard-to-analyze CE loss and inspires this paper ’ s main theoretical contributions explicitly characterizing MSE-NC . 1.3 CONTRIBUTIONS . Our main contributions are as follows : • We propose a new decomposition of the MSE loss : L = LNC1 + LNC2/3 + L⊥LS . The terms LNC1 and LNC2/3 possess interpretations attributable to NC phenomena and assume the classifier W is exactly the least-squares optimal classifier WLS ( relative to the given H ) ; and L⊥LS captures the deviation of W from WLS . ( Section 2 ) • We provide empirical measurements of our decomposition ( Figure 2 ) on realistic datasetnetwork combinations showing that L⊥LS becomes negligible during training—leading us to define the central path where L⊥LS=zero . ( Section 2 ) • We reveal a key invariance property on the central path . The invariance motivates the examination of a representative set of features , X=Σ− 1 2 W H , that we call renormalized features . ( Sections 3.1-3.2 ) • We study the gradient flow of those renormalized features along the central path and derive exact , closed-form dynamics that imply NC . The dynamics are explicit in terms of the captures the subtlety that the size of the within-class “ noise ” ( captured by ΣW ) should be viewed relative to the size of the class-means ( captured by ΣB ) . 4Traditional research on ETFs ( cf . Strohmer & Heath Jr ( 2003 ) ; Fickus & Mixon ( 2015 ) ) examines the P≤C setting where C ETF-vectors span their ambient RP space . However , in the P > C setting of classification deep nets , the vectors can not span RP . Following the precedent of Papyan , Han , and Donoho ( 2020 ) , we interpret ETFs as equal-length and maximally-equiangular vectors that are not necessarily spanning . 5Assume the stacking is performed in i-then-c order : the first column is ( i , c ) = ( 1 , 1 ) , the second is ( i , c ) = ( 2 , 1 ) , ... , the N -th is ( i , c ) = ( N , 1 ) , the ( N + 1 ) -st is ( i , c ) = ( 1 , 2 ) , and so on . This matters for formalizing Equation 11 in Corollary 2. singular value decomposition of the renormalized feature class-means at initialization . ( Section 3.3 ) Additionally , we complement this paper with new , extensive measurements on five benchmark datasets—in particular , the MNIST , FashionMNIST , CIFAR10 , SVHN , and STL10 datasets6 ( Deng et al. , 2009 ; Krizhevsky & Hinton , 2009 ; LeCun et al. , 2010 ; Xiao et al. , 2017 ) —and three canonical deep nets—in particular , the VGG , ResNet , and DenseNet networks ( He et al. , 2016 ; Huang et al. , 2017 ; Simonyan & Zisserman , 2014 ) —that verify the empirical reality of MSE-NC i.e . they show ( NC1 ) - ( NC4 ) indeed occur for networks trained with MSE loss . These experiments establish that theoretical modeling of MSE-NC is empirically well-motivated . They are lengthy—together spanning four pages with seven figures and a table—so we collectively defer them to Appendix A . 2 DECOMPOSITION OF MSE LOSS . Inspired by the community ’ s interest in the role of the MSE loss in deep net training ( Demirkaya et al. , 2020 ; Hui & Belkin , 2020 ; Mixon et al. , 2020 ; Poggio & Liao , 2020a ; b ) , we derive a new decomposition of the MSE loss that gives insights into the NC phenomenon . First , absorb the bias vector into the weight matrix—by defining the extended weight matrix W̃ = [ W , b ] ∈ RC× ( P+1 ) and the extended feature vector h̃i , c = [ hi , c ; 1 ] ∈ RP+1—so that Equation 2 can be rewritten as L ( W̃ , H̃ ) = 1 2 Ave i , c ∥W̃ h̃i , c − yi , c∥22 + λ 2 ∥W̃ ∥2F . ( 3 ) Using { h̃i , c } , define the entities µ̃c , µ̃G , H̃ , and Σ̃W analogously to those in Section 1.1 . We further define the extended total covariance and extended class-means matrices , respectively , as Σ̃T =Ave i , c ( h̃i , c − µ̃G ) ( h̃i , c − µ̃G ) ⊤ ∈ R ( P+1 ) × ( P+1 ) M̃ = [ µ̃1 , . . . , µ̃C ] ∈ R ( P+1 ) ×C . Next , we reformulate , with weight decay incorporated , a classic result of Webb & Lowe ( 1990 ) : Proposition 1 ( Webb & Lowe ( 1990 ) with Weight Decay ) . For fixed extended features H̃ , the optimal classifier minimizing the MSE loss L ( W̃ , H̃ ) is W̃LS = 1 C M̃⊤ ( Σ̃T + λI ) −1 , where I is the identity matrix . Note that W̃LS depends on H̃ only . Note that we can interpret W̃LS as a function of H̃ , leading us to identify the following decomposition of L where one term depends on H̃ only . Theorem 1 . ( Decomposition of MSE Loss ; Proof in Appendix B ) The MSE loss , L ( W̃ , H̃ ) , can be decomposed into two terms , L ( W̃ , H̃ ) = LLS ( H̃ ) + L⊥LS ( W̃ , H̃ ) , where LLS ( H̃ ) = 1 2 Ave i , c ∥W̃LSh̃i , c − yi , c∥22 + λ 2 ∥W̃LS∥2F , and L⊥LS ( W̃ , H̃ ) = 1 2 tr { ( W̃ − W̃LS ) ( Σ̃T + µ̃Gµ̃ ⊤ G + λI ) ( W̃ − W̃LS ) ⊤ } . In the above , LLS ( H̃ ) is independent of W̃ . Intuitively , it is the MSE-performance of the optimal classifiers W̃LS ( rather than the “ real classifiers ” W̃ ) on input H̃ . 6CIFAR100 and ImageNet are omitted because Demirkaya et al . ( 2020 ) ; Hui & Belkin ( 2020 ) showed that they require an additional scaling-heuristic on top of the traditional MSE loss . Rigorous investigation into this scaling-heuristic would have demanded mathematical analysis outside the scope of this paper as well as experimentation that would have consumed expensive amounts of computational resources . The component L⊥LS ( W̃ , H̃ ) is non-negative7 and is zero only when W̃ = W̃LS . Therefore , L⊥LS ( H̃ ) quantifies the distance of W̃ from W̃LS . In short , the least-squares component , LLS ( H̃ ) , captures the behavior of the network when the classifier possesses optimal , least squares behavior . Meanwhile , the deviation component , L⊥LS ( W̃ , H̃ ) , captures the divergence from that behavior . We can further decompose LLS ( H̃ ) into two terms , one capturing activation collapse ( NC1 ) and the other capturing convergence to Simplex ETF of both features and classifiers ( ( NC2 ) and ( NC3 ) ) . Theorem 2 . ( Decomposition of Least-Squares Component ; Proof in Appendix C ) The least-squares component , LLS ( H̃ ) , of the MSE decomposition in Theorem 1 can be further decomposed into LLS ( H̃ ) = LNC1 ( H̃ ) + LNC2/3 ( H̃ ) , where LNC1 ( H̃ ) = 1 2 tr { W̃LS [ Σ̃W + λI ] W̃⊤LS } , LNC2/3 ( H̃ ) = 1 2C ∥W̃LSM̃ − I∥2F . Inspection of these terms is revealing . First , observe that LNC2/3 is a function of the class-means and MSE-optimal classifiers . Minimizing LNC2/3 will push the ( unextended ) class-means and classifiers towards the same Simpex ETF matrix8 i.e . ( NC2 ) - ( NC3 ) . Next , note that the within-class variation is independent of the means . Thus , despite the fact that classifiers are converging towards some ( potentially large ) ETF matrix , we can always reduce LNC1 by pushing ΣW towards zero , which corresponds to ( NC1 ) . Figure 2 measures the empirical values of the above decomposition terms on five canonical datasets and three prototypical networks . It shows that L⊥LS ( W̃ , H̃ ) becomes negligible compared to LLS ( H̃ ) when training canonical deep nets on benchmark datasets . In other words , L ( W̃ , H̃ ) ≈LLS ( H̃ ) starting from an early epoch in training and persists into TPT . This motivates us to formulate the central path : P = { ( W̃LS ( H̃ ) , H̃ ) | H̃ ∈ R ( P+1 ) ×CN } , ( 4 ) where the notation W̃LS ( · ) makes explicit the fact that W̃LS is a function of H̃ only . Intuitively , for a classifier-features pair to lie on the central path , i.e . ( W̃ , H̃ ) ∈ P , means that the “ real classifier ” W̃ will exactly equal W̃LS ( H̃ ) —the optimal classifier that would result from fixing H̃ and minimizing L w.r.t . just the classifier . Combined with Theorem 1 , we see ( W̃ , H̃ ) ∈ P if and only if L ( W̃ , H̃ ) =LLS ( H̃ ) . Figure 2 shows that classifier-features pairs lie approximately on the central path during TPT , allowing us to shift focus from L to LLS .

Summary: Recently, Papyan, Han and Donoho (2020) have observed that training neural networks beyond zero training error leads to simplex arrangements of the features. The submission studies this phenomenon, called neural collapse, from a theoretical perspective, when training is done via minimizing the square loss.  In particular, it is shown that the square loss can be split, such that one summand corresponds to the quality of the features (quantified by the loss of the MSE classifier on them) and that the other summand corresponds to the quality of the classification layer (quantified by the its deviationfrom the least squares classifier). Furthermore, the first summand quantifies the closeness of the features to neural collapse. In a second step, the authors investigate the gradient flow induced from the first summand and derive closed form solutions.