How Gradient Descent Separates Data with Neural Collapse: A Layer-Peeled Perspective

In this paper , we study the inductive bias of the neural features and parameters1 from neural networks with cross-entropy loss . We study a surrogate model named2 unconstrained layer peeled model ( ULPM ) , which helps us to illustrate that the3 features and classifiers in the last layer of the neural network will converge to4 a certain neural collapse structure [ 28 ] , where the cross-example within-class5 variability of the last-layer features collapse to zero and the class-means converge6 to a Simplex Equiangular Tight Frame ( ETF ) . We illustrate that the ULPM with7 cross-entropy loss enjoys a benign global landscape on this model where all the8 critical points are strict saddle points except the only global minimizers which9 exhibit neural collapse phenomenon . Empirically we show that our results also10 hold during the training of neural networks in real world tasks when explicit11 regularization or weight decay is not included.12 1 Introduction13 Deep learning has achieved state-of-the-art per-14 formances in various applications [ 20 ] , from15 computer vision [ 16 ] , to natural language16 processing [ 6 ] and even scientific discovery [ 23,17 41 ] . Despite the empirical successes achieved,18 how gradient descent or its variants leads deep19 neural networks to be biased towards solutions20 with good generalization performance on the21 test set is still a major open question . To de-22 velop a theoretical foundation for deep learn-23 ing , many works have studied the implicit24 bias of gradient descent in different settings25 [ 21 , 1 , 37 , 33 , 25 , 3 ] .26 It is well-acknowledged that well-trained end-27 to-end deep architectures have the ability to ef-28 fectively extract features relevant to the given label . Although theoretical analysis of deep learning29 has several achievements in recent years [ 2 , 13 ] , most of the works that aim to analyze properties30 of the final output function fail to understand the feature learned . Recently in [ 28 ] , authors observe31 that the within-class cross-sample features will collapse to the mean and the mean will converge32 to an Equiangular Tight Frame ( ETF ) during the terminal phase of training , i.e . after achieving33 zero training error and interpolating the in-sample training data . Such phenomenon , namely Neural34 Collapse ( NC ) [ 28 ] , provides a clear view of how the last layer features in the neural network involve35 after interpolation and enables us to understand the benefit of training after achieving zero training36 error to achieve better properties in generalization and robustness . To theoretically analyze the neuron37 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . collapse phenomenon , [ 9 , 24 , 39 ] propose the Layer-Peeled Model ( LPM ) as a simplification for38 neural networks , where the last-layer features are modeled as free optimization variables . In particular,39 in a K-class classification problem using a neural network with d neurons in the last hidden layer , a40 corresponding class of LPMs can be defined through the form41 min W , H 1 N N∑ i=1 L ( Whi , yi ) s.t . 1 2 ||W ||2F ≤ C1 , 1 2 ||H||2F ≤ C2 ( 1 ) for some positive constant C1 , C2 . Here W = [ w1 , w2 , · · · , wK ] > ∈ RK×d is the weight of the42 final linear classifier , H = [ h1 , h2 , · · · , hN ] ∈ Rd×N is the feature of the last layer and yi is the43 corresponding label . The intuition behind LPM is that the modern deep networks are often highly44 over-parameterized , with the capacity to learn any representations of the input data . It has been shown45 that equiangular tight frame ( ETF ) , i.e . feature with neural collapse , is the only global optimum46 of the LPM objective ( 1 ) [ 9 , 24 , 39 ] . However , even for this simplified model , the non-convexity47 nature of it makes the analysis highly non-trivial . In this paper we aim to understand how gradient48 descent separates data with neural collapse . To do this , we build a connection between the neural49 collapse with the recently proposed normalized margin [ 25 , 38 ] . In [ 25 ] , the authors shows that , using50 gradient descent , the direction of the weight converges to the direction that maximizes the ` 2-margin51 of the data while the norm of the weight diverges to +∞ in homogeneous neural networks . Based on52 these results , we introduce neural collapse margin and use it provide a convergence result to the first53 order stationary point of the minimum-norm separation problem . Furthermore , we illustrate that the54 cross-entropy loss enjoys a benign global landscape where all the critical points are strict saddles55 in the tangent space except the only global minimizers which exhibit neural collapse phenomenon.56 The analysis provides insights on how gradient descent separates data during the training of neural57 networks with neural collapse and the benefit of training after interpolation on generalization and58 robustness . We verify our insights via empirical experiments . 59 Besides , [ 26 ] and a concurrent paper [ 43 ] also provide landscape and optimization analysis to study60 neural collapse phenomenon , we summarize the connection and difference with our paper in Table 1.61 Our result doesn ’ t introduce any extra feature norm constraint or feature norm regularization , which62 are not commonly used in the realistic deep learning . We put the detailed discussion in Section 5.2.63 1.1 Contribution64 We summarize our contribution as follows.65 • We build a relationship between the max-margin analysis [ 33 , 27 , 25 ] with the neural66 collapse and provide the inductive bias analysis to the feature rather than the output function.67 • Previous works only prove that Gradient Descent on homogeneous neural networks will68 converge to the KKT point of the corresponding minimum-norm separation problem . How-69 ever , the minimum-norm separation problem is still a highly non-convex problem . In this70 paper , we prove that the ULPM cases enjoys a benign landscape and characterize the neural71 collapse property of the global minimizer.72 • We show that although the gradient descent on cross entropy loss will push the parameters to73 infinity , the landscape in the tangent space has no spurious minimum thus many optimization74 algorithms will converge only along the neural collapse directions .75 1.2 Related Work76 Inductive Bias of Gradient Descent : To understand how gradient or its variants descent helps77 deep learning to find solutions with good generalization performance on the test set . A recent line of78 research have studied the implicit bias of gradient descent in different settings . As example , gradient79 descent is biased towards model have smaller weight [ 21 , 1 , 37 ] and will converge to large margin80 solution [ 33 , 27 , 25 , 7 , 14 ] while using logistic loss . For linear networks , [ 3 , 31 , 12 ] have shown that81 gradient descent will find out a low rank approximation.82 Loss Landscape Analysis : Although the practical optimization problems encountered in machine83 learning are often nonconvex , recent works have shown that critical points other than the good ones84 always lies in the balanced superpositions of symmetric copies of the ground truth according to85 the hidden symmetries in the objective function [ 34 , 42 ] which leads to a benign global landscape.86 In particular , these landscapes do not exhibit spurious local minimizers or flat saddles and can be87 optimized easily via gradient based methods [ 10 ] . The examples including phase retrieval [ 36 ] ,88 low-rank matrix recovery [ 11 , 10 ] , dictionary learning [ 35 , 30 , 19 ] , blind deconvolution [ 18 ] .89 2 Preliminaries and Problem Setup90 2.1 Preliminaries91 We consider a dataset with K classes : ⋃K k=1 { xk , i } nk i=1 . For simplicity , we assume the dataset is92 balanced , i.e . n1 = · · · = nK = n. A standard fully connected neural network can be represented as:93 94 f ( x ; W full ) = bL + WLσ ( bL−1 + WL−1σ ( · · ·σ ( b1 + W 1x ) ) ) . ( 2 ) Here W full = ( W 1 , W 2 , · · · , WL ) denote the weight matrices in each layer and ( b1 , b2 , · · · , bL ) 95 are the bias terms , σ ( · ) stands for the nonlinear activation function , for example , ReLU or sigmoid.96 Let hk , i = σ ( bL−1 + WL−1σ ( · · ·σ ( b1 + W 1xk , i ) ) ) ∈ Rd denote the last layer feature for data97 xk , i and h̄k = 1n ∑n i=1 hk , i the feature mean within in the k-th class . Without loss of generality , we98 can absorb the bias term into the weight matrix by adding a scalar into each feature vectors , so we99 will ignore the bias term in the following analysis . Let W ∈ RK×d = WL = [ w1 , w2 , · · · , wK ] > 100 be the weight of the final linear classifier . Neural collapse is the phenomenon that the final layer101 feature will convergence to a simplex equiangular tight frame ( ETF ) :102 Definition 2.1 . A symmetric matrix M ∈ RK×K is said to be simplex equiangular tight frame103 ( ETF ) if104 M = √ K K − 1 Q ( IK − 1 K 1K1 > K ) . ( 3 ) Where Q ∈ RK×K is an orthogonal matrix.105 The four criteria of neural collapse can be formulated precisely as106 • ( NC1 ) Variability collapse : As training progresses , the within-class variation of the activation becomes negligible as these activation collapse to their class-means 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-sized 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 each other , up to rescaling . ∃C s.t . wk = Ch̄k , ∀1 ≤ k ≤ K • ( NC4 ) Simplification to Nearest Class-Center For a given deepnet activation h = σ ( bL−1 + WL−1σ ( · · ·σ ( b1 + W 1x ) ) ) ∈ 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 a point W ∈ RK×d , H ∈ Rd×nK satisfies neural collapse conditions or is107 neural collapse solution if these four criteria are all satisfied for ( W , H ) .108 2.2 Problem Setup109 In this paper , we mainly focus on the neural collapse phenomenon , which is only related to the110 classifiers and features in the last layer . Since general analysis on the highly non-smooth and non-111 convex neural network is difficult , here we peel down the last layer of neural network and propose112 the following Unconstrained Layer-Peeled Model ( ULPM ) as a simplification to capture the main113 characteristic related to neural collapse during the training dynamics . Similar simplification is114 common used in previous theoretical works [ 24 , 9 , 39 , 43 ] , but ours don ’ t have any constraint or115 regularization on features and stands closer to realistic neural network models . We need to mention116 that although [ 26 ] also study the unconstrained model , their analysis is highly dependent on the ` 2117 loss function which is rarely used in classification task while ours can address the most popular cross118 entropy loss.119 Let W = [ w1 , w2 , · · · , wK ] > ∈ RK×d and H = [ h1,1 , · · · , h1 , N , h2,1 , · · · , hK , N ] ∈ Rd×KN120 be the matrices of classifiers and features in the last layer , where K is the number of classes and N121 is the number of data points in each classes . The Unconstrained Layer-Peeled Model is defined as122 following:123 min W , H L ( W , H ) = − K∑ k=1 n∑ i=1 log ( exp ( w > k hk , i ) ∑K j=1 exp ( w > j hk , i ) ) ( 4 ) Here we do not have any constrain or regularization on features , which corresponds to the absence124 of weight decay in deep learning training . The objective function ( 4 ) is generally non-convex on125 ( W , H ) and we aim to study the landscape of the objective function ( 4 ) . Furthermore , we consider126 the gradient flow of the the objective function127 dW ( t ) dt = ∂L ( W ( t ) , H ( t ) ) ∂W , dH dt = ∂L ( W ( t ) , H ( t ) ) ∂H . We also trace the the dynamic of the loss function L ( t ) : = L ( W ( t ) , H ( t ) ) and study the convergence128 of ( W ( t ) , H ( t ) ) .129 Notations . We denote || · ||F the Frobenius norm , ‖ · ‖2 the matrix spectral norm , ‖ · ‖∗ the nuclear130 norm , ‖ · ‖ the vector l2 norm and tr ( · ) the trace of matrices . We use [ K ] : = { 1 , 2 , · · · , K } to denote131 the set of indices up to K.132 3 Main Results133 In this section , we present our main results about the training dynamics and landscape analysis about134 ( 4 ) . We organize the section as follows : First in Section 3.1.1 , we show the relationship between135 margin and neural collapse in our surrogate model . Inspired by this relationship , we propose a136 minimum-norm separation problem ( 5 ) and show the connection between the convergence direction137 of gradient flow and the KKT point of ( 5 ) . In addition , we explicitly solve the global optimum of138 ( 5 ) and show it must satisfy neural collapse conditions . However , due to the non-convexity , we find139 an Example 3.1 in Section 3.2 which shows that there exist some bad KKT points such that simple140 gradient flow will get stuck in them and not converge to neural collapse solution which is proved141 to be optimal in Theorem 3.3 . Then we present our second–order analysis result in Theorem 3.4 to142 show that those bad points will exhibit decreasing directions in the tangent space thus if we add some143 noise in the training algorithm ( e.g . use stochastic gradient descent ) , our algorithm can escape from144 those directions and can only converge to the neural collapse solutions.145 3.1 Convergence To The First–Order Stationary Point146 3.1.1 Neural Collapse Margin147 Before we state our convergence result , let ’ s first discuss the relationship between margin and neural148 collapse . By building the relationship between them we can have a better intuition about why gradient149 flow can converge to neural collapse solution since the convergence to max-margin solutions has been150 studied in many literature [ 21 , 25 , 1 , 37 ] . Recall the margin of a single data point xk , i and associated151 feature hk , i as qk , i ( W , H ) : = w > k hk , i−maxj 6=k w > j hk , i . [ 5 , 4 ] . To bridge the margin theory with152 neural collapse phenomenon , we define the following neural collapse margin:153 Definition 3.1 . We define the the Neural Collapse Margin for the entire dataset as qmin ( W , H ) =154 mink∈ [ 1 , K ] , i∈ [ 1 , n ] qk , i ( W , H ) .155 The following lemma shows that the neural collapse margin is an indicator of the neural collapse156 phenomenon in the sense that collapsed margin minimize the neural collapse margin . Thus we can157 trace the neural collapse margin to study the convergence to the neural collapse solution.158 Lemma 3.1 ( Neural Collapse Margin as an Indicator of Neural Collapse ) . The neural collapse margin always smaller than qmin ( W , H ) ≤ ‖W ‖2F + ‖H‖2F 2 ( K − 1 ) √ n and ( W , H ) must satisfies the neural collapse conditions when the inequality above is reduced to an159 equality.160 3.1.2 Convergence Results161 Now we present our result about the convergence of gradient flow on the ULPM ( 4 ) . Following [ 25 ] ,162 we link gradient flow on cross-entropy loss with a minimum-norm separation problem.163 Theorem 3.1 . For problem ( 4 ) , let ( W ( t ) , H ( t ) ) be the path of gradient flow at time t , if there164 exist a time t0 such that LCE ( W ( t0 ) , H ( t0 ) ) < log 2 , then any limit point of { ( Ĥ ( t ) , Ŵ ( t ) ) : =165 ( H ( t ) √ ‖W ( t ) ‖22+‖H ( t ) ‖22 , W ( t ) √ ‖W ( t ) ‖22+‖H ( t ) ‖22 ) } is along the direction of an Karush-Kuhn-Tucker ( KKT ) 166 point of the following minimum-norm separation problem:167 min W , H 1 2 ||W ||2F + 1 2 ||H||2F s.t . ∀k 6= j ∈ [ K ] , i ∈ [ n ] , w > k hk , i −w > j hk , i ≥ 1 . ( 5 ) Remark 3.1 . Indeed , the problem ( 5 ) can be reorganized to maximize neural collapse margin such168 that the norm is constrained to be lower than a certain value . The proof is as follows , for all feasible169 solutions ( W , H ) , we can find that ∀α ≥ qmin ( W , H ) −1/2 , α ( W , H ) are still feasible thus the170 minimum objective value is 1 2 ||W || 2 F+ 1 2 ||H|| 2 F qmin ( W , H ) 1/2 along the direction of ( W , H ) . Then take minimum171 among all the directions we can find the minimum is attained if and only if ( W , H ) attains the172 maximum neural collapse margin on the sphere { ( W , H ) : ||W ||2F + ||H||2F ≤ C } 173 The Theorem 3.1 indicates that the convergent direction of gradient flow is restricted to those174 max-margin directions , which usually enjoy some good properties on robustness or generalization175 performance . Generally speaking , the KKT conditions are not sufficient to obtain global optimality176 since the minimum-norm separation problem ( 5 ) is non-convex . Moreover , in some certain occasions,177 KKT conditions may be even not necessary for global optimum . However , we can have a precise178 characterization about the optimum from another perspective , the following result shows that the179 global optimum of this problem satisfies neural collapse conditions.180 Theorem 3.2 . Every global optimum of the minimum-norm separation problem ( 5 ) is also a KKT181 point and it satisfies the neural collapse conditions.182 To illustrate how does ( 5 ) related to ( 4 ) and gain insight about Theorem 3.1 , we provided the following183 lemmas to show that when t is sufficient large , the ( W ( t ) , H ( t ) ) is an ( , δ ) approximate KKT point184 after appropriate scaling , where the ( , δ ) converges to zero when t→∞ . Then as shown in [ 8 ] we185 know that the limit of these ( , δ ) approximate KKT point is exact KKT point . Detailed definition of186 KKT points and approximate KKT points can be found in appendix.187 Lemma 3.2 . If there exist a time t0 such that L ( W ( t0 ) , H ( t0 ) ) < log 2 , then for any t > t0 ( ˜W ( t ) , ˜H ( t ) ) : = ( W ( t ) , H ( t ) ) /qmin ( W ( t ) , H ( t ) ) 1/2 is a ( , δ ) - approximate KKT point of the following minimum-norm separation problem . More precisely , we have = √ 2 ( 1− β ( t ) ) C , δ = K 2Cqmin ( t ) where : β = tr ( W > ∇WL ( W , H ) ) + tr ( H > ∇HL ( W , H ) ) √ ||W ||2F ||+ ||H||2F √ ||∇WL ( W , H ) ||2F ||+ ||∇HL ( W , H ) ||2F is the angle between ( W , H ) and its corresponding gradient and C is a positive constant.188 Lemma 3.3 . If there exist a time t0 such that LCE ( W ( t0 ) , H ( t0 ) ) < log 2 , then we have:189 β ( t ) → 1 , qmin ( t ) →∞ as t→∞ ( 6 ) which implies that → 0 and δ → 0 when time t goes to infinity.190 3.2 Second–Order Landscape Analysis191 Due to the non-convex nature of the objective ( 4 ) , we can ’ t achieve such global solution efficiently.192 The global optimality condition shown in Theorem 3.2 still can ’ t guarantee convergence to neural193 collapse . In this section , we aim to show that this non-convex optimization problem is actually not194 scary.195 Different from previous landscape analysis of non-convex problem , where people aim to show that196 the objective has a negative directional curvature around any stationary point [ 34 , 42 ] , once features197 can be perfectly separated , the ULPM objective ( 4 ) will always decrease along the direction of the198 current point and the optimum is attained only in infinity . Although growing along all of those199 perfectly separation directions can let the loss function decreasing to 0 , the speed of decreasing are200 quite different and there exists an optimal direction with fastest decreasing speed . However , simple201 first–order analysis may fail to interpret how does gradient flow move among these directions and we202 need second–order analysis to help us fully characterize the realistic training dynamics . Here is an203 example illustrating our motivation.204 Example 3.1 ( A Motivating Example ) . Consider the case when K = 4 , n = 1 , let ( W , H ) be the205 following point:206 W = H = C 1 −1 0 0−1 1 0 00 0 1 −1 0 0 −1 1 ( 7 ) One can easily verify that this ( W , H ) enables our model to classify all of the features perfectly.207 Further more , we can show it is along the direction of a KKT point of the minimum-norm separation208 problem ( 5 ) by construct the Lagrangian multiplier Λ = ( λij ) Ki , j=1 as following:209 Λ = 0 0 12 1 2 0 0 12 1 2 1 2 1 2 0 0 1 2 1 2 0 0 ( 8 ) And the gradient of ( W , H ) is210 ∇WL ( W , H ) = ∇HL ( W , H ) = −C 2 + 2e−2C 2 2 + 2e−2C2 + 2e2C2 1 −1 0 0−1 1 0 00 0 1 −1 0 0 −1 1 ( 9 ) We can find that the directions of gradient and the parameter align with each other ( i.e.211 W //∇WL ( W , H ) , H//∇HL ( W , H ) ) , which implies simple gradient descent get stuck in this212 direction and only grow the parameter norm . However , if we construct:213 W ′ = H ′ = C 1 α β βα 1 β ββ β 1 α β β α 1 , α2 + 2β2 = 1 , α < 0 , β < 0 ( 10 ) Then ∀ > 0 , we can choose appropriate α , β such that ( see detailed computation in Appendix ) :214 ||W ′||2F = ||W ||2F , ||H ′||2F = ||H||2F , ||W ′ −W ||2F + ||H ′ −H||2F < , L ( W ′ , H ′ ) ≤ L ( W , H ) ( 11 ) The results in ( 11 ) indicate that ( W ′ , H ′ ) is a saddle point on the sphere and there exists many better215 direction ( W ′ , H ′ ) staying very close to the original direction ( W , H ) . Although simple gradient216 descent will always move along the original direction , once we add some noise in the training ( e.g.217 stochastic gradient descent ) , the optimization algorithm can find this better direction and escape the218 original bad direction.219 In Example 3.1 , we show that there does exist some suboptimal KKT point of the minimum-norm220 separation problem ( 5 ) , but there also exist some better points close to it thus stochastic gradient221 method can easily escape form them . In the following theorem , we will show that the best directions222 are neural collapse solutions in the sense that the loss function is lowest among all the growing223 directions.224 Theorem 3.3 . The optimal value of loss function ( 4 ) on a sphere is attained ( i.e . L ( W , H ) ≤225 L ( W ′ , H ′ ) , ∀||W ′||2F +||H ′||2F = ||W ||2F +||H||2F ) if only if the ( W , H ) satisfies neural collapse226 conditions and ||W ||F = ||H||F .227 Remark 3.2 . Note that the second conditions is necessary since neural collapse conditions don ’ t228 specify the norm ratio of W and H . That is , if ( W , H ) satisfies neural collapse conditions,229 ( αW , βH ) , ∀α , β ∈ R will also satisfies them but only some certain α , β are optimal.230 Now we turns to those points that don ’ t satisfy neural collapse conditions . To formalize our discussion231 in the motivating Example 3.1 , we first introduce the tangent space:232 Definition 3.2 ( tangent space ) . The tangent space of ( W , H ) is defined to be a set of directions that233 are orthogonal to ( W , H ) :234 T ( W , H ) = { ∆W ∈ RK×d , ∆H ∈ Rd×nK ) : tr ( W > ∆W ) + tr ( H > ∆H ) = 0 } ( 12 ) Our next result justify our observation in the Example 3.1 that for every suboptimal points , there exist235 a direction in the tangent space such that move along this direction will leads to a lower objective236 value.237 Theorem 3.4 . If ( W , H ) is not the optimal solutions in Theorem 3.3 , then ∃ ( ∆W , ∆H ) ∈238 T ( W , H ) , M > 0 such that239 ∀0 < δ < M , L ( W + δ∆W , H + δ∆H ) ≤ L ( W , H ) ( 13 ) . Further more , it implies that ∀ > 0 , ∃ ( W ′ , H ′ ) such that:240 ||W ′||2F + ||H ′||2F = ||W ||2F + ||H||2F , ||W ′ −W ||2F + ||H ′ −H||2F < , L ( W ′ , H ′ ) ≤ L ( W , H ) ( 14 ) Remark 3.3 . The result in ( 13 ) give us a decreasing direction orthogonal to the direction of ( W , H ) ,241 as shown in Example 3.1 , the gradient might be parallel to ( W , H ) , the decreasing direction must be242 obtained by analyze the Hessian matrices and it further indicates that these points are exactly saddle243 points in the tangent space , a formal statement and definition can be found in appendix . For a large244 family of stochastic optimization algorithm , the projection of noise onto this decreasing direction245 is not zero with probability 1 , so its those algorithms will escape the bad point and no longer move246 along this direction within a small number of iterations.247 4 Empirical Results248 Gradient Descent on the ULPM Objective . We first conduct experiments on the ULPM objective249 ( 4 ) to support the results of convergence towards Neural Collapse in our theories . We set N = 10,250 K = 5 , d = 20 and use gradient descent with learning rate 5 to run 105 epochs . We characterize251 the dynamics of the training procedure in Figure 2 , through four aspects : ( 1 ) variation of the252 centered class-mean features ’ norms ( i.e. , Std ( ‖h̄k − h̄‖ ) /Avg ( ‖h̄k − h̄‖ ) ) and the variation of the253 classifier ’ s norms ( i.e. , Std ( ‖w̄k‖ ) /Avg ( ‖w̄k‖ ) ) . ( 2 ) Within-class variation of last layer features254 ( i.e. , Avg ( ‖hk , i − hk‖ ) /Avg ( ‖hk , i − h̄‖ ) ) . ( 3 ) The cosines between pairs of last layer features ( i.e.,255 Avg ( | cos ( h̄k , h̄k′ ) +1/ ( K−1 ) | ) ) and that of the classifiers ( i.e. , Avg ( | cos ( w̄k , w̄k′ ) +1/ ( K−1 ) | ) ) .256 ( 4 ) The distance between normalized centered classifier and normalized last layer feature ( i.e.,257 Avg ( | ( h̄k − h̄ ) /‖h̄k − h̄‖ − w̄k/‖w̄k‖| ) ) . Empirically we observe that logarithm of the two258 variations of norms ( in the first aspect ) decrease approximately at rate O ( 1/ ( log ( t ) ) ) , and the259 remaining quantities decrease approximately at rate O ( 1/ ( log ( t ) ) ) .260 Realistic Training . We also extend our theory to realistic neural network training on benchmark261 dataset . To evaluate our theory , we train the VGG-13 [ 32 ] on FashionMNIST [ 40 ] without weight262 decay and track the convergence speed of the last layer feature to the neural collapse solution every few263 epochs to see how it changes during the terminal phase training . Observe that all the aforementioned264 quantities either decrease or stay in small values during the training process , providing implications265 that neural collapse can occur with sufficient training epochs.266 267 5 Conclusion and Discussion268 5.1 Conclusion269 To understand the inductive bias of neural feature from gradient descent training , we build a connection270 between large margin inductive bias with neural collapse phenomenon and study a unconstrained271 layer-peeled model in this paper . We proved that the gradient flow of the ULPM convergences272 to KKT point of a minimum-norm separation problem where the global optimum satisfies neural273 collapse conditions . Although the ULPM is nonconvex , we show that ULPM have a nice landscape274 where all the stationary point is a strict saddle point in the tangent space except the global neural275 collapse solution . Our study helps to demystify the neural collapse phenomenon , which shed light on276 the generalization and robustness properties during the terminal phase of training deep networks in277 classification problems.278 5.2 Relationship with Other Results279 Theoretical analysis of neural collapse are first provided by [ 24 , 39 , 9 ] , they show that the neural280 collapse solution is the only global minimum of the simplified non-convex objective function . In281 particular , [ 39 , 24 ] study a continuous integral form of the loss function and show that the feature282 learnt should be a uniform distribution on sphere . A more realistic discrete setting are studied in283 [ 9 ] , where the constraint is on the whole feature matrix rather than individual features . All these284 results only relies on Jensen inequality on output logits thus can be generalized to other convex in285 logit losses . Our result utilize the implicit bias of the exponential like loss function to remove the286 feature norm constraint which is not practicable in real applications.287 Though the global optimum shares good property [ 9 ] , the ULPM objective is still highly non-convex.288 Regards optimization , [ 26 , 29 ] analyze the unconstrained feature model with ` 2 loss and establish289 convergence results to collapsed feature for gradient descent . However they fail to generalize on290 other more practical loss functions used in classification tasks . The analysis highly relies on the ` 2291 loss which turns the training dynamic to an ODE in eigenvalues.292 The most relevant paper is a concurrent breakthrough work [ 43 ] , which provide a landscape analysis293 about the regularized unconstrained feature model . [ 43 ] turns the feature norm constraint in [ 9 ] into294 feature norm regularization and still preserves the neural collapse global optimum . At the same295 time , [ 43 ] also show that the modified regularized objective shares a benign landscape , where all296 the critical points are strict saddles except the global one . Although our paper and [ 43 ] discover297 similar landscape results , we believe our characterization stays closer to the real algorithms used in298 the following two ways299 • The same as [ 24 , 39 , 9 ] , [ 43 ] only utilize the convexity in logits of the loss function . However,300 our analysis also explores the exponential-like property of the cross-entropy loss which will301 enlarge the norm of the feature . The large feature will provide better approximation to the302 true neural collapse problem of the normalized feature via approximating the max function303 via gradually scaled exponential function.304 • We doesn ’ t introduce any constraints or regularization on the feature norm , which is not305 applied in the realist training . Regularization on feature introduce in [ 43 ] is still different306 from the weight decay regularization [ 17 ] . However weight decay on homogeneous neural307 network is equivalent to gradient descent with scaling step size on unregularized objective308 [ 22 , 41 ] .309 We summarize analysis of neural collapse in Table 1.310 5.3 Limitation and Future Work311 The convergence to neural collapse is super slow . [ 15 ] provide a loss dependent learning rate schedule312 and leads to O ( 1/t ) convergence rate for linear regression . It ’ s interesting to investigate can this313 methodology being generalized to our setting . On the other hand , although we have shown that the314 ULPM have a nice landscape , we still leave the global convergence of ( stochastic ) gradient descent315 as future work for we want to provide global convergence of gradient descent combined with a plug316 in feature extractor.317 References318 [ 1 ] Shun-ichi Amari , Jimmy Ba , Roger Grosse , Xuechen Li , Atsushi Nitanda , Taiji Suzuki , Denny319 Wu , and Ji Xu . When does preconditioning help or hurt generalization ? arXiv preprint320 arXiv:2006.10732 , 2020.321 [ 2 ] Raman Arora , Sanjeev Arora , Joan Bruna , Nadav Cohen , Rong Ge , Suriya Gunasekar,322 Chi Jin , Jason Lee , Tengyu Ma , Behnam Neyshabua , and Zhao Song . Theory of323 deep learning . https : //www.cs.princeton.edu/courses/archive/fall19/cos597B/324 lecnotes/bookdraft.pdf/.325 [ 3 ] Sanjeev Arora , Nadav Cohen , Wei Hu , and Yuping Luo . Implicit regularization in deep matrix326 factorization . arXiv preprint arXiv:1905.13655 , 2019.327 [ 4 ] Peter Bartlett , Dylan J Foster , and Matus Telgarsky . Spectrally-normalized margin bounds for328 neural networks . arXiv preprint arXiv:1706.08498 , 2017.329 [ 5 ] Peter Bartlett , Yoav Freund , Wee Sun Lee , and Robert E Schapire . Boosting the margin : A new330 explanation for the effectiveness of voting methods . The annals of statistics , 26 ( 5 ) :1651–1686,331 1998.332 [ 6 ] Tom B Brown , Benjamin Mann , Nick Ryder , Melanie Subbiah , Jared Kaplan , Prafulla Dhariwal,333 Arvind Neelakantan , Pranav Shyam , Girish Sastry , Amanda Askell , et al . Language models are334 few-shot learners . arXiv preprint arXiv:2005.14165 , 2020.335 [ 7 ] Lenaic Chizat and Francis Bach . Implicit bias of gradient descent for wide two-layer neural336 networks trained with the logistic loss . In Conference on Learning Theory , pages 1305–1338.337 PMLR , 2020.338 [ 8 ] J. Dutta , K. Deb , Rupesh Tulshyan , and Ramnik Arora . Approximate kkt points and a proximity339 measure for termination . Journal of Global Optimization , 56:1463–1499 , 2013.340 [ 9 ] Cong Fang , Hangfeng He , Qi Long , and Weijie J Su . Layer-peeled model : Toward understanding341 well-trained deep neural networks . arXiv preprint arXiv:2101.12699 , 2021.342 [ 10 ] Rong Ge , Furong Huang , Chi Jin , and Yang Yuan . Escaping from saddle points—online343 stochastic gradient for tensor decomposition . In Conference on learning theory , pages 797–842.344 PMLR , 2015.345 [ 11 ] Rong Ge , Jason D Lee , and Tengyu Ma . Matrix completion has no spurious local minimum.346 arXiv preprint arXiv:1605.07272 , 2016.347 [ 12 ] Gauthier Gidel , Francis Bach , and Simon Lacoste-Julien . Implicit regularization of discrete348 gradient dynamics in linear neural networks . arXiv preprint arXiv:1904.13262 , 2019.349 [ 13 ] Micah Goldblum , Jonas Geiping , Avi Schwarzschild , Michael Moeller , and Tom Goldstein.350 Truth or backpropaganda ? an empirical investigation of deep learning theory . arXiv preprint351 arXiv:1910.00359 , 2019.352 [ 14 ] Ziwei Ji , Miroslav Dudík , Robert E Schapire , and Matus Telgarsky . Gradient descent follows353 the regularization path for general losses . In Conference on Learning Theory , pages 2109–2136.354 PMLR , 2020.355 [ 15 ] Ziwei Ji and Matus Telgarsky . Characterizing the implicit bias via a primal-dual analysis . In356 Algorithmic Learning Theory , pages 772–804 . PMLR , 2021.357 [ 16 ] Alex Krizhevsky , Ilya Sutskever , and Geoffrey E Hinton . Imagenet classification with deep358 convolutional neural networks . Advances in neural information processing systems , 25:1097–359 1105 , 2012.360 [ 17 ] Anders Krogh and John A Hertz . A simple weight decay can improve generalization . In361 Advances in neural information processing systems , pages 950–957 , 1992.362 [ 18 ] Yenson Lau , Qing Qu , Han-Wen Kuo , Pengcheng Zhou , Yuqian Zhang , and John Wright.363 Short-and-sparse deconvolution–a geometric approach . arXiv preprint arXiv:1908.10959 , 2019.364 [ 19 ] Thomas Laurent and James Brecht . Deep linear networks with arbitrary loss : All local minima365 are global . In International conference on machine learning , pages 2902–2907 . PMLR , 2018.366 [ 20 ] Yann LeCun , Yoshua Bengio , and Geoffrey Hinton . Deep learning . nature , 521 ( 7553 ) :436–444,367 2015.368 [ 21 ] Yuanzhi Li , Tengyu Ma , and Hongyang Zhang . Algorithmic regularization in over-parameterized369 matrix sensing and neural networks with quadratic activations . In Conference On Learning370 Theory , pages 2–47 . PMLR , 2018.371 [ 22 ] Zhiyuan Li and Sanjeev Arora . An exponential learning rate schedule for deep learning . arXiv372 preprint arXiv:1910.07454 , 2019.373 [ 23 ] Zichao Long , Yiping Lu , Xianzhong Ma , and Bin Dong . Pde-net : Learning pdes from data . In374 International Conference on Machine Learning , pages 3208–3216 . PMLR , 2018.375 [ 24 ] Jianfeng Lu and Stefan Steinerberger . Neural collapse with cross-entropy loss . arXiv preprint376 arXiv:2012.08465 , 2020.377 [ 25 ] Kaifeng Lyu and Jian Li . Gradient descent maximizes the margin of homogeneous neural378 networks . arXiv preprint arXiv:1906.05890 , 2019.379 [ 26 ] Dustin G Mixon , Hans Parshall , and Jianzong Pi . Neural collapse with unconstrained features.380 arXiv preprint arXiv:2011.11619 , 2020.381 [ 27 ] Mor Shpigel Nacson , Jason Lee , Suriya Gunasekar , Pedro Henrique Pamplona Savarese , Nathan382 Srebro , and Daniel Soudry . Convergence of gradient descent on separable data . In The 22nd383 International Conference on Artificial Intelligence and Statistics , pages 3420–3428 . PMLR,384 2019.385 [ 28 ] Vardan Papyan , XY Han , and David L Donoho . Prevalence of neural collapse during the386 terminal phase of deep learning training . Proceedings of the National Academy of Sciences,387 117 ( 40 ) :24652–24663 , 2020.388 [ 29 ] Tomaso Poggio and Qianli Liao . Explicit regularization and implicit bias in deep network389 classifiers trained with the square loss . arXiv preprint arXiv:2101.00072 , 2020.390 [ 30 ] Qing Qu , Yuexiang Zhai , Xiao Li , Yuqian Zhang , and Zhihui Zhu . Analysis of the optimization391 landscapes for overcomplete representation learning . arXiv preprint arXiv:1912.02427 , 2019.392 [ 31 ] Noam Razin and Nadav Cohen . Implicit regularization in deep learning may not be explainable393 by norms . arXiv preprint arXiv:2005.06398 , 2020.394 [ 32 ] Karen Simonyan and Andrew Zisserman . Very deep convolutional networks for large-scale395 image recognition . arXiv preprint arXiv:1409.1556 , 2014.396 [ 33 ] Daniel Soudry , Elad Hoffer , Mor Shpigel Nacson , Suriya Gunasekar , and Nathan Srebro . The397 implicit bias of gradient descent on separable data . The Journal of Machine Learning Research,398 19 ( 1 ) :2822–2878 , 2018.399 [ 34 ] Ju Sun , Qing Qu , and John Wright . When are nonconvex problems not scary ? arXiv preprint400 arXiv:1510.06096 , 2015.401 [ 35 ] Ju Sun , Qing Qu , and John Wright . Complete dictionary recovery over the sphere i : Overview402 and the geometric picture . IEEE Transactions on Information Theory , 63 ( 2 ) :853–884 , 2016.403 [ 36 ] Ju Sun , Qing Qu , and John Wright . A geometric analysis of phase retrieval . Foundations of404 Computational Mathematics , 18 ( 5 ) :1131–1198 , 2018.405 [ 37 ] Sharan Vaswani , Reza Babanezhad , Jose Gallego , Aaron Mishkin , Simon Lacoste-Julien , and406 Nicolas Le Roux . To each optimizer a norm , to each norm its generalization . arXiv preprint407 arXiv:2006.06821 , 2020.408 [ 38 ] Colin Wei , Jason Lee , Qiang Liu , and Tengyu Ma . On the margin theory of feedforward neural409 networks . 2018.410 [ 39 ] Stephan Wojtowytsch and Weinan E. On the emergence of tetrahedral symmetry in the final411 and penultimate layers of neural network classifiers . arXiv preprint arXiv:2012.05420 , 2020.412 [ 40 ] Han Xiao , Kashif Rasul , and Roland Vollgraf . Fashion-mnist : a novel image dataset for413 benchmarking machine learning algorithms . arXiv preprint arXiv:1708.07747 , 2017.414 [ 41 ] Linfeng Zhang , Jiequn Han , Han Wang , Roberto Car , and Weinan E. Deep potential molecular415 dynamics : a scalable model with the accuracy of quantum mechanics . Physical Review Letters,416 120 ( 14 ) :143001 , 2018.417 [ 42 ] Yuqian Zhang , Qing Qu , and John Wright . From symmetry to geometry : Tractable nonconvex418 problems . arXiv preprint arXiv:2007.06753 , 2020.419 [ 43 ] Zhihui Zhu , Tianyu Ding , Jinxin Zhou , Xiao Li , Chong You , Jeremias Sulam , and Qing420 Qu . A geometric analysis of neural collapse with unconstrained features . arXiv preprint421 arXiv:2105.02375 , 2021.422 Checklist423 1 . For all authors ... 424 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s425 contributions and scope ? [ Yes ] 426 ( b ) Did you describe the limitations of your work ? [ Yes ] 427 ( c ) Did you discuss any potential negative societal impacts of your work ? [ N/A ] 428 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to429 them ? [ Yes ] 430 2 . If you are including theoretical results ... 431 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] 432 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] 433 3 . If you ran experiments ... 434 ( a ) Did you include the code , data , and instructions needed to reproduce the main experi-435 mental results ( either in the supplemental material or as a URL ) ? [ Yes ] 436 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they are437 chosen ) ? [ Yes ] 438 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-439 ments multiple times ) ? [ No ] 440 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type441 of GPUs , internal cluster , or cloud provider ) ? [ No ] 442 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 443 ( a ) If your work uses existing assets , did you cite the creators ? [ N/A ] 444 ( b ) Did you mention the license of the assets ? [ N/A ] 445 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ N/A ] 446 447 ( d ) Did you discuss whether and how consent is obtained from people whose data you ’ re448 using/curating ? [ N/A ] 449 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable450 information or offensive content ? [ N/A ] 451 5 . If you used crowdsourcing or conducted research with human subjects ... 452 ( a ) Did you include the full text of instructions given to participants and screenshots , if453 applicable ? [ N/A ] 454 ( b ) Did you describe any potential participant risks , with links to Institutional Review455 Board ( IRB ) approvals , if applicable ? [ N/A ] 456 ( c ) Did you include the estimated hourly wage paid to participants and the total amount457 spent on participant compensation ? [ N/A ] 458

Summary: This paper studies representations of neural networks trained with cross entropy, when trained beyond zero training error, utilizing the unconstrained layer peeled model (ULPM). ULPM is motivated by the fact that many neural networks are highly overparametrized and assumes that the representations, on which the last linear layer acts, can be optimized freely, i.e.  they are unconstrained. In this model, the authors derive several theoretical results, related to the recently observed neural collapse [28]. In particular, they derive a relation to max-margin analysis, show that neural networks converge to Karush–Kuhn–Tucker (KKT) points of a certain minimum-norm separation problem, and show that the loss landscape is benign, as it has no spurious minima.