Year: 2021

Abstract:
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 classifiers 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 flow 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.

heading: 1 INTRODUCTION
text: Deep learning has achieved state-of-the-art performance in various applications (LeCun et al., 2015), such as computer vision (Krizhevsky et al., 2012), natural language processing (Brown et al., 2020), and scientific discovery (Long et al., 2018; Zhang et al., 2018). Despite the empirical success of deep learning, how gradient descent or its variants lead deep neural networks to be biased towards solutions with good generalization performance on the test set is still a major open question. To develop a theoretical foundation for deep learning, many studies have investigated the implicit bias of gradient descent in different settings (Li et al., 2018; Amari et al., 2020; Vaswani et al., 2020; Soudry et al., 2018; Lyu & Li, 2019; Arora et al., 2019).
It is well acknowledged that well-trained end-to-end deep architectures can effectively extract features relevant to a given label. Although theoretical analysis of deep learning has been successful in recent years (Arora et al.; Goldblum et al., 2019), most of the studies that aim to analyze the properties of the final output function fail to understand the features learned by neural networks. Recently, in Papyan et al. (2020), the authors observed that the features in the same class will collapse to their mean and the mean will converge to an equiangular tight frame (ETF) during the terminal phase of training, that is, the stage after achieving zero training error. This phenomenon, namely, neural collapse (Papyan et al., 2020), provides a clear view of how the last-layer features in the neural network evolve after interpolation and enables us to understand the benefit of training after achieving zero training error to achieve better performance in terms of generalization and robustness. To theoretically analyze the neural collapse phenomenon, Fang et al. (2021) proposed the layer-peeled model (LPM) as a simple surrogate for neural networks, where the last-layer features are modeled as free optimization variables. In particular, in a balanced K-class classification problem using a neural network with d neurons in the last hidden layer, the LPM takes the following form:
min W ,H
1
n n∑ i=1 L (Whi,yi) , s.t. 1 2 ||W ||2F ≤ C1, 1 2 ||H||2F ≤ C2, (1.1)
where C1, C2 are positive constants. Here, W = [w1,w2, · · · ,wK ]> ∈ RK×d is the weight of the final linear classifier, H = [h1,h2, · · · ,hn] ∈ Rd×n is the feature of the last layer and yi is the corresponding label. The intuition behind the LPM is that modern deep networks are often highly over-parameterized, with the capacity to learn any representations of the input data. It has been shown
that an equiangular tight frame (ETF), i.e., feature with neural collapse, is the only global optimum of the LPM objective (1.1) (Fang et al., 2021; Lu & Steinerberger, 2020; Wojtowytsch & E, 2020; Zhu et al., 2021).
However, feature constraints in LPMs are not equivalent to weight decay used in practice. In this study, we directly deal with the unconstrained model and show that gradient flow can find those neural collapse solutions without the help of explicit constraints and regularization. To do this, we build a connection between the neural collapse and recent theories on max-margin implicit regularization (Lyu & Li, 2019; Wei et al., 2018), and use it to provide a convergence result to the first-order stationary point of a minimum-norm separation problem. Furthermore, we illustrate that the crossentropy loss enjoys a benign global landscape where all the critical points are strict saddles in the tangent space, except for the only global minimizers that exhibit the neural collapse phenomenon. Finally, we verify our insights via empirical experiments.
In contrast to previous theoretical works on neural collapse, our analysis does not incorporate any explicit regularization or constraint on features. A comparison with other results can be found in Table 1 and we defer a detailed discussion to Section 5.2. The reasons we investigate the unregularized objective are summarized as follows:
1. Feature regularization or constrain is still not equivalent to weight decay used in practice. However, previous studies have justified that neural networks continue to perform well without any regularization or constraint (Zhang et al., 2021). Moreover, it is proved that SGD with exponential learning rate on unconstrained objective is equivalent to SGD with weight decay. (Li & Arora, 2019).
2. As shown in this study, neural collapse exists even under an unconstrained setting, which implies the emergence of neural collapse should be attributed to gradient descent and cross-entropy loss rather than explicit regularization.
3. Regularization or constraint feature constraint can be barriers for existing theories of neural networks (Jacot et al., 2018; Lyu & Li, 2019). By allowing features to be totally free, we hope our results can inspire further analysis to plug in a realistic neural network.

heading: 1.1 CONTRIBUTIONS
text: The contributions of the present study can be summarized as follows.
• We build a relationship between the max-margin analysis (Soudry et al., 2018; Nacson et al., 2019b; Lyu & Li, 2019) and the neural collapse and provide the implicit bias analysis to the
feature rather than the output function. Although both parameters and features diverge to infinity, we prove that the convergent direction is along the direction of the minimum-norm separation problem.
• Previous works (Lyu & Li, 2019; Ji et al., 2020) only prove that gradient flow on homogeneous neural networks will converge to the KKT point of the corresponding minimum-norm separation problem. However, the minimum-norm separation problem remains a highly non-convex problem and a local KKT point may not be the neural collapse solution. In this study, we perform a more detailed characterization of the convergence direction via landscape analysis.
• Previous analysis about neural collapse relies on the explicit regularization or constraint. In this study, we show that the implicit regularization effect of gradient flow is sufficient to lead to a neural collapse solution. The emergence of neural collapse should be attributed to gradient descent and loss function, rather than explicit regularization or constraint. We put detailed discussion in Section 5.2.

