A N UNCONSTRAINED LAYER - PEELED PERSPECTIVE ON NEURAL COLLAPSE
Anonymous authors
Paper under double-blind review
A BSTRACT
Neural collapse is a highly symmetric geometry of neural networks that emerges during the terminal phase of training, with profound implications on the generalization performance and robustness of the trained networks.
To understand how the last-layer features and classiﬁers exhibit this recently discovered implicit bias, in this paper we introduce a surrogate model called the unconstrained layer-peeled model (ULPM).
We prove that gradient ﬂow on this model converges to critical points of a minimum-norm separation problem exhibiting neural collapse in its global minimizer.
Moreover, we show that the ULPM with the cross-entropy loss has a benign global landscape for its loss function which allows us to prove that all the critical points are strict saddle points except the global minimizers that exhibit the neural collapse phenomenon.
Empirically, we show that our results also hold during the training of neural networks in real-world tasks when explicit regularization or weight decay is not used
I NTRODUCTION
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 scientiﬁc 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 ﬁnal output function fail to understand the features learned by neural networks.
Recently, in Papyan et al.
, 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 beneﬁt 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 classiﬁcation problem using a neural network with d neurons in the last hidden layer, the LPM takes the following form:
[Equation]
where C 1 , C 2 are positive constants.
Here, W = [ w 1 , w 2 , · · · , w K ] (cid:62) ∈ R K
×
d is the weight of the ﬁnal linear classiﬁer, H =
[ h 1 , h 2 , · · · , h n ] ∈ R d × n is the feature of the last layer and y i 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 ﬂow can ﬁnd 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 ﬁrst-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
and we defer a detailed discussion to Section 5.2.
The reasons we investigate the unregularized objective are summarized as follows:
Feature regularization or constrain is still not equivalent to weight decay used in practice.
However, previous studies have justiﬁed 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).
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.
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.
Reference
Contribution
Feature Norm
Constraint
Feature Norm
Regularization
Loss
Function (Papyan et al., 2020)
Empirical Results (cid:37) (cid:37)
CE Loss (Wojtowytsch & E, 2020) (Lu & Steinerberger, 2020)
(Fang et al., 2021)
Global Optimum (cid:34) (cid:37)
CE Loss (Mixon et al., 2020) (Poggio & Liao, 2020) (Han et al., 2021)
Modiﬁed
Training Dynamics (cid:37) (cid:37) (cid:96) 2 Loss (Zhu et al., 2021) concurrent
Landscape Analysis (cid:37) (cid:34)
CE Loss
This paper
Training Dynamics+
Landscape Analysis (cid:37) (cid:37)
CE Loss
1. C ONTRIBUTIONS
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 inﬁnity, 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 ﬂow 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 ﬂow is sufﬁcient 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. R ELATED W ORKS
Implicit Bias of Gradient Descent: To understand how gradient descent or its variants helps deep learning to ﬁnd 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 (cid:96) 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.
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 ﬂat 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).
P RELIMINARIES AND P ROBLEM S ETUP
In this paper, || · || F denotes the Frobenius norm, (cid:107) · (cid:107) 2 denotes the matrix spectral norm, (cid:107) · (cid:107) ∗ denotes the nuclear norm, (cid:107) · (cid:107) denotes the vector l 2 norm and tr ( · ) is the trace of matrices.
We use [ K ] := { 1 , 2 , · · · , K } to denote the set of indices up to K .
2. P RELIMINARIES
We consider a balanced dataset with K classes (cid:83)
Kk =1 { x k,i } ni
=1 .
A standard fully connected neural network can be represented as:
[Equation]
Here W full =
( W 1 , W 2 , · · · , W L ) denote the weight matrices in each layer, ( b 1 , b 2 , · · · , b L ) denote the bias terms, and σ ( · ) denotes the nonlinear activation function, for example, ReLU or sigmoid.
Let h k,i =
σ ( b L − 1 + W L − 1 σ ( · · · σ ( b 1 + W 1 x k,i )))
∈
R d denote the last layer feature for data x k,i and ¯ h k = 1 n (cid:80) ni
=1 h k,i denotes the feature mean within the k-th class.
To provide a formal deﬁnition of neural collapse, we ﬁrst introduce the concept of a simplex equiangular tight frame (ETF):
Deﬁnition 2.1 (Simplex ETF) .
A symmetric matrix M ∈ R K × K is said to be a simplex equiangular tight frame (ETF) if
[Equation]
Where Q ∈ R d × K is a matrix with orthogonal columns.
Let W ∈ R K × d = W L
=
[ w 1 , w 2 , · · · , w K ] (cid:62) be the weight of the ﬁnal layer classiﬁer, 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 = 1 n (cid:80) ni
=1 h k,i .
[Equation]
k [Equation]
• (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 conﬁguration constrained to the previous two properties
[Equation]
• (NC3) Convergence to self-duality:
The linear classiﬁers and class-means will converge to align with each other, up to appropriate rescaling, that is, there exist a universal constant [Equation] > 0 such that
[Equation]
• (NC4) Simpliﬁcation to Nearest Class-Center For a given deepnet activation:
[Equation]
the network classiﬁer converges to choose whichever class has the nearest train class-mean
[Equation]
In this paper
[Equation]
we say that a point W ∈ R K × d , H ∈ R d × nK satisﬁes the neural collapse conditions or is a neural collapse solution if the above four criteria are all satisﬁed for ( W , H ) .
2. P ROBLEM S ETUP
We mainly focus on the neural collapse phenomenon, which is only related to the classiﬁers and features in the last layer.
Since the general analysis of the highly non-smooth and non-convex neural network is difﬁcult, we peel down the last layer of the neural network and propose the following unconstrained layer-peeled model (ULPM) as a simpliﬁcation to capture the main characters related to neural collapse during training dynamics.
A similar simpliﬁcation is commonly used in previous theoretical works (Lu & Steinerberger, 2020; Fang et al., 2021; Wojtowytsch & E, 2020;
Zhu et al., 2021), but ours does not have any constraint or regularization on features and is closer to realistic neural network models.
It should be mentioned that although Mixon et al.
; Han et al.
(2021) also studied the unconstrained model, their analysis adopted an approximation to real dynamics and was highly dependent on the (cid:96) 2 loss function, which is rarely used in classiﬁcation tasks.
Compared with their works, we directly deal with the real training dynamics and cover the most popular cross-entropy loss in classiﬁcation tasks.
Let W =
[ w 1 , w 2 , · · · , w K ] (cid:62) ∈ R K × d and H =
[ h 1 , 1 , · · · , h 1 ,n , h 2 , 1 , · · · , h K,n ] ∈ R d × Kn be the matrices of classiﬁers and features in the last layer, where K is the number of classes and n is the number of data points in each class.
The ULPM is deﬁned as follows:
[Equation]
Here, we do not have any constraint or regularization on features, which corresponds to the absence of weight decay in deep learning training.
The objective function (2.3) is generally non-convex on ( W , H ) and we aim to study the landscape of the objective function (2.3).
Furthermore, we consider the gradient ﬂow of the objective function:
[Equation]
M AIN R ESULTS
In this section, we present our main results regarding the training dynamics and landscape analysis of (2.3).
This section is organized as follows: First, in Section 3.1, we show the relationship between the margin and neural collapse in our surrogate model.
Inspired by this relationship, we propose a minimum-norm separation problem (3.1) and show the connection between the convergence direction of the gradient ﬂow and the KKT point of (3.1).
In addition, we explicitly solve the global optimum of (3.1) and show that it must satisfy the neural collapse conditions.
However, owing to the nonconvexity, we ﬁnd an Example 3.1 in Section 3.2, which shows that there exist some bad KKT points such that a simple gradient ﬂow will get stuck in them and does not converge to the neural collapse solution, which is proved to be optimal in Theorem 3.3.
Then, we present our second–order analysis result in Theorem 3.4 to show that those bad points will exhibit decreasing directions in the tangent space thus gradient descent and its variants can avoid those bad directions and only converge to the neural collapse solutions (Lee et al., 2016;
Ge et al., 2015).
3. C ONVERGENCE T O T HE F IRST –O RDER S TATIONARY P OINT
Heuristically speaking, the simplex ETF (Deﬁnition 2.1) gives a set of vectors with the maximum average angle between them.
As a result, neural collapse implies that the neural networks tend to maximize the angles between each class and the corresponding classiﬁers.
At a high level, such behavior is quite similar to margin maximization which is known to be an implicit regularization effect of gradient descent and has been extensively studied in Soudry et al.
; Nacson et al.
(2019b)
; Lyu & Li (2019); Ji et al.
First, we illustrate the connection between the margin and neural collapse.
Recall that the margin of a single data point x k,i and the associated feature h k,i is q k,i ( W , H ) :
= w (cid:62) k h k,i
[Equation]
i .
Then, the margin of the entire dataset can be deﬁned as
[Equation]
The following theorem demonstrates that neural collapse yields the maximum margin solution of our
ULPM model:
Theorem 3.1 (Neural collapse as max-margin solution) .
For the ULPM model (2.3), the margin of the entire dataset always satisﬁes
[Equation]
and the equality holds if and only if ( W , H ) satisﬁes the neural collapse conditions with (cid:107) W (cid:107) F = (cid:107) H (cid:107) F .
Based on this ﬁnding, we present an analysis of the convergence of the gradient ﬂow on the ULPM (2.3).
Following Lyu & Li (2019), we link the gradient ﬂow on cross-entropy loss to a minimum-norm separation problem.
Theorem 3.2.
For problem (2.3), let ( W ( t ) , H ( t )) be the path of gradient ﬂow at time t , if there exist a time t
such that L ( W ( t 0 ) , H ( t 0 )) < log 2 , then any limit point of { ( ˆ H ( t ) , ˆ
[Equation]
[Equation] [Equation] is along the direction (i.e., a constant multiple) of a Karush-Kuhn-Tucker (KKT) point of the following minimum-norm separation problem:
[Equation]
Remark 3.1 .
Here we write (3.1) as a constraint problem, but the constraint is introduced by the implicit regularization effect of gradient ﬂow on our ULPM objective (2.3).
Since the training dynamics will diverge to inﬁnity, we hope to justify that the diverge direction is highly related to neural collapse and an appropriate normalization is needed, which is why it appears to be a constraint optimization form.
Our goal is to justify that the neural collapse phenomenon is caused by the properties of the loss function and training dynamics rather than an explicit regularization or constraint, which seems to be necessary for previous studies (Fang et al., 2021; Lu & Steinerberger, 2020; Wojtowytsch & E, 2020).
Remark 3.2 .
In Theorem 3.2, we assume that there exists a time t 0 such that L ( W ( t 0 ) , H ( t 0 )) < log 2 .
Note that the loss function can be rewritten as:
[Equation]
the requirement L ( W , H )
< log 2 implies that w k h k,i − w j h k,i ≥ 0 for any k, j ∈
[ K ] , i ∈ [ n ] , which is equivalent to q min ( W , H ) > 0 ; that is, every feature is separated perfectly by the classiﬁer.
This assumption is common in the study of implicit bias in the nonlinear setting (Nacson et al., 2019a;
Lyu & Li, 2019; Ji et al., 2020) and its validity can be justiﬁed by the fact that neural collapse is found only in the terminal phase of training in the deep neural network, where the training accuracy has achieved 100%.
It is also an interesting direction to remove this assumption and study the early-stage dynamics of training, which is beyond the scope of this study and we leave it to future exploration.
Theorem 3.2 indicates that the convergent direction of the gradient ﬂow is restricted to the max-margin directions, which usually have good robustness and generalization performance.
In general, the
KKT conditions are not sufﬁcient to obtain global optimality because the minimum-norm separation problem (3.1) is non-convex.
On the other hand, we can precisely characterize its global optimum in the ULPM case based on Theorem 3.1:
Corollary 3.1.
Every global optimum of the minimum-norm separation problem (3.1) is also a KKT point that satisﬁes the neural collapse conditions.
With Theorem 3.2 bridging dynamics with KKT points of (3.1) and Corollary 3.1 associating global optimum of (3.1) with neural collapse, the remaining work is to close the gap between the KKT point and the global optimum.
3. S ECOND –O RDER L ANDSCAPE A NALYSIS
In the convex optimization problem, the KKT conditions are usually equivalent to global optimality.
Unfortunately, owing to the non-convex nature of the objective (2.3), the KKT points can also be saddle points or local optimum other than the global optimum.
In this section, we aim to show that this non-convex optimization problem is actually not scary via landscape analysis.
To be more speciﬁc, we prove that except for the global optimum given by neural collapse, all the other KKT points are actually saddle points that can be avoided by gradient ﬂow.
In contrast to previous landscape analysis of non-convex problems, where people aim to show that the objective has a negative directional curvature around any stationary point (Sun et al., 2015; Zhang et al., 2020), our analysis is slightly different.
Note that since the model is unconstrained, once features can be perfectly separated, the ULPM objective (2.3) will always decrease along the direction of the current point and the optimum is attained only in inﬁnity.
Although growing along all of those perfectly separating directions can let the loss function decrease to 0, the speed of decrease is quite different and there exists an optimal direction with the fastest decreasing speed.
As shown in Section 3.1, ﬁrst-order analysis of training dynamics fails to distinguish such an optimal direction from KKT points, and we need second-order analysis to help us fully characterize the realistic training dynamics.
First, we provide an example to illustrate the motivation and necessity of second-order landscape analysis.
Example 3.1 (A Motivating Example) .
Consider the case where K = 4 , n = 1 , let ( W , H ) be the following point:
[Equation]
−      −  − 
 .
One can easily verify that ( W , H ) enables our model to classify all features perfectly.
Furthermore, we can show that it is along the direction of a KKT point of the minimum-norm separation problem
(3.1) by constructing the Lagrangian multiplier Λ = ( λ ij )
Ki,j =1 as follows:
[Equation]
                  .
To see this, simply write down the corresponding Lagrangian (note that we aim to justify ( W , H ) is along the direction of a KKT point of (3.1), to make it a true KKT point , one needs to multiple 1 / √ 2 on W , H ):
[Equation]
Simply take derivatives for W , H and Λ we ﬁnd that it satisﬁes the KKT conditions.
However, the gradient of ( W , H ) is:
[Equation]
[Equation] [Equation]
−
[Equation]
We can see that the directions of the gradient and the parameter align with each other (i.e., W is parallel to ∇ W L ( W , H ) , and H is parallel to ∇ H L ( W , H ) ), which implies that simple gradient descent may get stuck in this direction and only grow the parameter norm.
However, if we construct:
[Equation]
− α 1 − α − 1 − α − α − α − 1 − α 1 − α 
 ,
By simple calculation, we ﬁnd that f ( α ) :
= L ( W (cid:48) , H (cid:48) )
satisﬁes f (cid:48)
( α ) = 0 , f (cid:48)(cid:48)
( α ) < 0 .
Since (cid:107) W (cid:48) (cid:107) F = (cid:107) W (cid:107) F , (cid:107) H (cid:48) (cid:107) F = (cid:107) H (cid:107) F and || W (cid:48) − W || 2 F + || H (cid:48)
− H ||
2 F → 0 as α → 0 , this result implies that for any (cid:15) > 0 , we can choose appropriate α such that:
[Equation]
In Example 3.1, it is shown that there are some KKT points of the minimum-norm separation problem (3.1) that are not globally optimal, but there exist some better points close to it; thus, the gradientbased method can easily avoid them (see Lee et al.
Ge et al.
for a detailed discussion).
In the following theorem, we will show that the best directions are neural collapse solutions in the sense that the loss function is the lowest among all the growing directions.
Theorem 3.3.
The optimal value of the loss function (2.3) on a sphere is obtained (i.e., L ( W , H ) ≤
L ( W (cid:48) , H (cid:48) ) for any || W (cid:48)
||
2 F + || H (cid:48)
||
2 F = ||
W || 2 F + ||
H || 2 F )
if only if ( W , H ) satisﬁes neural collapse conditions and || W || F = ||
H || F .
Remark 3.3 .
Note that the second condition is necessary because neural collapse conditions do not specify the norm ratio of W and H .
That is, if ( W , H ) satisﬁes the neural collapse conditions, then for any α, β ∈ R , ( α W , β H ) will also satisfy them, but only some certain α, β are optimal.
Now, we turn to those points that are not globally optimal.
To formalize our discussion in the motivating example 3.1, we ﬁrst introduce the tangent space:
Deﬁnition 3.1 (tangent space) .
The tangent space of ( W , H ) is deﬁned as a set of directions that are orthogonal to ( W , H ) :
[Equation]
Our next result justiﬁes our observation in Example 3.1 that for those non-optimal points, there exists a direction in the tangent space such that moving along this direction will lead to a lower objective value.
Theorem 3.4.
If ( W , H ) is not the optimal solution in Theorem 3.3 (i.e., ( W , H ) is not a neural collapse solution or it is a neural collapse solution but (cid:107) W (cid:107) F (cid:54)
= (cid:107) H (cid:107) F ), then there exists a direction (∆ W ,
∆ H ) ∈ T ( W , H ) and constant M > 0 such that for any 0 < δ < M ,
[Equation]
Further more, it implies that for any (cid:15) > 0 there exists a point ( W (cid:48) , H (cid:48) )
such that:
[Equation]
Remark 3.4 .
The result in (3.3) gives us a decreasing direction orthogonal to the direction of ( W , H ) .
As shown in Example 3.1, the gradient on these non-optimal points might be parallel to ( W , H ) ; thus, the ﬁrst-order analysis fails to explain the prevalence of neural collapse.
Here the decreasing direction is obtained by analyzing the Riemannian Hessian matrix and ﬁnding its eigenvector corresponding to a negative eigenvalue, which further indicates that these points are ﬁrst-order saddle points in the tangent space. That is why we name it second-order landscape analysis.
A formal statement and proof are presented in the appendix C.
Previous works have shown that for a large family of gradient-based methods, they can avoid saddle points and only converge to minimizers (Lee et al., 2016;
Ge et al., 2015; Panageas et al., 2019), thus our landscape analysis indicates that the gradient ﬂow dynamics only ﬁnd neural collapse directions [Figure]
E MPIRICAL R ESULTS
[Figure]
To evaluate our theory
[Figure]
we trained the [-VGG-13 (Simonyan & Zisserman, 2014) and ResNet
-]ResNet18 (He et al., 2016) on both MNIST (LeCun et al., 1998) and CIFAR-10 (Krizhevsky et al., 2009) datasets without weight decay [Figure] and tracked how the last layer features and classiﬁers converge to neural collapse solutions.
The results are plotted in Figure 1.
Here we reported the following metrics to measure the level of neural collapse: The variation of both feature norm (i.e., Std ( (cid:107) ¯ h k − ¯ h (cid:107) ) / Avg ( (cid:107) ¯ h k − ¯ h (cid:107) ) ) and classiﬁer norm in the last layer ( i
e., Std ( (cid:107) ¯ w k (cid:107) ) /
Avg ( (cid:107) ¯ w k (cid:107) ) ).
Within class variation for last layer features (i.e., Avg ( (cid:107) h k,i − h k (cid:107) ) /
Avg ( (cid:107) h k
i − ¯ h (cid:107) ) ).
Average cosine between last layer features (i
e., Avg ( | cos(¯ h k , ¯ h k (cid:48) )
+ 1 / ( K − 1) | ) ) and that of last layer classiﬁers (i.e., Avg ( | cos( ¯ w k , ¯ w k (cid:48) ) + 1 / ( K − 1) | ) ). Self-duality distance between features and classiﬁers corresponding to the same class in the last layer. (i.e., Avg ( | (¯ h k − ¯ h ) /
(cid:107) ¯ h k − ¯ h (cid:107) − ¯ w k / (cid:107) ¯ w k (cid:107)| ) )
Here a lower value of each metric represents closer to a neural collapse solution. Speciﬁcally, the within class variation measures (NC1) ,the variation of norms measures and cosine between pairs measure (NC2), the self-duality measures (NC3) and (NC4) has been shown to be a corollary of (NC1)-(NC3) (Papyan et al., 2020). More experiment results under various settings and training details can be found in Appendix D. As we can see, all of these metric decreases along the training epochs. These results reveal strong evidence for the emergence of neural collapse in the unconstrained setting and provide sound support for our theory.
C ONCLUSION AND D ISCUSSION 5. C ONCLUSION
To understand the implicit bias of neural features from gradient descent training, we built a connection between max-margin implicit bias and the neural collapse phenomenon and studied the ULPM in this study.
We proved that the gradient ﬂow of the ULPM converges to the KKT point of a minimum-norm separation problem, where the global optimum satisﬁes the neural collapse conditions.
Although the
ULPM is non-convex, we show that ULPM has a benign landscape where all the stationary points are strict saddle points in the tangent space, except for the global neural collapse solution.
Our study helps to demystify the neural collapse phenomenon, which sheds light on the generalization and robustness during the terminal phase of training deep networks in classiﬁcation problems.
5. R ELATIONSHIP WITH O THER R ESULTS ON N EURAL C OLLAPSE
Theoretical analysis of neural collapse was ﬁrst provided by Lu & Steinerberger (2020)
; Wojtowytsch & E (2020); Fang et al.
(2021), who showed that the neural collapse solution is the only global minimum of the simpliﬁed non-convex objective function.
In particular, Wojtowytsch & E (2020); Lu & Steinerberger (2020) studied a continuous integral form of the loss function and showed that the features learned should have a uniform distribution on the sphere.
A more realistic discrete setting was studied in Fang et al.
(2021), where the constraint is on the entire feature matrix rather than individual features.
Our result utilizes the implicit bias of the cross-entropy loss function to remove the feature norm constraint, which is not practical in real applications.
Although the global optimum can be fully characterized by neural collapse conditions, the ULPM objective is still highly non-convex.
Regarding optimization, Mixon et al.
(2020); Poggio & Liao (2020); Han et al.
(2021) analyzed the unconstrained feature model with (cid:96) 2 loss and established convergence results for collapsed features for gradient descent.
However, they fail to generalize to the more practical cross-entropy loss functions used in classiﬁcation tasks.
The analysis relies highly on the (cid:96) 2 loss to obtain a closed-form gradient ﬂow and still requires some additional approximation to guarantee global convergence.
The most relevant study is a concurrent work (Zhu et al., 2021), which provides a landscape analysis of the regularized unconstrained feature model.
Zhu et al.
(2021) turns the feature norm constraint in Fang et al.
(2021) into feature norm regularization and still preserves the neural collapse global optimum.
At the same time, it shows that the modiﬁed regularized objective shares a benign landscape, where all the critical points are strict saddles except for the global one.
Although our study and Zhu et al.
(2021) discover similar landscape results, we believe our characterization remains closer to the real algorithms since we do not introduce any constraints or regularization on the feature norm following the conventional setting in realist training.
The regularization of features introduced in Zhu et al.
(2021) is still different from weight decay regularization (Krogh & Hertz, 1992).
However, weight decay on homogeneous neural networks is equivalent to gradient descent with scaling step size on an unregularized objective (Li & Arora, 2019; Zhang et al., 2018).
Moreover, neural networks are found to perform well in the absence of weight decay, which highlights the importance of implicit regularization (Zhang et al., 2021).
As a result, instead of explicit feature norm constraint/regularization, in this study, we consider implicit regularization brought by the gradient ﬂow on cross-entropy.
We show that implicit regularization is sufﬁcient to lead the dynamics to converge to the neural collapse solution without the explicit regularization.
A E LEMENTS OF OPTIMIZATION
In this section, we introduce some basic deﬁnitions and theory about optimization.
In the following discussion, we consider a standard form inequality constrained optimization problem:
[Equation]
In addition, we assume all of those functions f
and g i are twice differentiable.
A point x ∈ R d is said to be feasible if and only if it satisﬁes all of the constraints in (A.1), i.e., g i ( x ) ≤ 0 , i ∈ [ n ] .
And the Lagrangian of problem (A.1) is deﬁned as following:
[Equation]
A. K ARUSH K UHN T UCKER C ONDITIONS
Now let’s ﬁrst introduce the deﬁnition of Karush Kuhn Tucker (KKT) point and approximate KKT point.
Here we follows the deﬁnition of ( (cid:15), δ ) -KKT point as in Lyu & Li (2019).
Deﬁnition A.1 (Deﬁnition of KKT point) .
A feasible point is said to be KKT point of problem (A.1) if there exist λ 1 , λ 2 , · · · , λ n ≥ 0 such that the following Karush Kuhn Tucker (KKT) conditions hold:
∇
[Equation]
( x )
[Equation]
λ
i g i ( x ) = 0 , i ∈
[ n ]
Deﬁnition A.2 (Deﬁnition of ( (cid:15), δ ) -KKT point) .
For any (cid:15), δ > 0 , a feasible point of (A.1) is said to be ( (cid:15), δ ) -KKT point of problem (A.1) if there exist λ 1 , λ 2 , · · · , λ n ≥ 0 such that:
||∇ f ( x ) + (cid:80) ni
=
λ i ∇ g
i ( x ) || ≤ (cid:15) 2.
λ
i g i ( x ) ≥ − δ, i ∈ [ n ]
Generally speaking, KKT conditions might not be necessary for global optimality.
We need some additional regular conditions to make it necessary.
For example, as shown in Dutta et al.
(2013) we can require the problem to satisfy the following Mangasarian-Fromovitz constraint qualiﬁcation (MFCQ):
Deﬁnition A.3 (Mangasarian-Fromovitz constraint qualiﬁcation (MFCQ) ) .
For a feasible point x of (A.1), problem (A.1) is said to satisfy (MFCQ) at x if there exist a vector v ∈ R d such that:
[Equation]
Moreover, when MFCQ holds we can build a connection between approximate KKT point and KKT point, see detailed proof in Dutta et al.
Theorem A.1 (Relationship between Approximate KKT point and KKT point) .
Let (cid:8) x k ∈ R
d : k ∈ N (cid:9) be a sequence of feasible points of (A.1) , { (cid:15) k > 0 : k ∈ N } and { δ k > 0 : k ∈ N } be two sequences of real numbers such that x k is an ( (cid:15) k , δ k )
-KKT point for every k , and (cid:15) k
→ 0 , δ k → 0 .
If x k → x as k → + ∞ .
If MFCQ (A.3) holds at x , then x is a KKT point of ( P ) .
B O MITTED PROOFS FROM S ECTION 3.
Recall our ULPM problem:
[Equation] [Equation]
Let’s review some basic notations deﬁned in the main body.
Let s k,i,j = w (cid:62) k h k,i − w (cid:62) j h k,i ,
∀ k ∈
[ K ] , i ∈ [ n ] , j ∈ [ K ] , the margin of a single feature h k,i is deﬁned to be q k,i ( W , H ) :
= min j (cid:54) =
k s k,i,j = w (cid:62) k h k,i − max j (cid:54) =
k w (cid:62) j h k,i .
We deﬁne the margin of entire dataset as q min =
q min ( W , H ) =
min k ∈ [1 ,K ] ,i ∈ [1 ,n ] q k,i ( W , H ) .
We ﬁrst prove Theorem 3.1 in the mainbody.
Proof of Theorem 3.1.
First we can ﬁnd that the margin will not change if we minus a vector a for all w j , so if we denote the mean of classiﬁer ˜ w i =
w i − 1 K (cid:80) Ki
=1 w i
and then we have w [Equation] k
[Equation]
k,i [Equation] [Equation] min
[Equation]
W , H ) , that is:
[Equation]
Note that (cid:80) Kj =1 ˜ w (cid:62) j h k,i = 0 then sum this inequality over j we have:
[Equation]
By Cauchy inequality, we have:
[Equation]
Sum (B.2) over k and i we have:
[Equation]
On the other hand, we know that:
[Equation]
Then we can conclude that:
[Equation]
as desired.
When the equality holds, ﬁrst we have || ˜ W
|| 2 F =
|| W || 2 F which is equivalent to 1 K (cid:80)
Ki
=
w i = 0 , ˜ w
i = w i .
Take it back into (B.3), then we must have all of the equality holds in (B.2) and (B.1), which give us:
[Equation]
Take this into (B.1) we have:
[Equation]
which implies neural collapse conditions.
Now let’s turn to the training dynamics and prove Theorem 3.2.
Before starting the proof, we need to introduce some additional notations.
Since the W and H are all optimization variables here, we denote θ = vec ( W , H ) as the whole parameter for simplicity and all of previous function can be deﬁned on θ by matching the corresponding parameters.
Denote ρ = (cid:107) θ (cid:107) as the norm of θ and ˜ γ = − log( e L ( θ ) − 1) ρ.
Now we can state our ﬁrst lemma to show how training dynamics of gradient ﬂow on ULPM objective (B) is related to a KKT point of (3.1).
Lemma B.1.
If there exist a time t 0 such that L ( θ ( t 0 )) < log 2 , then for any t > t 0 ˜ θ := θ/q min ( θ ) 1 / 2 is an ( (cid:15), δ ) - approximate KKT point of the following minimum-norm separation problem.
More precisely, we have
[Equation]
where:
[Equation]
is the angle between θ and its corresponding gradient.
Proof.
The training dynamics is given by gradient ﬂow:
[Equation]
Then by the chain rule we have:
[Equation]
It indicates that the loss function L is monotonically decreasing.
If L ( θ ( t 0 )) < log 2 , we have
L ( θ ( t )) < log 2 , ∀ t > t 0 .
On the other hand, note that
[Equation]
which gives us q min ( t ) > 0 , ∀ t > t 0 .
Let g = dθdt , note that we can rewrite the ULPM objective function (B) as L ( θ )
= (cid:80)
Kk =1 (cid:80) ni
=1 log(1+ (cid:80)
j
[Equation]
= k e − s k,i,j ) .
By the chain rule and the gradient ﬂow equation we have
[Equation]
where g k,i,j is the gradient of s k,i,j , i.e. g k,i,j = ∇ θ s k,i,j ( θ ) .
Now let ˜ g k,i,j = g k,i,j /q 1 / 2 min = ∇ θ s k,i,j (˜ θ ) and construct λ k,i,j =
[Equation]
, we only need to show:
[Equation]
To prove (B.6), we only need to compute (Recall that ˜ θ
= θ/q min ( θ ) 1 / 2 ):
|
[Equation]
Note that:
[Equation]
Then we have the following inequality:
[Equation]
Take this back into (B.8) we have (B.6) as desired.
To prove (B.7), ﬁrst by our construction:
[Equation]
B.10)
Note that
||
g || 2 ≥ (cid:68) g, θ ||
θ
|| 2 (cid:69) =
1 ρ (cid:104) g, θ (cid:105) and (cid:104)
g k,i,j , θ (cid:105) =
2 s k,i,j since s k,i,j = w (cid:62) k h k,i − w (cid:62) j h k,i , we have:
[Equation]
since s k
[Equation]
i
[Equation]
j ≥ q min
[Equation]
and e − s k
[Equation]
i
[Equation]
≤ 1)
[Equation]
Take this inequality back into the (B.10) we have:
[Equation]
n (cid:88) i = (cid:88)
j (cid:54) =
k e q min − s k,i,j
+ (cid:80) l (cid:54) =
k e − s k,i,l ( s k,i,j − q min ) ≤ Kρ2 q 2 min
K (cid:88) k
= n (cid:88) i = (cid:88)
j (cid:54) =
k e q min − s k,i,j ( s k,i,j − q min ) ≤
K
[Equation]
K (cid:88) k
= n (cid:88) i = (cid:88)
j (cid:54) =
k e q min − s k,i,j ( s k,i,j
[Equation]
q min )
[Equation]
Where the last inequality is obtained from the fact xe − x ≤ 1 e , ∀ x > 0 , which can be proved by some elementary calculus.
Based on Lemma B.1, we have shown that the ( W , H ) will be a ( (cid:15), δ ) -KKT point, if we can show ( (cid:15), δ ) converge to zero, then by Theorem A.1 we know the limit point will be along the direction of a KKT point.
Ignoring the constant term, we only have to show how ˜ γ ( t ) , β ( t ) and q min ( t ) evolve along time.
Now we provide the following lemmas to illustrate their dynamics.
The ﬁrst lemma aims at proving that the norms of parameter ρ ( t ) and ˜ γ ( t ) are monotonically increasing.
Lemma B.2.
If there exist t 0 such that L ( θ ( t 0 )) < log 2 , then for any t > t 0 we have:
[Equation]
B
[Equation] [Equation]
Proof.
We can disentangle the whole training dynamics into the following two parts:
• the radial part: v :=
ˆ θ ˆ θ (cid:62) dθdt , • the tangent part: u
= ( I − ˆ θ ˆ θ (cid:62) ) dθdt .
First analyze the radial part, by the chain rule: || v || 2 = | ˆ θ (cid:62) dθdt | = | 1 ρ (cid:10) θ, dθdt (cid:11) | = | 1 ρ 12 dρ 2 dt | .
For dρ 2 dt , we have the following equation:
[Equation]
where the last equality holds by equation (B.11).Then when t > t 0 , we have shown that q min ( t ) ≥ 0 , combine this with the fact s k,i,j ≥ q min we obtain the ﬁrst inequality in (B.12)
[Equation]
Next, we aim to prove the monotonicity of ˜ γ ( t ) , compute the derivative of ˜ γ ( t ) we have:
[Equation]
Recall that we have:
[Equation]
≥
[Equation]
Then second last line is because the deﬁnition of the loss function L ( θ ) = (cid:80)
Kk =1 (cid:80) ni
=1 log(1 + (cid:80) j (cid:54) =
k e − s k,i,j ) ≥ log(1 + (cid:80) j (cid:54) =
k e − s k,i,j ) and the monotonicity of e x − 1 xe x (in fact, ddx e x − 1 xe x = e − x ( x − e x +1)
x≤ 0 , ∀ x > 0 ).
As a result, we notice that
[Equation]
At the same time, we notice that (cid:107) v (cid:107) 2 = 1 ρ 2 (cid:16) 12 dρ 2 dt (cid:17)
2 = 12 dρ 2 dt · ddt log ρ on the one hand, and by the chain rule:
[Equation]
Combine this with the radial term:
[Equation]
Dividing 12 dρ 2 dt on both sides, we have
[Equation] [Equation]
[Equation]
dt (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ − d L ( θ ) dt q min e L ( θ ) 2( e L ( θ ) − 1) .
Now by equation (B.15) and (B.16) we obtain:
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) .
[Equation]
By (B.14) we know the ρ is monotonically increasing, we have ddt log ρ > 0 and then we get the second inequality in (B.12).
Lemma B.2 gives us the monotonicity of ˜ γ , note that since the loss function L ( θ ( t 0 )) < log 2 , we have ˜ γ ( t ) ≥ ˜ γ ( t 0 ) > 0 , then we can treat ˜ γ ( t ) in Lemma B.1 as a positive constant.
The remaining work is to show q min ( t ) grows to inﬁnity and β ( t ) → 1 .
To show q min ( t ) → ∞ , it’s equivalent to show L ( t ) → 0 and we have the following lemma:
Lemma B.3.
If there exist t 0 such that L ( θ ( t 0 )) < log 2 , then L ( θ ( t ))
→ 0 and q min ( θ ( t ))
→ ∞ as t → ∞ , moreover we have the convergence rate L ( θ ( t ))
=
O (1 /t )
Proof.
By (B.4) and (B.13), the evolution of loss function L ( θ ) can be written as:
[Equation]
Combine it with (B.5),(B.14) and (B.17) we have:
[Equation]
which indicates:
[Equation] [Equation]
Since 0 < L ( θ ( t )) <
L ( θ ( t 0 )) < log 2 , note that:
[Equation]
which implies:
[Equation]
On the other hand, we can ﬁnd that:
[Equation]
Combine equation (B.21) and (B.22) together, one can conclude that there exist a constant C >such that for t > t 0 :
[Equation]
It further implies that:
[Equation]
then integral on both side we have:
[Equation]
and
[Equation]
Thus we must have L ( θ ( t ))
→ 0 and combine this with q min ≥ − log( e L ( θ ( t )) − 1) we know q min ( θ ( t ))
→ ∞ as desired.
To bound β ( t ) , we ﬁrst need a useful lemma to bound the changes of the direction of θ .
Lemma B.4.
If there exist t 0 such that L ( θ ( t 0 )) < log 2 , then for any t > t
[Equation]
Proof.
First we know that:
[Equation]
i = (cid:88)
[Equation]
(cid:54) =
k e
[Equation]
(B.23)
Recall our g k,i,j = ∂s k,i,j ∂θ and s k,i,j = w (cid:62) k h k,i − w (cid:62) j h k,i , we can ﬁnd that (cid:107) g k,i,j (cid:107) ≤ 2 ρ .
On the other hand, Combine it with (B.13) and (B.9) we have:
[Equation]
as desired.
Where the second inequality holds by multiple ρ q min on the right hand of (B.23) and the formula of dρ 2 dt in equation (B.13), and the last inequality holds since we have shown that ˜ γ ( t ) is monotonically increasing in (B.2) and ˜ γ ( t 0 ) > 0 since L ( t 0 ) < log 2 .
Let’s turn back to β , though we can not show directly that it increase to one, we can ﬁnd a sequence of time { t m } for each limit point such that β ( t m ) →
Lemma B.5.
If there exist t 0 such that L ( θ ( t 0 )) < log 2 , then for every limit point ¯ θ of { ˆ θ ( t ) : t ≥ 0 } , there exists a sequence of time { t m > 0 : m ∈ N } such that t m → ∞ , ˆ θ ( t m )
→ ¯ θ , and β ( t m ) → 1 .
Proof.
Recall that in equation (B.19) we have shown that:
[Equation]
Since ddt log ρ = 1 ρ dρdt = 12 ρ 2 dρ 2 dt = 1 ρ 2 (cid:10) θ, dθdt (cid:11) and:
[Equation]
Plug them into (B.24) we have:
[Equation]
For any t 2 > t 1 > t 0 , integrate both sides from time t 2 to t 1 we have: log ˜ γ ( t 2 ) − log ˜ γ ( t 1 ) ≥ 2 (cid:82) t 2 t 1 ( β ( t ) − 2 − 1) ddt log ρdt .
By the continuity of β we know there exist a time t ∗ such that:
[Equation]
By (B.9) we know that ˜ γ ≤ q min ρand the right hand is bounded, and the ˜ γ is monotonically increasing, then there exist ˜ γ ∞ such that ˜ γ ( t ) ↑ ˜ γ ∞ .
Now we are ready to construct the sequence of t m , ﬁrst take a sequence of { (cid:15) m > 0 , m ∈ N } such that (cid:15) m → 0 .
We construct t m by induction, suppose we have already ﬁnd t 1 < t 2 < · · · < t m −satisfy our requirement, since ¯ θ is a limit point of { ˆ θ ( t ) : t > 0 } , then we can ﬁnd a time s m such that:
[Equation]
By the monotonicity and continuity of ρ we can ﬁnd a time s (cid:48) m such that log ρ ( s (cid:48) m ) − log ρ
[Equation]
s m ) ≤ (cid:15) m .
Take t 2 = s (cid:48) m , t 1 = s m in (B.25), there exist a time t m such that:
[Equation]
[Equation]
on the other hand, by Lemma B.4 we have:
[Equation]
Note that (cid:10) θ, dθdt (cid:11) > 0 , then by deﬁnition we know β > 0 .
Combine equations (B.26) and (B.27) we have β ( t m ) → 1 and ˆ θ ( t m )
→ ¯ θ as desired.
Now we are ready to prove Theorem 3.2:
Proof of Theorem 3.2.
By
Lemma
B.1, we know that once t > t 0 , ( W ( t ) , H ( t )) /q
min ( W
[Equation]
t ) , H ( t )) is an
[Equation]
-approximate
KKT point.
We have shown that ˜ γ ( t ) >
˜ γ ( t 0 ) > 0 in Lemma B.2, q min → ∞ in Lemma
B.3 and from Lemma B.5 we know for any limit point ( ¯ W , ¯ H ) of { ( ˆ H ( t ) , ˆ W ( t )) :
= (
[Equation]
) } , there exists a sequence of time { t m > 0 : m ∈ N } such that t m → ∞ , β ( t m ) → 1 and ( ˆ H ( t m ) , ˆ W ( t m ))
→ ( ¯ W , ¯ H ) .
Then ( ¯ W , ¯ H ) is along the direction of a limit point of a sequence of ( (cid:15), δ ) -approximate
KKT point with (cid:15), δ → 0 .
On the other hand, we can verify that the problem (3.1) satisﬁes MCFQ (A.3) by simply setting v = θ , then:
[Equation]
Now by Theorem A.1 we know ( ¯ W , ¯ H ) is along the direction of a KKT point of problem (3.1)
Theorem 3.2 characterize the convergent behaviour of gradient ﬂow, under separable conditions the limit point is along the direction of a KKT point of (3.1), next we show that the global minimum of (3.1) must satisfy neural collapse conditions by proving Corollary 3.1:
Proof of Corollary 3.1.
Since we have shown that the problem (3.1) satisfy MCFQ, then the KKT conditions are necessary for global optimality, we only need to show the global optimum satisﬁes neural collapse conditions.
First the constraints in (3.1) can be transformed to be a single constraint by the deﬁnition of margin:
[Equation]
Note that the margin is homogeneous:
[Equation]
Then for any point ( W , H ) satisﬁes q min ( W , H ) > 0 , after an appropriate scaling α , ( α W , α H ) , ∀ α 2 ≥ 1 /q min ( W , H ) is feasible for (3.1).
Take optimum among all scaling factor α we know the minimum norm is attained if and only if α 2 = 1 /q min ( W , H ) .
And the optimum norm is:
[Equation]
Then by Theorem 3.1 we have:
[Equation]
And the global optimum is attained only when ( W , H ) satisﬁes neural collapse conditions
C O MITTED PROOFS FROM S ECTION 3.
To begin with, let’s ﬁnish the computation in the motivating example (Example 3.1)
Proof in Example 3.1.
Consider the case where K = 4 , n = 1 , let ( W , H ) be the following point:
[Equation]
One can easily verify that this ( W , H ) enables our model to classify all of the features perfectly.
Further more, we can show it is along the direction of a KKT point of the minimum-norm separation problem (3.1) by construct the Lagrangian multiplier Λ = ( λ ij )
Ki,j =1 as following: Λ =                    .
To see this, just write down the corresponding Lagrangian (note that to make it to be a true KKT point of (3.1), one needs to multiple 1 / √ 2 on W , H ):
[Equation]
Simply take derivatives for W , H and Λ
we can ﬁnd it satisﬁes KKT conditions.
On the other hand, the gradient of ( W , H ) is
[Equation]
= ∇ H L ( W , H )
= − 2 + 2 e −2 + 2 e − 2 + 2 e  −  −      −  − 
 .