heading: 1.2 RELATED WORKS
text: Implicit Bias of Gradient Descent: To understand how gradient descent or its variants helps deep learning to find solutions with good generalization performance on the test set, a recent line of research have studied the implicit bias of gradient descent in different settings. For example, gradient descent is biased toward solutions with smaller weights under `2 loss (Li et al., 2018; Amari et al., 2020; Vaswani et al., 2020) and will converge to large margin solution while using logistic loss (Soudry et al., 2018; Nacson et al., 2019b; Lyu & Li, 2019; Chizat & Bach, 2020; Ji et al., 2020). For linear networks, Arora et al. (2019); Razin & Cohen (2020); Gidel et al. (2019) have shown that gradient descent determines a low-rank approximation.
Loss Landscape Analysis: Although the practical optimization problems encountered in machine learning are often nonconvex, recent works have shown the landscape can enjoy benign properties which allow further analysis. In particular, these landscapes do not exhibit spurious local minimizers or flat saddles and can be easily optimized via gradient-based methods (Ge et al., 2015). Examples include phase retrieval (Sun et al., 2018), low-rank matrix recovery (Ge et al., 2016; 2015), dictionary learning (Sun et al., 2016; Qu et al., 2019; Laurent & Brecht, 2018) and blind deconvolution (Lau et al., 2019).

heading: 2 PRELIMINARIES AND PROBLEM SETUP
text: In this paper, || · ||F denotes the Frobenius norm, ‖ · ‖2 denotes the matrix spectral norm, ‖ · ‖∗ denotes the nuclear norm, ‖ · ‖ denotes the vector l2 norm and tr(·) is the trace of matrices. We use [K] := {1, 2, · · · ,K} to denote the set of indices up to K.

heading: 2.1 PRELIMINARIES
text: We consider a balanced dataset with K classes ⋃K k=1{xk,i}ni=1. A standard fully connected neural network can be represented as: f (x;Wfull) = bL + WLσ (bL−1 + WL−1σ (· · ·σ (b1 + W1x))) . (2.1)
Here Wfull = (W1,W2, · · · ,WL) denote the weight matrices in each layer, (b1, b2, · · · , bL) denote the bias terms, and σ(·) denotes the nonlinear activation function, for example, ReLU or sigmoid. Let hk,i = σ (bL−1 + WL−1σ (· · ·σ (b1 + W1xk,i))) ∈ Rd denote the last layer feature for data xk,i and h̄k = 1n ∑n i=1 hk,i denotes the feature mean within the k-th class. To provide a formal definition of neural collapse, we first introduce the concept of a simplex equiangular tight frame (ETF): Definition 2.1 (Simplex ETF). A symmetric matrix M ∈ RK×K is said to be a simplex equiangular tight frame (ETF) if
M =
√ K
K − 1 Q(IK −
1
K 1K1
> K). (2.2)
Where Q ∈ Rd×K is a matrix with orthogonal columns.
Let W ∈ RK×d = WL = [w1,w2, · · · ,wK ]> be the weight of the final layer classifier, the four criteria of neural collapse can be formulated precisely as:
• (NC1) Variability collapse: As training progresses, the within-class variation of the activation becomes negligible as these activation collapse to their class mean h̄k = 1n ∑n i=1 hk,i.
||hk,i − h̄k|| = 0, 1 ≤ k ≤ K.
• (NC2) Convergence to Simplex ETF: The vectors of the class-means (after centering by their global-mean) converge to having equal length, forming equal-size angles between any given pair, and being the maximally pairwise-distanced configuration constrained to the previous two properties.
cos(h̄k, h̄j) = − 1
K − 1 , ||h̄k|| = ||h̄j ||, k 6= j.
• (NC3) Convergence to self-duality: The linear classifiers and class-means will converge to align with each other, up to appropriate rescaling, that is, there exist a universal constant C > 0 such that
wk = Ch̄k, k ∈ [K].
• (NC4) Simplification to Nearest Class-Center For a given deepnet activation:
h = σ (bL−1 + WL−1σ (· · ·σ (b1 + W1x))) ∈ Rd,
the network classifier converges to choose whichever class has the nearest train class-mean
arg min k 〈wk,h〉 → arg min k ∥∥h− h̄k∥∥ , In this paper, we say that a point W ∈ RK×d,H ∈ Rd×nK satisfies the neural collapse conditions or is a neural collapse solution if the above four criteria are all satisfied for (W ,H).

heading: 2.2 PROBLEM SETUP
text: We mainly focus on the neural collapse phenomenon, which is only related to the classifiers and features in the last layer. Since the general analysis of the highly non-smooth and non-convex neural network is difficult, we peel down the last layer of the neural network and propose the following unconstrained layer-peeled model (ULPM) as a simplification to capture the main characters related to neural collapse during training dynamics. A similar simplification 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. It should be mentioned that although Mixon et al. (2020); Han et al. (2021) also studied the unconstrained model, their analysis adopted an approximation to real dynamics and was highly dependent on the `2 loss function, which is rarely used in classification tasks. Compared with their works, we directly deal with the real training dynamics and cover the most popular cross-entropy loss in classification tasks.
Let W = [w1,w2, · · · ,wK ]> ∈ RK×d and H = [h1,1, · · · ,h1,n,h2,1, · · · ,hK,n] ∈ Rd×Kn be the matrices of classifiers 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 defined as follows:
min W ,H L(W ,H) = min W ,H − K∑ k=1 n∑ i=1 log ( exp(w>k hk,i)∑K j=1 exp(w > j hk,i) ) . (2.3)
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 flow of the objective function:
dW (t) dt = −∂L(W (t),H(t)) ∂W , dH dt = −∂L(W (t),H(t)) ∂H .

heading: 3 MAIN RESULTS
text: 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 flow 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 find an Example 3.1 in Section 3.2, which shows that there exist some bad KKT points such that a simple gradient flow 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).

heading: 3.1 CONVERGENCE TO THE FIRST–ORDER STATIONARY POINT
text: Heuristically speaking, the simplex ETF (Definition 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 classifiers. 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. (2018); Nacson et al. (2019b); Lyu & Li (2019); Ji et al. (2020). First, we illustrate the connection between the margin and neural collapse. Recall that the margin of a single data point xk,i and the associated feature hk,i is qk,i(W ,H) := w > k hk,i−maxj 6=kw>j hk,i. Then, the margin of the entire dataset can be defined as
qmin(W ,H) := min k∈[K],i∈[n] qk,i(W ,H).
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 satisfies
qmin(W ,H) ≤ ‖W ‖2F + ‖H‖2F
2(K − 1) √ n
and the equality holds if and only if (W ,H) satisfies the neural collapse conditions with ‖W‖F = ‖H‖F .
Based on this finding, we present an analysis of the convergence of the gradient flow on the ULPM (2.3). Following Lyu & Li (2019), we link the gradient flow 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 flow at time t, if there exist a time t0 such that L(W (t0),H(t0)) < log 2, then any limit point of
{(Ĥ(t), Ŵ (t)) := ( H(t)√ ‖W (t)‖2F + ‖H(t)‖2F , W (t)√ ‖W (t)‖2F + ‖H(t)‖2F )}
is along the direction (i.e., a constant multiple) of a Karush-Kuhn-Tucker (KKT) point of the following minimum-norm separation problem:
min W,H
1 2 ||W ||2F + 1 2 ||H||2F
s.t.w>k hk,i −w>j hk,i ≥ 1, k 6= j ∈ [K], i ∈ [n]. (3.1)
Remark 3.1. Here we write (3.1) as a constraint problem, but the constraint is introduced by the implicit regularization effect of gradient flow on our ULPM objective (2.3). Since the training dynamics will diverge to infinity, 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 t0 such that L(W (t0),H(t0)) < log 2. Note that the loss function can be rewritten as:
L(W ,H) = K∑ k=1 n∑ i=1 log(1 + ∑ j 6=k exp(wjhk,i −wkhk,i)) (3.2)
the requirement L(W ,H) < log 2 implies that wkhk,i −wjhk,i ≥ 0 for any k, j ∈ [K], i ∈ [n], which is equivalent to qmin(W ,H) > 0; that is, every feature is separated perfectly by the classifier. 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 justified 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 flow is restricted to the max-margin directions, which usually have good robustness and generalization performance. In general, the KKT conditions are not sufficient 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 satisfies 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.

heading: 3.2 SECOND–ORDER LANDSCAPE ANALYSIS
text: 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 specific, 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 flow.
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 infinity. 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, first-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:
W = H =  1 −1 0 0−1 1 0 00 0 1 −1 0 0 −1 1  . 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:
Λ =  0 0 12 1 2 0 0 12 1 2
1 2 1 2 0 0 1 2 1 2 0 0  . 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):
L(W ,H,Λ) = 1 4 ‖W ‖2F + 1 4 ‖H‖2F − 4∑ i=1 ∑ j 6=i λi,j( 1 2 wihi − 1 2 wjhi − 1).
Simply take derivatives for W ,H and Λ we find that it satisfies the KKT conditions. However, the gradient of (W ,H) is:
∇WL(W ,H) = ∇HL(W ,H) = − 2 + 2e−2
2 + 2e−2 + 2e2  1 −1 0 0−1 1 0 00 0 1 −1 0 0 −1 1  . We can see that the directions of the gradient and the parameter align with each other (i.e., W is parallel to ∇WL(W ,H), and H is parallel to ∇HL(W ,H)), which implies that simple gradient descent may get stuck in this direction and only grow the parameter norm. However, if we construct:
W ′ = H ′> =
√ 1
1 + 2α2  1 + α −1 + α α α−1 + α 1 + α α α−α −α 1− α −1− α −α −α −1− α 1− α  , By simple calculation, we find that f(α) := L(W ′,H ′) satisfies f ′(α) = 0, f ′′(α) < 0. Since ‖W ′‖F = ‖W ‖F , ‖H ′‖F = ‖H‖F and ||W ′ −W ||2F + ||H ′ −H||2F → 0 as α→ 0, this result implies that for any > 0, we can choose appropriate α such that:
||W ′||2F = ||W ||2F , ||H ′||2F = ||H||2F , ||W ′ −W ||2F + ||H ′ −H||2F < ,L(W ′,H ′) < L(W ,H).
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. (2016); Ge et al. (2015) 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 ′,H ′) for any ||W ′||2F + ||H ′||2F = ||W ||2F + ||H||2F ) if only if (W ,H) satisfies 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) satisfies 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 first introduce the tangent space: Definition 3.1 (tangent space). The tangent space of (W ,H) is defined as a set of directions that are orthogonal to (W ,H) :
T (W ,H) = {∆W ∈ RK×d,∆H ∈ Rd×nK) : tr(W>∆W ) + tr(H>∆H) = 0}
Our next result justifies 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 ‖W ‖F 6= ‖H‖F ), then there exists a direction (∆W ,∆H) ∈ T (W ,H) and constant M > 0 such that for any 0 < δ < M ,
L(W + δ∆W ,H + δ∆H) < L(W ,H). (3.3) Further more, it implies that for any > 0 there exists a point (W ′,H ′) such that:
||W ′||2F + ||H ′||2F = ||W ||2F + ||H||2F , ||W ′ −W ||2F + ||H ′ −H||2F < ,L(W ′,H ′) < L(W ,H).
(3.4)
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 first-order analysis fails to explain the prevalence of neural collapse. Here the decreasing direction is obtained by analyzing the Riemannian Hessian matrix and finding its eigenvector corresponding to a negative eigenvalue, which further indicates that these points are first-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 flow dynamics only find neural collapse directions.

heading: 4 EMPIRICAL RESULTS
text: To evaluate our theory, we trained the ResNet18 (He et al., 2016) on both MNIST (LeCun et al., 1998) and CIFAR-10 (Krizhevsky et al., 2009) datasets without weight decay, and tracked how the last layer features and classifiers 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:
1. The variation of both feature norm (i.e., Std(‖h̄k − h̄‖)/Avg(‖h̄k − h̄‖)) and classifier norm in the last layer (i.e., Std(‖w̄k‖)/Avg(‖w̄k‖)).
2. Within-class variation for last layer features (i.e., Avg(‖hk,i − hk‖)/Avg(‖hk,i − h̄‖)). 3. Average cosine between last layer features (i.e., Avg(| cos(h̄k, h̄k′) + 1/(K− 1)|)) and that
of last layer classifiers (i.e., Avg(| cos(w̄k, w̄k′) + 1/(K − 1)|)). 4. Self-duality distance between features and classifiers corresponding to the same class in the
last layer. (i.e., Avg(|(h̄k − h̄)/‖h̄k − h̄‖ − w̄k/‖w̄k‖|))
Here a smaller value of each metric indicates a closer state to a neural collapse solution. Specifically, 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 shown in the Figure 1, 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.

heading: 5 CONCLUSION AND DISCUSSION

heading: 5.1 CONCLUSION
text: 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 flow of the ULPM converges to the KKT point of a minimum-norm separation problem, where the global optimum satisfies 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 classification problems.

heading: 5.2 RELATIONSHIP WITH OTHER RESULTS ON NEURAL COLLAPSE
text: Theoretical analysis of neural collapse was first 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 simplified 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 `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 classification tasks. The analysis relies highly on the `2 loss to obtain a closed-form gradient flow 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 modified 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 flow on cross-entropy. We show that implicit regularization is sufficient to lead the dynamics to converge to the neural collapse solution without the explicit regularization.