We can ﬁnd that the directions of gradient and the parameter align with each other (i.e., W is parallel to ∇ W L ( W , H ) , and H is parallel to ∇ H L ( W , H ) ), which implies that simple gradient descent may gets stuck in this direction and only grows the parameters’ norm.
However, if we construct:
W (cid:48)
= (cid:114) 1 + 2 α  1 + α − 1 + α α α −
+ α 1 + α α
α − α − α 1 − α − 1 − α − α − α − 1 − α 1 − α 
 ,
H (cid:48) =
(cid:114) 1 + 2 α  1 + α − 1 + α − α − α −
+ α 1 + α −
α − α
α
α 1 − α − 1 − α α α − 1 − α 1 − α 
 .
Note that (cid:107) W (cid:48) (cid:107) F = (cid:107) W (cid:107) F , (cid:107) H (cid:48) (cid:107) F = (cid:107) H (cid:107) F and || W (cid:48) − W || 2 F + || H (cid:48)
− H ||
2 F → 0 as α → 0 .
First we can compute:
[Equation]
−
[Equation]
and:
[Equation]
Our aim is to show that for any (cid:15) > 0 , there exist α such that |
α | < (cid:15) and L ( W (cid:48) , H (cid:48) ) <
L ( W , H ) .
By the formulation of L ( W (cid:48) , H (cid:48) ) , it’s sufﬁcient to show that f ( α ) (cid:44) e 2 2 α 2 − 1 1+2 α 2 + 2 e − 2 2 α 2
+2 α 2 < f (0) .
Then take the derivative of f ( α ) we have:
[Equation]
Now we can ﬁnd
that f (cid:48) (0) =
and f (cid:48)(cid:48) (0) =
16( 1 e 2 − 1) < 0 .
Since the function f ( α ) is continuously twice differentiable, we can conclude that for any (cid:15) > 0 , we can choose appropriate α such that:
[Equation]
Now we prove Theorem 3.3 by some similar strategies as used in proving Theorem 3.1.
Proof of Theorem 3.3.
Again we rewrite the ULPM objective by introducing s k,i,j = w (cid:62) k h k,i − w (cid:62) j h k,i :
L ( W , H ) =
K (cid:88) k
= n (cid:88) i =
[Equation] [Equation]
j (cid:54)
[Equation]
k exp( − s k,i,j )) .
In addition, we can ﬁnd that centralizing w i does not change the value of s k,i,j .
Let ˜ w i =
w i −
K (cid:80) Kk
=1 w k , ∀ i ∈
[ K ] , then s k,i,j = w (cid:62)
k h k,i − w
[Equation]
First by the strict convexity of e x and Jensen Inequality:
[Equation]
Where the last equality is obtained from:
[Equation]
Now again by the strict convexity of log(1 + ( K − 1) exp( − x )) and Jensen inequality, we have:
[Equation]
Where the last inequality holds since (cid:107) W (cid:107)
[Equation]
When all of the above inequality reduce to equality
[Equation]
we must have:
(cid:80)
Ki
=
w i = 0 , ˜ w
i = w i (the last inequality in (C.
w k = √ n h k,i , ∀ i ∈ [ n ] (the third inequality in (C.2)
[Equation]
(cid:107) w k (cid:107) =
(cid:107)
[Equation]
[ K ] , i,
[Equation]
∈ [ n ] (the second inequality in (C.2))
[Equation]
the ﬁrst inequality in
[Equation]
[Equation]
These four conditions are exactly equivalent to neural collapse conditions and (cid:107) W (cid:107) F = (cid:107) H (cid:107) F
The global optimality is not enough to illustrate how does gradient ﬂow converges to neural collapse since there may exist some bad local minimum.
We will provide the following second-order analysis to eliminate spurious local minimum.
First, deﬁne the cross-entropy loss on a matrix Z ∈ R K × nK :
[Equation]
where z k,i,j denote the j -th row and [( k − 1)
K
+ i ] -th column elements of Z .
Then we have
L ( W , H ) =
L ( W H ) .
Now compute the gradient of L ( Z ) to each element:
[Equation]
If u ∈ R K satisﬁes u (cid:62) ∇ L ( Z ) = 0 , denote u p as the maximum element of u , then we have:
[Equation]
Where the last inequality holds if and only if u q = u p , ∀ q ∈
[ K ] .
Which indicates that the rank of ∇ L ( Z ) is K − 1 and u (cid:62) ∇ L ( Z ) = 0 ⇔ u = 1 .
Now we are ready to prove Theorem 3.4 in the main body.
Proof of Theorem 3.4.
First compute the gradient of ULPM objective (B), by the chain rule we have:
[Equation]
If there exist a vector (∆ W , ∆ H ) ∈
T ( W , H ) such that (cid:104)∇ W L ( W , H ) , ∆ W (cid:105) +
[Equation]
H L ( W , H ) , ∆ H
[Equation]
= 0 , moreover we can assume (cid:104)∇
[Equation]
+
[Equation]
H L ( W , H ) , ∆ H (cid:105) < 0 since we can take the negative direction if the formula is greater than zero, then by Taylor expansion:
[Equation]
we know that (∆ W , ∆ H ) satisﬁes our requirement.
Now let’s discuss the case when:
[Equation]
by deﬁnition of T ( W , H ) , it contains all vectors that are orthogonal to ( W , H ) , so ( ∇ W L ( W , H ) , ∇ H L ( W , H )) is parallel to ( W , H ) , that is, there exist λ such that:
[Equation]
If there does not exist (∆ W , ∆ H ) satisfying the requirement, we know that for any feasible curve φ ( t ) =
( W ( t ) , H ( t )) with φ (0) =
( W , H ) on the sphere S
= { ( W (cid:48) , H (cid:48) )
[Equation]
= (cid:107) W (cid:107) 2 F + (cid:107) H (cid:107) 2 F } , t = 0 admits the local minimum of L ( φ ( t )) and thus:
[Equation]
On the other hand, since the curve lies on the sphere S , denote h ( W , H )
= (cid:107) W (cid:107) 2 F + (cid:107) H (cid:107) 2 F , then h ( φ ( t )) must stay as a constant, take twice derivative we have:
[Equation]
Then sum these two conditions together, we have:
[Equation]
By equation (C.4) we know that ∇ ( L + | λ | 2 h )( W , H ) =
Note that φ (cid:48) (0) ∈ T ( W , H ) since the curve lies on S and for any (∆ W , ∆ H ) ∈ T ( W , H ) we can construct a curve φ ( t ) such that φ (cid:48) (0) =
(∆ W , ∆ H ) .
Then (C.7) indicates that ∀ (∆ W , ∆ H ) ∈ T ( W , H ) we have:
[Equation]
When λ = 0 , by equation (C.3) we know that (cid:107)∇ L ( WH ) (cid:107) 2 > 0 , which gives us | λ | ≤ (cid:107)∇ L ( W H ) (cid:107)
When λ (cid:54) = 0 , combine the two equations in (C.4)
we know:
[Equation]
which further implies:
[Equation]
Thus we also have (Note that when W = H =
we can take λ to be zero):
[Equation]
Now when | λ | < ||∇ L ( W H ) || 2 , we can show that it will contradict with (C.8): We have shown that the rank of ∇ L ( Z ) is K − 1 , so by (C.4) and (C.9) there exist a vector a such that W a = H (cid:62) a = 0 , let u and v are the left and right singular vectors corresponding to the largest singular value of ∇ L ( W H ) , construct ∆ W = ua (cid:62) , ∆ H = − av (cid:62) , then (∆ W , ∆ H ) ∈ T ( W , H ) and:
[Equation]
Then it only remains to analyze the | λ | =
||∇ L ( W H )
||
2 cases, construct another convex optimization problem:
[Equation]
suppose Z has SVD Z = U Σ V (cid:62) , as we know that the subgradient of || Z || ∗ can be written as (see
Watson (1992) for a proof):
[Equation]
On the other hand, we know that:
[Equation]
which indicates that H (cid:62) H
= V Σ V (cid:62) and W W (cid:62) =
U Σ U (cid:62) .
Combine them with (C.4) we have:
[Equation]
Note that | λ | =
||∇ L ( W H )
|| 2 , then by (C.11) we know that −∇ L ( W H ) ∈ | λ | ∂ (cid:107) W H (cid:107) ∗ .
Then by the strict convexity of (C.10) we know W H is the global minimum of it.
In addition, we have (cid:107) W (cid:107)
[Equation]
F + (cid:107)
[Equation]
Σ
[Equation]
) = 2 (cid:107)
[Equation]
(cid:107) ∗ .
In addition, previous works (Haeffele & Vidal, 2015) have shown that:
[Equation]
which is equivalent to:
[Equation]
Now for any ( W (cid:48) , H (cid:48) ) , we know that:
[Equation]
which indicates ( W , H ) must attain global minimum of the following optimization problem:
[Equation]
If ( W , H ) does not satisfy neural collapse conditions, by the optimality of neural collapse solution (Theorem 3.3) we know there exists another point ( W (cid:48) , H (cid:48) )
such that L ( W (cid:48) , H (cid:48) )
<
L ( W , H ) and (cid:107) W (cid:107) 2 F + (cid:107)
[Equation]
(cid:107)
[Equation]
F
[Equation]
(cid:107)
[Equation]
(cid:48) (cid:107) 2 F thus:
[Equation]
which contradicts with the global optimality of ( W , H ) , thus ( W , H ) must satisfy all of the neural collapse conditions and we ﬁnish the proof.
D A DDITIONAL E MPIRICAL R ESULTS
Gradient Descent on the ULPM Objective.
We conduct experiments on the ULPM objective (2.3) to support the results of convergence toward neural collapse in our theories.
We set N = 10 , K = 5 , and d = 20 , and used gradient descent with a learning rate of 5 to run 10 5 epochs.
We characterize the dynamics of the training procedure in Figure 2 based on four aspects: (1) Relative variation of the centered class-mean feature norms (i.e., Std ( (cid:107) ¯ h k − ¯ h (cid:107) ) /
Avg ( (cid:107) ¯ h k − ¯ h (cid:107) ) ) and the variation of the classiﬁer’s norms (i.e., Std ( (cid:107) ¯ w k (cid:107) ) /
Avg ( (cid:107) ¯ w k (cid:107) ) ).
(2) Within-class variation of the last layer features (i.e., Avg ( (cid:107) h k,i − h k (cid:107) ) /
Avg ( (cid:107) h k
i − ¯ h (cid:107) ) ).
The cosines between pairs of last layer features (i.e.,
Avg ( | cos(¯ h k , ¯ h k (cid:48) )
+1 /
( K − 1) | ) ) and that of the classiﬁers (i.e., Avg ( | cos( ¯ w k , ¯ w k (cid:48) )
+1 /
( K − 1) | ) ).
(4) The distance between the normalized centered classiﬁer and the normalized last layer feature (i.e.,
Avg ( | (¯ h k − ¯ h ) / (cid:107) ¯ h k − ¯ h (cid:107) − ¯ w k / (cid:107) ¯ w k
cid:107)| ) ).
Empirically, we observe that all four quantities decrease at approximately the rate
O
(1 / (log( t ))) .
[Figure]
Details of Realistic Training.
In the real data experiments, we trained the VGG-13 (Simonyan & {+Zisserman, 2014) and ResNet
+}
(He et al., 2016) on MNIST (LeCun et al., 1998), KMNIST (Clanuwat et al., 2018), FashionMNIST (Xiao et al., 2017) and CIFAR-10 datasets (Krizhevsky et al., 2009) without weight decay, and with a learning rate of 0.01, momentum of 0.3, and batch size of 128.
The metrics are deﬁned similarly to the ULPM case and the experiment results are reported in Figure 1, 3, 4, 5, 6, [Figure] and 8.
All experiments were run in Python (version 3.6.9) on Google Colab. It [Figure] [Figure]
[Figure]
The x -axis in the ﬁgures are set to have log(log( t )) scales and the y -axis in the ﬁgures are set to have log scales.
[Figure] [Figure] [Figure] should be mentioned that we can observe that the variation of classiﬁer norm stays at a low value and do not decrease in some settings (e.g., Figure 1 (a)), this phenomenon is also found in Figure 2 of (Papyan et al., 2020), which might attribute to the network architecture (e.g., batch normalization) and characteristics of real-world datasets.
Overall, we can ﬁnd that neural collapse occurs for various network architectures and datasets under an unconstrained setting, which provides sound support for our theory.
E L IMITATION AND FUTURE DIRECTIONS
Currently, our analysis is still limited in the terminal phase of training and only studies the training behavior after data is perfectly separated.
Such assumption is necessary for all implicit regularization analysis (Lyu & Li, 2019; Ji et al., 2020; Nacson et al., 2019a) under a nonlinear setting, and how to fully characterize the early time training dynamics is still an open problem.
Moreover, we have provided a loss landscape analysis in the paper showing that there is no spurious minimum in the tangent space.
To guarantee the training dynamics do not get stuck in saddle points and converge to neural collapse solution, we need to introduce randomness in the training dynamics (e.g. use stochastic gradient descent (Ge et al., 2015)),
otherwise simple gradient ﬂow may get stuck in saddle points as we have shown in Example 3.1.
How to characterize the training dynamics of stochastic gradient descent in the nonlinear setting is also an open problem.
It is an interesting direction to further build a complete characterization of convergence to neural collapse solutions by addressing the early time training dynamics and studying the stochastic gradient descent.