heading: A ELEMENTS OF OPTIMIZATION
text: In this section, we introduce some basic definitions and theory about optimization. In the following discussion, we consider a standard form inequality constrained optimization problem:
min x∈Rd f(x)
s.t. gi(x) ≤ 0, i ∈ [n]. (A.1)
In addition, we assume all of those functions f and gi are twice differentiable. A point x ∈ Rd is said to be feasible if and only if it satisfies all of the constraints in (A.1), i.e., gi(x) ≤ 0, i ∈ [n]. And the Lagrangian of problem (A.1) is defined as following:
L(x, λ) = f(x) + n∑ i=1 λigi(x).
A.1 KARUSH KUHN TUCKER CONDITIONS
Now let’s first introduce the definition of Karush Kuhn Tucker (KKT) point and approximate KKT point. Here we follows the definition of ( , δ)-KKT point as in Lyu & Li (2019). Definition A.1 (Definition 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:
1. ∇f(x) + ∑n i=1 λi∇gi(x) = 0
2. λigi(x) = 0, i ∈ [n] Definition A.2 (Definition of ( , δ)-KKT point). For any , δ > 0, a feasible point of (A.1) is said to be ( , δ)-KKT point of problem (A.1) if there exist λ1, λ2, · · · , λn ≥ 0 such that:
1. ||∇f(x) + ∑n i=1 λi∇gi(x)|| ≤
2. λigi(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 qualification (MFCQ): Definition A.3 (Mangasarian-Fromovitz constraint qualification (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 ∈ Rd such that:
〈∇xgi(x), v〉 > 0, i ∈ [n]. (A.2)
Moreover, when MFCQ holds we can build a connection between approximate KKT point and KKT point, see detailed proof in Dutta et al. (2013): Theorem A.1 (Relationship between Approximate KKT point and KKT point). Let{ xk ∈ Rd : k ∈ N } be a sequence of feasible points of (A.1) , { k > 0 : k ∈ N} and {δk > 0 : k ∈ N} be two sequences of real numbers such that xk is an ( k, δk)-KKT point for every k, and k → 0, δk → 0. If xk → x as k → +∞. If MFCQ (A.3) holds at x, then x is a KKT point of (P ).

heading: B OMITTED PROOFS FROM SECTION 3.1
text: Recall our ULPM problem:
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) ) .
Let’s review some basic notations defined in the main body. Let sk,i,j = w>k hk,i −w>j hk,i,∀k ∈ [K], i ∈ [n], j ∈ [K], the margin of a single feature hk,i is defined to be qk,i(W ,H) := minj 6=k sk,i,j = w > k hk,i − maxj 6=kw>j hk,i. We define the margin of entire dataset as qmin = qmin(W ,H) = mink∈[1,K],i∈[1,n] qk,i(W ,H). We first prove Theorem 3.1 in the mainbody.
Proof of Theorem 3.1. First we can find that the margin will not change if we minus a vector a for all wj , so if we denote the mean of classifier w̃i = wi − 1K ∑K i=1 wi and then we have w>k hk,i −maxj 6=kw>j hk,i = w̃>k hk,i −maxj 6=k w̃>j hk,i ≥ qmin(W ,H), that is:
w̃>k hk,i − w̃>j hk,i ≥ qmin(W ,H),∀j 6= k ∈ [K], i ∈ [n]. (B.1) Note that ∑K j=1 w̃ > j hk,i = 0 then sum this inequality over j we have:
(K − 1)w̃>k hk,i − ∑ j 6=k w̃>j hk,i = Kw̃ > k hk,i ≥ (K − 1)qmin(W ,H),∀k ∈ [K], i ∈ [n].
By Cauchy inequality, we have: 1
2 ( 1√ n ||w̃k||22 + √ n||hk,i||22) ≥ w̃>k hk,i ≥ K − 1 K qmin(W ,H). (B.2)
Sum (B.2) over k and i we have: 1
2
√ n(||W̃ ||2F + ||H||2F ) ≥ n(K − 1)qmin(W ,H). (B.3)
On the other hand, we know that:
||W̃ ||2F = K∑ i=1 ||wi − 1 K K∑ i=1 wi||22 ≤ K∑ i=1 ||wi||22 = ||W ||2F .
Then we can conclude that:
qmin(W ,H) ≤ ‖W ‖2F + ‖H‖2F
2(K − 1) √ n
as desired. When the equality holds, first we have ||W̃ ||2F = ||W ||2F which is equivalent to 1 K ∑K i=1 wi = 0, w̃i = wi. Take it back into (B.3), then we must have all of the equality holds in (B.2) and (B.1), which give us:
wk = √ nhk,i, ||wk||22 = n||hk,i||22 = ||W ||2F + ||H||2F 2K .
Take this into (B.1) we have:
hk,i = hk,i′ ,h > k,ihj,i′ = w > k wj = − ||W ||2F + ||H||2F 2K(K − 1) √ n ,
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 defined on θ by matching the corresponding parameters. Denote ρ = ‖θ‖ as the norm of θ and γ̃ = − log(e
L(θ)−1) ρ2 . Now we can state our first lemma to show how training dynamics of gradient
flow on ULPM objective (B) is related to a KKT point of (3.1). Lemma B.1. If there exist a time t0 such that L(θ(t0)) < log 2, then for any t > t0 θ̃ := θ/qmin(θ)
1/2 is an ( , δ) - approximate KKT point of the following minimum-norm separation problem. More precisely, we have
=
√ 2(1− β(t))
γ̃(t) , δ = K2(K − 1)n 2eγ̃(t)qmin(t) ,
where: β = 〈 θ
||θ||2 , dθ dt /||dθ dt ||2〉
is the angle between θ and its corresponding gradient.
Proof. The training dynamics is given by gradient flow:
dθ dt = −∂L(θ) ∂θ .
Then by the chain rule we have:
− dL(θ) dt = −∂L ∂θ dθ dt = ∥∥∥∥dθdt ∥∥∥∥2 . (B.4)
It indicates that the loss function L is monotonically decreasing. If L(θ(t0)) < log 2, we have L(θ(t)) < log 2,∀t > t0. On the other hand, note that
L(θ(t)) = K∑ k=1 n∑ i=1 log(1 + ∑ j 6=k e−sk,i,j(t)) ≥ log(1 + exp(−qmin(t))), (B.5)
which gives us qmin(t) > 0,∀t > t0. Let g = dθdt , note that we can rewrite the ULPM objective function (B) asL(θ) = ∑K k=1 ∑n i=1 log(1+∑
j 6=k e −sk,i,j ). By the chain rule and the gradient flow equation we have
g = −∂L ∂θ K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l gk,i,j ,
where gk,i,j is the gradient of sk,i,j , i.e. gk,i,j = ∇θsk,i,j(θ). Now let g̃k,i,j = gk,i,j/q1/2min = ∇θsk,i,j(θ̃) and construct λk,i,j = ρ||g||2 e−sk,i,j 1+ ∑ l 6=k e −sk,i,l , we only need to show:
||θ̃ − K∑ k=1 n∑ i=1 ∑ j 6=k λk,i,j g̃k,i,j ||22 ≤ 1− β γ̃ , (B.6)
K∑ k=1 n∑ i=1 ∑ j 6=k λk,i,j(sk,i,j(θ̃)− 1) ≤ K2(K − 1)n 2eqminγ̃ . (B.7)
To prove (B.6), we only need to compute (Recall that θ̃ = θ/qmin(θ)1/2):
||θ̃ − K∑ k=1 n∑ i=1 ∑ j 6=yn λk,i,j g̃k,i,j ||22 = ρ2 qmin || θ ||θ||2 − g ||g||2 ||22 = ρ2 qmin (2− 2β).
Note that:
γ̃ = − log(eL(θ) − 1)
ρ2 , L(θ) = N∑ n=1 log(1 + ∑ j 6=yn e−snj ) ≥ log(1 + exp(−qmin)). (B.8)
Then we have the following inequality:
γ̃ ≤ qmin ρ2 . (B.9)
Take this back into (B.8) we have (B.6) as desired.
To prove (B.7), first by our construction:
K∑ k=1 n∑ i=1 ∑ j 6=k λk,i,j(sk,i,j(θ̃)− 1) = ρ qmin||g||2 K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l (sk,i,j − qmin).
(B.10)
Note that ||g||2 ≥ 〈 g, θ||θ||2 〉 = 1ρ 〈g, θ〉 and 〈gk,i,j , θ〉 = 2sk,i,j since sk,i,j = w > k hk,i −w>j hk,i, we have:
||g||2 ≥ 1 ρ 〈g, θ〉 =1 ρ K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l 〈gk,i,j , θ〉
= 2
ρ K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l sk,i,j
≥2 ρ qmin K K∑ k=1 n∑ i=1 ∑ j 6=k e−sk,i,j (since sk,i,j ≥ qmin > 0 and e−sk,i,l ≤ 1)
≥2 ρ qmin K e−qmin .
(B.11) Take this inequality back into the (B.10) we have:
K∑ k=1 n∑ i=1 ∑ j 6=k λk,i,j(sk,i,j(θ̃)− 1) ≤ Kρ2 2q2min K∑ k=1 n∑ i=1 ∑ j 6=k
eqmin−sk,i,j 1 + ∑ l 6=k e −sk,i,l (sk,i,j − qmin)
≤ Kρ 2
2q2min K∑ k=1 n∑ i=1 ∑ j 6=k eqmin−sk,i,j (sk,i,j − qmin)
≤ K 2qminγ̃ K∑ k=1 n∑ i=1 ∑ j 6=k eqmin−sk,i,j (sk,i,j − qmin)
≤ K 2(K − 1)n 2eqminγ̃ .
Where the last inequality is obtained from the fact xe−x ≤ 1e ,∀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 ( , δ)-KKT point, if we can show ( , δ) 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 qmin(t) evolve along time. Now we provide the following lemmas to illustrate their dynamics. The first lemma aims at proving that the norms of parameter ρ(t) and γ̃(t) are monotonically increasing.
Lemma B.2. If there exist t0 such that L(θ(t0)) < log 2, then for any t > t0 we have:
dρ2
dt > 0,
dγ̃ dt ≥ 0. (B.12)
Proof. We can disentangle the whole training dynamics into the following two parts:
• the radial part: v := θ̂θ̂> dθdt ,
• the tangent part: u = (I − θ̂θ̂>)dθdt .
First analyze the radial part, by the chain rule: ||v||2 = |θ̂> dθdt | = | 1 ρ 〈 θ, dθdt 〉 | = | 1ρ 1 2 dρ2 dt |. For dρ2
dt , we have the following equation:
1
2
dρ2
dt =
〈 θ, dθ
dt
〉 = 2 K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l sk,i,j , (B.13)
where the last equality holds by equation (B.11).Then when t > t0, we have shown that qmin(t) ≥ 0, combine this with the fact sk,i,j ≥ qmin we obtain the first inequality in (B.12)
1
2
dρ2
dt =2 K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l sk,i,j
≥2 K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l qmin ≥ 0.
(B.14)
Next, we aim to prove the monotonicity of γ̃(t), compute the derivative of γ̃(t) we have:
γ̃ = − log(eL(θ) − 1)
ρ2 ,
d dt log γ̃ = d dt (log(− log(eL(θ) − 1))− 2 log ρ), (B.15)
d dt log(− log(eL(θ) − 1)) = 1 log(eL(θ) − 1) eL(θ) eL(θ) − 1 dL(θ) dt ≥ −dL(θ) dt 1
qmin
eL(θ)
eL(θ) − 1 . (B.16)
Recall that we have:
1
2
dρ2
dt =2 K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l sk,i,j
≥2 K∑ k=1 n∑ i=1
∑ j 6=k e −sk,i,j
1 + ∑ j 6=k e −sk,i,j qmin
= K∑ k=1 n∑ i=1 log(1 + ∑ j 6=k e−sk,i,j ) 1 log(1 + ∑ j 6=k e −sk,i,j )
∑ j 6=k e −sk,i,j
1 + ∑ j 6=k e −sk,i,j qmin
≥2 K∑ k=1 n∑ i=1 log(1 + ∑ j 6=k e−sk,i,j ) eL(θ) − 1 L(θ)eL(θ) qmin
=2 eL(θ) − 1 eL(θ) qmin.
(B.17)
Then second last line is because the definition of the loss function L(θ) = ∑K k=1 ∑n i=1 log(1 +∑
j 6=k e −sk,i,j ) ≥ log(1 + ∑ j 6=k e −sk,i,j ) and the monotonicity of e x−1 xex (in fact, d dx ex−1 xex =
e−x(x−ex+1) x2 ≤ 0,∀x > 0).
As a result, we notice that
1
2
d dt log γ̃(t) ≥ −
( 1
2
dρ2
dt )−1 dL dt − d dt log ρ.
At the same time, we notice that ‖v‖2 = 1ρ2 ( 1 2 dρ2 dt )2 = 12 dρ2 dt · d dt log ρ on the one hand, and by the chain rule: d
dt θ̂ =
1
ρ2
( ρ dθ
dt − dρ dt θ
) = 1
ρ2
( ρ dθ dt − ( θ̂> dθ dt ) θ ) = u ρ .
Combine this with the radial term:
−dL dt = ∥∥∥∥dθdt ∥∥∥∥2 = ‖v‖2 + ‖u‖2 = 12 dρ2dt · ddt log ρ+ ρ2 ∥∥∥∥∥dθ̂dt ∥∥∥∥∥ 2 .
Dividing 12 dρ2 dt on both sides, we have
−dL dt · ( 1 2 dρ2 dt )−1 = d dt log ρ+ ( d dt log ρ )−1 ∥∥∥∥∥dθ̂dt ∥∥∥∥∥ 2 , (B.18)
d dt log ρ+
( d
dt log ρ )−1 ∥∥∥∥∥dθ̂dt ∥∥∥∥∥ 2 ≤ −dL(θ) dt 1
qmin
eL(θ)
2(eL(θ) − 1) .
Now by equation (B.15) and (B.16) we obtain:
1
2
d dt log γ̃ ≥ −dL(θ) dt 1
qmin
eL(θ) 2 ( eL(θ) − 1 ) − d dt
log ρ ≥ ( d
dt log ρ )−1 ∥∥∥∥∥dθ̂dt ∥∥∥∥∥ 2 . (B.19)
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(θ(t0)) < log 2, we have γ̃(t) ≥ γ̃(t0) > 0, then we can treat γ̃(t) in Lemma B.1 as a positive constant. The remaining work is to show qmin(t) grows to infinity and β(t)→ 1. To show qmin(t)→∞, it’s equivalent to show L(t)→ 0 and we have the following lemma: Lemma B.3. If there exist t0 such that L(θ(t0)) < log 2, then L(θ(t))→ 0 and qmin(θ(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:
dL(θ) dt = − ∥∥∥∥dθdt ∥∥∥∥2 ≤ −〈dθdt , θ||θ||2 〉2 = −(2 K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l sk,i,j) 2.
Combine it with (B.5),(B.14) and (B.17) we have:
2 K∑ k=1 n∑ i=1 ∑ j 6=k
e−sk,i,j 1 + ∑ l 6=k e −sk,i,l sk,i,j ≥ 2 eL(θ) − 1 eL(θ) qmin ≥ −2 eL(θ) − 1 eL(θ) log(eL(θ) − 1),
which indicates: dL(θ(t))
dt ≤ −4(e L(θ(t)) − 1 eL(θ(t)) log(eL(θ(t)) − 1))2. (B.20)
Since 0 < L(θ(t)) < L(θ(t0)) < log 2, note that: d
dx ( ex − 1 x ) > 0, lim x→0 ex − 1 x = 1,
which implies:
1 < eL(θ(t)) − 1 L(θ(t)) < 1 log(2) , (B.21)
On the other hand, we can find that:
log(eL(θ(t)) − 1) < log(eL(θ(t0)) − 1) < 0, 1 < eL(θ(t)) < 2 (B.22) Combine equation (B.21) and (B.22) together, one can conclude that there exist a constant C > 0 such that for t > t0:
dL(θ(t)) dt ≤ −C(L(θ(t)))2
It further implies that:
C < d
dt (
1
L(θ(t)) ),
then integral on both side we have:
C(t− t0) < 1 L(θ(t)) − 1 L(θ(t0)) ,
and L(θ(t)) = O(1/t).
Thus we must have L(θ(t)) → 0 and combine this with qmin ≥ − log(eL(θ(t)) − 1) we know qmin(θ(t))→∞ as desired.
To bound β(t), we first need a useful lemma to bound the changes of the direction of θ. Lemma B.4. If there exist t0 such that L(θ(t0)) < log 2, then for any t > t0∥∥∥∥∥dθ̂dt
∥∥∥∥∥ ≤ 1γ̃(t0) ddt log ρ. Proof. First we know that: ∥∥∥∥∥dθ̂dt ∥∥∥∥∥ = 1ρ ∥∥∥∥(I − θ̂θ̂>)dθdt ∥∥∥∥ ≤ 1ρ ∥∥∥∥dθdt
∥∥∥∥ , ∥∥∥∥dθdt ∥∥∥∥ = ∥∥∥∥∥∥ K∑ k=1 n∑ i=1 ∑ j 6=k e−sk,i,j 1 + ∑ l 6=k e −sk,i,l gk,i,j ∥∥∥∥∥∥ ≤ K∑ k=1 n∑ i=1 ∑ j 6=k e−sk,i,j 1 + ∑ l 6=k e −sk,i,l ‖gk,i,j‖ .
(B.23) Recall our gk,i,j = ∂sk,i,j ∂θ and sk,i,j = w > k hk,i −w>j hk,i, we can find that ‖gk,i,j‖ ≤ 2ρ. On the other hand, Combine it with (B.13) and (B.9) we have:∥∥∥∥∥dθ̂dt ∥∥∥∥∥ ≤ 1ρ ∥∥∥∥dθdt ∥∥∥∥ ≤ 12qmin dρ 2 dt = ρ2 qmin d dt log ρ ≤ 1 γ̃ d dt log ρ ≤ 1 γ̃(t0) d dt log ρ
as desired. Where the second inequality holds by multiple ρqmin 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 γ̃(t0) > 0 since L(t0) < log 2.
Let’s turn back to β, though we can not show directly that it increase to one, we can find a sequence of time {tm} for each limit point such that β(tm)→ 1 Lemma B.5. If there exist t0 such that L(θ(t0)) < log 2, then for every limit point θ̄ of {θ̂(t) : t ≥ 0}, there exists a sequence of time {tm > 0 : m ∈ N} such that tm → ∞, θ̂ (tm) → θ̄, and β (tm)→ 1.
Proof. Recall that in equation (B.19) we have shown that:
d dt log γ̃ ≥ 2
( d
dt log ρ )−1 ∥∥∥∥∥dθ̂dt ∥∥∥∥∥ 2 . (B.24)
Since ddt log ρ = 1 ρ dρ dt = 1 2ρ2
dρ2
dt = 1 ρ2 〈 θ, dθdt 〉 and:
dθ̂ dt = d dt θ ‖θ‖ = 1 ρ2 (ρ dθ dt − 1 ρ θθ> dθ dt ) = 1 ρ (I − θ̂θ̂>)dθ dt .
Plug them into (B.24) we have:
d dt log γ̃ ≥ 2
∥∥dθ dt ∥∥2 − 〈θ̂, dθdt〉2〈 θ̂, dθdt
〉2 ddt log ρ = 2(β−2 − 1) ddt log ρ. For any t2 > t1 > t0, integrate both sides from time t2 to t1 we have: log γ̃(t2) − log γ̃(t1) ≥ 2 ∫ t2 t1 (β(t)−2 − 1) ddt log ρdt. By the continuity of β we know there exist a time t ∗ such that:
log γ̃(t2)− log γ̃(t1) ≥ 2 ∫ t2 t1 (β(t)−2 − 1) d dt log ρdt
= 2(β(t∗)−2 − 1) ∫ t2 t1 d dt log ρdt = 2(β(t∗)−2 − 1)(log ρ(t2)− log ρ(t1)).
(B.25)
By (B.9) we know that γ̃ ≤ qminρ2 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 tm, first take a sequence of { m > 0,m ∈ N} such that m → 0. We construct tm by induction, suppose we have already find t1 < t2 < · · · < tm−1 satisfy our requirement, since θ̄ is a limit point of {θ̂(t) : t > 0}, then we can find a time sm such that:
‖θ̂(sm)− θ̄‖ ≤ m, log γ̃∞ γ̃(sm) ≤ 3m.
By the monotonicity and continuity of ρ we can find a time s′m such that log ρ(s ′ m)− log ρ(sm) ≤ m. Take t2 = s′m, t1 = sm in (B.25), there exist a time tm such that:
2(β(tm) −2 − 1) ≤ log γ̃(t2)− log γ̃(t1)
log ρ(t2)− log ρ(t1) ≤ 2m (B.26)
on the other hand, by Lemma B.4 we have:
‖θ̂(tm)− θ̄‖ ≤‖θ̂(sm)− θ̄‖+ ‖θ̂(sm)− θ̂(tm)‖
≤ m + 1
γ̃(t0) (log ρ(tm)− log ρ(sm)) ≤ (1 +
1
γ̃(t0) ) m.
(B.27)
Note that 〈 θ, dθdt 〉 > 0, then by definition we know β > 0. Combine equations (B.26) and (B.27) we have β(tm)→ 1 and θ̂(tm)→ θ̄ 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 > t0, (W (t),H(t))/qmin(W (t),H(t)) is an ( √ 2(1−β(t)) γ̃(t) , K2(K−1)n 2γ̃(t)qmin(t)
) -approximate KKT point. We have shown that γ̃(t) > γ̃(t0) > 0 in Lemma B.2, qmin → ∞ in Lemma B.3 and from Lemma B.5 we know for any limit point (W̄ , H̄) of {(Ĥ(t), Ŵ (t)) := ( H(t)√ ‖W (t)‖2F+‖H(t)‖2F , W (t)√ ‖W (t)‖2F+‖H(t)‖2F )}, there exists a sequence of time {tm > 0 : m ∈ N} such that tm → ∞, β(tm) → 1 and (Ĥ(tm), Ŵ (tm)) → (W̄ , H̄). Then (W̄ , H̄) is along the direction of a limit point of a sequence of ( , δ)-approximate KKT point with , δ → 0. On the other hand, we can verify that the problem (3.1) satisfies MCFQ (A.3) by simply setting v = θ, then:
〈∇sk,i,j , θ〉 = 2sk,i,j ≥ 0. 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 flow, 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 satisfies neural collapse conditions. First the constraints in (3.1) can be transformed to be a single constraint by the definition of margin:
∀k 6= j ∈ [K], i ∈ [n], w>k hk,i −w>j hk,i ≥ 1.⇔ qmin(W ,H) ≥ 1. (B.28) Note that the margin is homogeneous:
qmin(αW , αH) = α 2qmin(W ,H),∀α ∈ R.
Then for any point (W ,H) satisfies qmin(W ,H) > 0, after an appropriate scaling α, (αW , αH),∀α2 ≥ 1/qmin(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/qmin(W ,H). And the optimum norm is:
1 2 ||αW ||2F + 1 2 ||αH||2F =
1
2qmin(W ,H) (||W ||2F + ||H||2F ).
Then by Theorem 3.1 we have:
1
2qmin(W ,H) (||W ||2F + ||H||2F ) ≥ 2(K − 1)
√ n.
And the global optimum is attained only when (W ,H) satisfies neural collapse conditions

heading: C OMITTED PROOFS FROM SECTION 3.2
text: To begin with, let’s finish 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:
W = H =  1 −1 0 0−1 1 0 00 0 1 −1 0 0 −1 1  . 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:
Λ =  0 0 12 1 2 0 0 12 1 2
1 2 1 2 0 0 1 2 1 2 0 0  . 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):
L(W ,H,Λ) = 1 4 ‖W ‖2F + 1 4 ‖H‖2F − 4∑ i=1 ∑ j 6=i λi,j( 1 2 wihi − 1 2 wjhi − 1).
Simply take derivatives for W ,H and Λ we can find it satisfies KKT conditions. On the other hand, the gradient of (W ,H) is
∇WL(W ,H) = ∇HL(W ,H) = − 2 + 2e−2
2 + 2e−2 + 2e2  1 −1 0 0−1 1 0 00 0 1 −1 0 0 −1 1  . We can find that the directions of gradient and the parameter align with each other (i.e., W is parallel to ∇WL(W ,H), and H is parallel to ∇HL(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 ′ =
√ 1
1 + 2α2  1 + α −1 + α α α−1 + α 1 + α α α−α −α 1− α −1− α −α −α −1− α 1− α  ,
H ′ =
√ 1
1 + 2α2  1 + α −1 + α −α −α−1 + α 1 + α −α −αα α 1− α −1− α α α −1− α 1− α  . Note that ‖W ′‖F = ‖W ‖F , ‖H ′‖F = ‖H‖F and ||W ′ −W ||2F + ||H ′ −H||2F → 0 as α→ 0. First we can compute:
W ′H ′ = 1
1 + 2α2
 2 + 4α 2 4α2 − 2 −4α2 −4α2
4α2 − 2 2 + 4α2 −4α2 −4α2 −4α2 −4α2 2 + 4α2 4α2 − 2 −4α2 −4α2 4α2 − 2 2 + 4α2
 ,
and:
L(W ′,H ′) =− 4 log e 2
e2 + e 2 2α 2−1 1+2α2 + 2e −2 2α2 1+2α2
.
Our aim is to show that for any > 0, there exist α such that |α| < and L(W ′,H ′) < L(W ,H). By the formulation of L(W ′,H ′), it’s sufficient to show that f(α) , e2 2α2−1 1+2α2 + 2e −2 2α2 1+2α2 < f(0). Then take the derivative of f(α) we have:
f ′(α) = e 4α2−2 1+2α2
( 8α
1 + 2α2 −
8α ( 2α2 − 1 ) (1 + 2α2) 2 ) + 2e − 4α2 1+2α2 ( 16α3 (1 + 2α2) 2 − 8α 1 + 2α2 ) ,
f ′′(α) = e 4α2−2 1+2α2
( 8α
1 + 2α2 −
8α ( 2α2 − 1 ) (1 + 2α2) 2 )2 + 2e − 4α2 1+2α2 ( 16α3 (1 + 2α2) 2 − 8α 1 + 2α2 )2
+ e 4α2−2 1+2α2 ( − 64α 2
(1 + 2α2) 2 +
64 ( 2α2 − 1 ) α2
(1 + 2α2) 3 +
8
1 + 2α2 −
8 ( 2α2 − 1 ) (1 + 2α2) 2 )
+ 2e − 4α2 1+2α2
( 80α2
(1 + 2α2) 2 −
8 1 + 2α2 − 128α
4
(1 + 2α2) 3
) .
Now we can find that f ′(0) = 0 and f ′′(0) = 16( 1e2−1) < 0. Since the function f(α) is continuously twice differentiable, we can conclude that for any > 0, we can choose appropriate α such that:
||W ′||2F = ||W ||2F , ||H ′||2F = ||H||2F , ||W ′ −W ||2F + ||H ′ −H||2F < ,L(W ′,H ′) < L(W ,H).
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 sk,i,j = w>k hk,i − w>j hk,i:
L(W ,H) = K∑ k=1 n∑ i=1 log(1 + ∑ j 6=k exp(−sk,i,j)).
In addition, we can find that centralizing wi does not change the value of sk,i,j . Let w̃i = wi − 1 K ∑K k=1 wk,∀i ∈ [K], then sk,i,j = w>k hk,i −w>j hk,i = w̃>k hk,i − w̃>j hk,i and ∑K i=1 w̃i = 0. First by the strict convexity of ex and Jensen Inequality:
L(W ,H) ≥ K∑ k=1 n∑ i=1 log(1 + (K − 1) exp( 1 K − 1 ∑ j 6=k −sk,i,j))
= K∑ k=1 n∑ i=1 log(1 + (K − 1) exp(−Kw̃ > k hk,i K − 1 )).
(C.1)
Where the last equality is obtained from:∑ j 6=k sk,i,j = ∑ j 6=k w̃>k hk,i − w̃>j hk,i = (K − 1)w̃>k hk,i − ∑ j 6=k w̃>j hk,i = Kw̃ > k hk,i.
Now again by the strict convexity of log(1 + (K − 1) exp(−x)) and Jensen inequality, we have:
L(W ,H) ≥ K∑ k=1 n∑ i=1 log(1 + (K − 1) exp(−Kw̃ > k hk,i K − 1 ))
≥ nK log(1 + (K − 1) exp(− 1 n(K − 1) K∑ k=1 n∑ i=1 w̃>k hk,i))
≥ nK log(1 + (K − 1) exp(− 1 2n(K − 1) K∑ k=1 n∑ i=1 1√ n ‖w̃k‖2 + √ n‖hk,i‖2))
≥ nK log(1 + (K − 1) exp(− 1 2 √ n(K − 1) (‖W ‖2F + ‖H‖2F ))).
(C.2)
Where the last inequality holds since ‖W ‖2F = ∑K k=1 ‖wk‖2 ≥ ∑K k=1 ‖wk‖2− 1 K ‖ ∑K k=1 wk‖2 =∑K
k=1 ‖w̃k‖2. When all of the above inequality reduce to equality, we must have:
1. ∑K i=1 wi = 0, w̃i = wi (the last inequality in (C.2)) 2. wk = √ nhk,i,∀i ∈ [n] (the third inequality in (C.2))
3. ‖wk‖ = ‖wk′‖, ‖hk,i‖ = ‖hk′,j‖,∀k, k′ ∈ [K], i, j ∈ [n] (the second inequality in (C.2))
4. sk,i,j = w>k hk,i − w>j hk,i = KK−1w > k hk,i,∀k, j ∈ [K], i ∈ [n] (the first inequality in
(C.1))
These four conditions are exactly equivalent to neural collapse conditions and ‖W ‖F = ‖H‖F
The global optimality is not enough to illustrate how does gradient flow 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, define the cross-entropy loss on a matrix Z ∈ RK×nK :
L(Z) = K∑ k=1 n∑ i=1 − log e zk,i,j∑K l=1 e zk,i,l ,
where zk,i,j denote the j-th row and [(k − 1)K + i]-th column elements of Z. Then we have L(W ,H) = L(WH) . Now compute the gradient of L(Z) to each element:
∂L(Z) ∂zk,i,k = −1 + zk,i,k∑K
l=1 e zk,i,l
∂L(Z) ∂zk,i,j = zk,i,j∑K l=1 e zk,i,l ,∀k 6= j.
If u ∈ RK satisfies u>∇L(Z) = 0, denote up as the maximum element of u, then we have:
0 = up ∂L(Z) ∂zp,i,p + ∑ q 6=p uq ∂L(Z) ∂zp,i,q =up(−1 + zp,i,p∑K l=1 e zp,i,l ) + ∑ q 6=p uq zp,i,q∑K l=1 e zp,i,l
=− ∑ q 6=p (up − uq) zp,i,q∑K l=1 e zp,i,l ≤ 0.
(C.3)
Where the last inequality holds if and only if uq = up,∀q ∈ [K]. Which indicates that the rank of ∇L(Z) is K − 1 and u>∇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:
∇WL(W ,H) = ∇L(WH)H>,∇HL(W ,H) = W>∇L(WH).
If there exist a vector (∆W ,∆H) ∈ T (W ,H) such that 〈∇WL(W ,H),∆W 〉 + 〈∇HL(W ,H),∆H〉 6= 0, moreover we can assume 〈∇WL(W ,H),∆W 〉 + 〈∇HL(W ,H),∆H〉 < 0 since we can take the negative direction if the formula is greater than zero, then by Taylor expansion:
L(W + δ∆W ,H + δ∆H) = L(W ,H) + δ〈∇WL(W ,H),∆W 〉+ δ〈∇HL(W ,H),∆H〉+O(δ2),
we know that (∆W ,∆H) satisfies our requirement.
Now let’s discuss the case when:
〈∇WL(W ,H),∆W 〉+ 〈∇HL(W ,H),∆H〉 = 0,∀(∆W ,∆H) ∈ T (W ,H),
by definition of T (W ,H), it contains all vectors that are orthogonal to (W ,H), so (∇WL(W ,H),∇HL(W ,H)) is parallel to (W ,H), that is, there exist λ such that:
∇L(WH)H> = λW ,W>∇L(WH) = λH. (C.4)
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 ′,H ′) : ‖W ′‖2F + ‖H ′‖2F = ‖W ‖2F + ‖H‖2F }, t = 0 admits the local minimum of L(φ(t)) and thus:
0 ≤ d 2
dt2 L(φ(t))|t=0 = φ′(0)T∇2L(W ,H)φ′(0) +∇L(W ,H)φ′′(0). (C.5)
On the other hand, since the curve lies on the sphere S , denote h(W ,H) = ‖W ‖2F + ‖H‖2F , then h(φ(t)) must stay as a constant, take twice derivative we have:
0 = d2
dt2 h(φ(t))|t=0 = φ′(0)T∇2h(W ,H)φ′(0) +∇h(W ,H)φ′′(0). (C.6)
Then sum these two conditions together, we have:
0 ≤ d 2 dt2 (L(φ(t)) + |λ| 2 h(φ(t)))|t=0 = φ′(0)T∇2(L+ |λ| 2 h)(W ,H)φ′(0) +∇(L+ |λ| 2 h)(W ,H)φ′′(0).
(C.7)
By equation (C.4) we know that ∇(L + |λ|2 h)(W ,H) = 0 Note that φ ′(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 φ′(0) = (∆W ,∆H). Then (C.7) indicates that ∀(∆W ,∆H) ∈ T (W ,H) we have:
0 ≤(∆W ,∆H)>∇2L(W ,H)(∆W ,∆H) + |λ| 2 (∆W ,∆H)>∇2h(W ,H)(∆W ,∆H)
=(∆W ,∆H)>∇2L(W ,H)(∆W ,∆H) + |λ|(‖∆W ‖2F + ‖∆H‖2F ). (C.8)
When λ = 0, by equation (C.3) we know that ‖∇L(WH)‖2 > 0, which gives us |λ| ≤ ‖∇L(WH)‖2 When λ 6= 0, combine the two equations in (C.4) we know:
λW>W = W>∇L(WH)H> = λHH> ⇒W>W = HH>, (C.9)
which further implies: ||W ||F = ||H||F , ||W ||2 = ||H||2.
Thus we also have (Note that when W = H = 0 we can take λ to be zero):
∇L(WH)H> = λW ⇒ |λ|||W ||2 ≤ ||∇L(WH)||2||H||2 ⇒ |λ| ≤ ||∇L(WH)||2.
Now when |λ| < ||∇L(WH)||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>a = 0,
let u and v are the left and right singular vectors corresponding to the largest singular value of ∇L(WH), construct ∆W = ua>,∆H = −av>, then (∆W ,∆H) ∈ T (W ,H) and:
(∆W ,∆H)>∇2L(W ,H)(∆W ,∆H)− λ(‖∆W ‖2F + ‖∆H‖2F ) =(W∆H + ∆WH)∇2L (WH) (W∆H + ∆WH) + 2 〈∇L(WH),∆W∆H〉+ |λ|(‖∆W ‖2F + ‖∆H‖2F ) ≤2||a||22(|λ| − u>∇L(WH)v) < 0.
Then it only remains to analyze the |λ| = ||∇L(WH)||2 cases, construct another convex optimization problem:
min Z L(Z) + |λ|||Z||∗, (C.10)
suppose Z has SVD Z = UΣV >, as we know that the subgradient of ||Z||∗ can be written as (see Watson (1992) for a proof):
∂‖Z‖∗ = { UV > + W ,W ∈ RK×nK | U>W = 0,WV = 0, ‖W ‖2 ≤ 1 } . (C.11)
On the other hand, we know that:
H>HH>H = H>W>WH = V Σ2V >
WW>WW> = WHH>W> = UΣ2U>,
which indicates that H>H = V ΣV > and WW> = UΣU>. Combine them with (C.4) we have:
∇L(WH)H>H = λWH ⇔ ∇L(WH)V ΣV > = λUΣV >
⇔ ∇L(WH)V = λU WW>∇L(WH) = λWH ⇔ UΣU>∇L(WH) = λUΣV >
⇔ U>∇L(WH) = λV >.
Note that |λ| = ||∇L(WH)||2, then by (C.11) we know that −∇L(WH) ∈ |λ|∂‖WH‖∗. Then by the strict convexity of (C.10) we know WH is the global minimum of it. In addition, we have ‖W ‖2F +‖H‖2F = 2trΣ2) = 2‖WH‖∗. In addition, previous works (Haeffele & Vidal, 2015) have shown that:
‖Z‖∗ = min Z=WH
1
2
( ‖W ‖2F + ‖H‖2F ) ,
which is equivalent to:
‖WH‖∗ ≤ 1
2 (‖W ‖2F + ‖H‖2F ).
Now for any (W ′,H ′), we know that:
L(W ,H) + |λ| 2 (‖W ‖2F + ‖H‖2F ) = L(WH) + |λ|‖WH‖∗ ≤ L(W ′H ′) + |λ|‖W ′H ′‖∗
≤ L(W ′,H ′) + |λ| 2 (‖W ′‖2F + ‖H ′‖2F ),
which indicates (W ,H) must attain global minimum of the following optimization problem:
min W ,H L(W ,H) + |λ| 2 (‖W ‖2F + ‖H‖2F ).
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 ′,H ′) such that L(W ′,H ′) < L(W ,H) and ‖W ‖2F + ‖H‖2F = ‖W ′‖2F + ‖H ′‖2F thus:
L(W ,H) + |λ| 2 (‖W ‖2F + ‖H‖2F ) > L(W ′,H ′) + |λ| 2 (‖W ′‖2F + ‖H ′‖2F ),
which contradicts with the global optimality of (W ,H), thus (W ,H) must satisfy all of the neural collapse conditions and we finish the proof.

heading: D ADDITIONAL EMPIRICAL RESULTS
text: 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 105 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(‖h̄k − h̄‖)/Avg(‖h̄k − h̄‖)) and the variation of the classifier’s norms (i.e., Std(‖w̄k‖)/Avg(‖w̄k‖)). (2) Within-class variation of the last layer features (i.e., Avg(‖hk,i − hk‖)/Avg(‖hk,i − h̄‖)). (3) The cosines between pairs of last layer features (i.e., 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)|)). (4) The distance between the normalized centered classifier and the normalized last layer feature (i.e., Avg(|(h̄k − h̄)/‖h̄k − h̄‖ − w̄k/‖w̄k‖|)). Empirically, we observe that all four quantities decrease at approximately the rate O(1/(log(t))).
Details of Realistic Training. In the real data experiments, we trained the VGG-13 (Simonyan & Zisserman, 2014) and ResNet18 (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 defined similarly to the ULPM case and the experiment results are reported in Figure 1, 3, 4, 5, 6, 7, and 8. All experiments were run in Python (version 3.6.9) on Google Colab. It
should be mentioned that we can observe that the variation of classifier norm stays at a low value and do not decrease in some settings (e.g., Figure 7a and Figure 8a), 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 find that neural collapse occurs for various network architectures and datasets under an unconstrained setting, which provides sound support for our theory.

heading: E LIMITATION AND FUTURE DIRECTIONS
text: 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 flow 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.

Review summaries:
This paper analyzes the phenomenon of neural collapse from the simplified perspective of an unconstrained features models, whereby the only two optimization variables in the model are the last layer classifier and features, which are fit to some labels by minimizing the cross entropy loss. The authors leverage their observation that maximizing the margin results in a model achieving neural collapse, and then study the corresponding problem of minimum-norm classifier, and show that gradient flow converges directionally to the neural collapse solution. The authors then show that the optimization landscape is benign, in that all critical points that are not global minimum have a decreasing direction in the tangent space. 

This paper provides new analyses for understanding the neural collapse phenomenon observed in training deep learning classifiers. In particular, to understand the last-layer features and classifiers, the authors study the gradient flow of the unconstrained layer-peeled model (ULPM) without any regularizers and show that gradient flow converges to neural collapse classifiers and features. Experiments are provided to demonstrate the neural collapse phenomenon when no weight decay is used.

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