Linear Convergence Analysis of Neural Collapse with Unconstrained Features
Abstract
In this work, we study the recently discovered neural collapse (NC) phenomenon, which is prevalent in training over-parameterized deep neural networks for classification tasks.
Existing work has shown that any optimal solution of the trained problem for classification tasks is an NC solution and has a benign landscape under the unconstrained feature model.
However, these results do not provide an answer to the question of how quickly gradient descent can find an NC solution.
To answer this question, we prove an error bound property of the trained loss, which refers to the inequality that bounds the distance of a point in the optimal solution set by the norm of its gradient, under the unconstrained feature model.
Using this error bound, we show linear convergence of gradient descent for finding an NC solution.
1. Introduction
Over the past years, a large amount of work is devoted to attempting to understand the underlying mechanism of deep neural networks (DNNs) from a theoretical point of view; see, e g., [1, 6, 12]. Towards the goal of understanding the representations learned by deep neural networks, a recent line of seminal works [3, 14] presents an intriguing phenomenon named neural collapse (NC) that occurs pervasively across a range of canonical classification problems during the terminal phase of training.
Specifically, the authors observed that the last-layer features and the last-layer linear classifiers of a trained DNN exhibit the following simple but elegant structure: (i) Variability Collapse: the individual features of each class concentrate to their class-means; (ii) Convergence to Simplex ETF: the class-means have the same length and are maximally distant.
In other words, they form a Simplex Equiangular Tight Frame (ETF); (iii) Convergence to Self-Duality: the last-layer linear classifiers perfectly match their class-means up to rescaling.
Study of NC under unconstrained feature model.
Recently, a line of works demystifies the NC phenomenon in theory based on the so-called unconstrained feature model [13, 23], which is also called layer-peeled model [3]; see, e.g., [3, 4, 14, 15, 18, 20, 21, 23] and the references therein.
Specifically, consider an L -layer fully connected neural network of the form
[Equation]
where σ ( · ) is an activation function, Θ
= { W k , b k } Lk =1 denotes all the network parameters, and θ = { W k , b k } L − 1 k =1 denotes the network parameters up to the last layer.
In particular, the output of the penultimate layer, denoted by ϕ θ ( x ) , is referred to as the feature of the sample x learned by the neural network.
Given training samples { ( x k,i , y k ) } ⊆
R d × R K drawn from the same data distribution, they studied the multi-class (e.g., K classes) classification problem by minimizing the empirical risk over these samples,
[Equation]
where N = nK is the total number of samples, y k ∈ R K is an one-hot vector with only the k -th entry being 1 and the remaining ones being 0 for all k ∈
[ K ] , x k,i ∈ R
d is the i -th sample in the k -th class, n denotes the number of training samples in each class, λ > 0 is the regularization parameter, and L : R K × R K
→ R is a loss function.
Note that the unconstrained feature model assumes that the last-layer features of the neural network are free optimization variables, which simplifies the non-linear interaction across layers.
Intuitively, this simplification is reasonable because modern deep neural networks are often highly over-parameterized such that last-layer features can approximate or interpolate any point in the feature space [5, 10].
We introduce an auxiliary variable h k,i =
ϕ θ ( x k,i ) , which denotes the last-layer feature corresponding to the sample x k,i .
Using the unconstrained feature model and letting W = W TL and b =
b L , we consider a variant of Problem (1) by treating h k,i for all k, i as free optimization variables, i.e.,
[Equation]
Recently, it has been shown that any global optimal solution of Problem (2) and its constrained counterparts satisfies the NC properties when the function L is the CE loss, MSE loss, or supervised contrastive loss function; see, e.g., [3, 4, 14, 15, 18, 20, 23].
In particular, Yaras et al.
[18], Zhou et al.
[20], Zhu et al.
[23] showed that their studied training loss function has a benign landscape with no spurious local minimum.
This implies that the NC solutions can be reached by gradient descent methods with random initialization [2, 7].
However, these results still cannot answer the question that how fast gradient descent can reach an NC solution.
Linear convergence analysis under the error bound.
A powerful approach to analyzing the convergence behavior of gradient descent type methods for optimizing convex optimization problems is to use an error bound property; see, e.g., [11, 16, 22] and the references therein. Recently, Liu et al. [8, 9], Wang et al.
[17], Zheng et al.
[19] extended this approach to study non-convex optimization problems.
Consider an optimization problem
[Equation]
where E is a finite-dimensional Euclidean space and F : E → ( −∞ , + ∞ ) is a continuously differentiable function.
Let X ⊆
E denote the set of optimal solutions of Problem (3).
Then, we say that an error bound holds for Problem (3) if there exist constants δ, κ > 0 such that for all x ∈ E with dist ( x , X ) ≤ δ ,
[Equation]
Suppose that the error bound holds for Problem (3).
A unified approach to analyzing convergence rate of first-order iterative methods have been established based on the error bound; see, e.g., [11, 16, 22].
Fact 1 Suppose that the optimal solution of Problem (3) is non-empty, i.e., X ̸ = ∅ , and the error bound (4) holds for Problem (3) . Suppose in addition that the sequence { x k } k ≥ k 1 for an index k 1 ≥ 0 satisfies the following conditions: (i). (Sufficient Decrease)
There exists a constant κ 1 > 0 such that
[Equation]
(ii). (Cost-to-Go Estimate) There exists a constant κ 2 > 0 such that [Equation]
(iii).
(Safeguard)
There exists a constant κ 3 > 0 such that [Equation] Then, the sequence { F ( x k ) } k ≥ 0 converges Q-linearly to v ∗ and { x k } k ≥ 0 converges R-linearly to some x ∗ ∈ X .
Contributions of this work.
Motivated by above discussions, we are devoted to establishing an error bound for Problem (2) and proving linear convergence of gradient descent for solving Problem (2) in this work.
Specifically, we study Problem (2) with the MSE loss and CE loss, respectively.
We first characterize the optimal solution set of Problem (2) and then prove the error bound using the structure of optimal solution set.
Our experimental results on synthetic datasets complement and support our theoretical developments.
2. Main Results
Before we proceed, we introduce some essential notation.
Given a point y ∈ R n and a nonempty and closed set X ⊆ R n , we denote the Euclidean distance of y to X by dist ( y , X ) =
min {∥ x − y ∥ 2 : x ∈ X } .
We use 1 n to denote the n -dimensional all-one vector.
Let P := I − 1 K 1 TK /K and P ⊥ :
= 1 K 1 TK /K .
Given a positive integer n , we denote by [ n ] the set { 1 , . . . , n } .
We write all the features and the classifiers in the matrix form
[Equation]
For ease of exposition, let X denote the optimal solution set of Problem (2). {+Without loss of generality, we assume the label matrix to be Y = 
 Tn ⊗ I K .+}
2.1. Mean Squared Error Loss
Suppose that the loss function L is the MSE loss of the form L ( z , y k )
=
12 ∥
z − y k ∥ 22 .
Substituting this into Problem (2) yields
[Equation]
Then, we characterize the optimal solution set of Problem (5) as follows.
Proposition 1 For the matrix P , let P =
V Σ V T be its eigenvalue decomposition such that
V =
[ ¯ V v ] ∈
O K with ¯ V ¯ V T =
P , vv T =
P ⊥ , and Σ = (cid:20)
I K −   (cid:21) .
The optimal solution set
X of Problem (5) can be characterized as follows:
(i)
If λ W λ H ≥ 1 / ( nK 2 ) , we have X = (cid:8)(cid:0) 0 , 0 , 1 K 1 K (cid:1)(cid:9) .
(ii)
If λ W λ H < 1 / ( nK 2 ) , we have X = (cid:26)(cid:18)
[Equation]
where γ := 1 − K √ nλ W λ H
This result indicates that when λ W λ H < 1 /K 2 , any global optimal solution of Problem (5) satisfies the NC properties.
Compared to [20, Theorem 3.1], we provide all the parameters that depend on λ W and λ H explicitly.
Based on the above result, we prove the error bound of Problem (5).
Theorem 2 For Problem (5) , if λ W λ H ̸ = 1 /K 2 , there exist constants δ 1 , κ 1 > 0 that depend on λ W and λ H such that for any ( W , H , b ) satisfying dist (( W , H , b ) , X ) ≤ δ 1 , it holds that dist
[Equation]
Before we proceed, some remarks are in order.
First, when λ W λ H ̸ = 1 /K 2 , we show that the error bound holds for Problem (5) for any point in the neighborhood of the optimal solution set.
Second, using this error bound and Fact 1, we can prove the linear convergence of gradient descent for solving Problem (5).
2.2. Cross-Entropy Loss
Suppose that the loss function L is the CE loss of the form
[Equation]
As done in [3, 4], we consider Problem (2) without the bias term when we use the CE loss.
Substituting (7) into Problem (2) yields
[Equation]
Then, we characterize the optimal solution set of Problem (8) as follows.
Proposition 3 Suppose that the dimension of features is no smaller than the number of classes, i.e., d ≥ K .
Then, the global optimal solution set X of Problem (8) can be characterized as follows:
(i)
If λ W λ H ≥ 1 / ( nK 2 ) , it holds that X
= { ( 0 , 0 ) } .
(ii)
If λ W λ H < 1 / ( nK 2 ) , it holds that
[Equation]
where γ := 1 / √ nλ W
[Equation] H − K + 1 . Based on this result, we prove the error bound for Problem (8).
Theorem 4 Suppose that the number of classes K = 2 and the dimension of features is no smaller than the number of classes, i.e., d ≥ K .
For Problem (8) , if λ W λ H ̸ = 1 /K 2 , there exist constants δ 2 , κ 2 > 0 that depend on λ W and λ H such that for any ( W , H ) satisfying dist (( W , H ) , X ) ≤ δ 2 , it holds that
[Equation]
Remark that since the structure of the CE loss is more complicated than that of the MSE loss, we can only establish the error bound for Problem (8) when K = 2 .
We leave the error bound when
K ≥ 2 for future work.
3. Experimental Results
In this section, we corroborate our theory with experimental results by solving Problem (2) and by solving Problem (1) with ψ Θ being a two-layer neural network on synthetic data.
First, we employ batch gradient descent using backtracking line search for optimizing Problem (2) with both MSE and CE losses and different weight and feature regularization parameters λ W , λ H , respectively.
In the tests, we set the number of classes as K = 10 , number of samples in each class as n = 10 , and the dimension of features as d = 20 . Then, we report the convergence performance of gradient descent in Figure 1 with different λ W and λ H . One can observe that the optimality gap measured by
F k − F ∗ , where F k is the function value at the k -th iteration and F ∗ is the optimal function value computed by the NC solutions, converges linearly independent of the regularization parameters λ W , λ H . We refer to the readers to Appendix for the convergence rate of the NC metrics.
Next, we consider training a two-layer neural network on synthetic data by solving Problem (1) using gradient descent with constant step size.
We generate the training samples as follows.
The samples x k,i are drawn i.i.d.
from a zero-mean D -dimensional Gaussian distribution with
D = 1000 and covariance I / 4 , and the labels are drawn uniformly from { 1 , . . .
, K } (so that we have balanced classes).
In the tests, we set the number of classes as K = 10 , number of samples in each class as n = 100 , width of the first layer as m = 256 , and dimension of features as d = 20 .
Then, we report the results in Figure 2.
Unlike the UFM, we found that the regularization parameter is important in the convergence rate and requires careful tuning - for particular settings of λ W , λ H , we can achieve linear convergence when training a shallow network with explicit regularization. We refer to the readers to Appendix for the convergence rate of the NC metrics.
4. Conclusions
In this work, we studied the NC phenomenon under the unconstrained feature model.
We first characterized the set of optimal solutions of the trained problem. Based on this characterization, we showed that error bound holds for the trained problem.
Using this, we further established the linear convergence of gradient descent for optimizing the trained problem.
Finally, we supported our theoretical results by experimental results.
As a future work, we would like to extend our analysis to the case of K ≥ 2 for the CE loss.
Another interesting direction is to analyze the error bound of the constrained counterpart of Problem (2); see, e.g., [3, 18]. Appendix Linear Convergence of the NC Metrics
In Section 3, we demonstrated that gradient descent achieves linear convergence in terms of function value for training Problem (1) under the UFM and Problem (2) with shallow networks, respectively.
A natural question is that whether we can empirically observe linear convergence of metrics that measure the various aspects of neural collapse.
For measuring different aspects of neural collapse as introduced in Section 1, we adopt similar NC metrics as that in [14, 20, 23],
[Equation]
where Σ W and Σ B are the within-class and between-class covariance matrices (see [14, 23] for more details), Σ † B denotes pseudo inverse of Σ B , and H is the centered class mean matrix.
More specifically, N C 1 measures within class variability collapse, N C 2 measures convergence to the simplex ETF, and N C 3 measures duality collapse.
We first investigate the convergence rate of the NC metrics under the UFM, using the same set-up as that in Section 3.
We report the result in Figure 3.
One can observe that all three metrics converge linearly to 0 for both MSE and CE losses.
In particular, N C 1 converges substantially faster than the other two metrics.
, λ H = 5 × 10 − 5 .
Analogously, we investigate the convergence rate of the NC metrics for a two-layer neural network using the same set-up and data as that in Section 3.
We report the result in Figure 4.
We again observe that all three metrics converge linearly to 0 for both MSE and CE losses using appropriate regularization.
Motivated by the above observations an interesting direction is to prove linear convergence of these NC metrics using the error bound condition.
This is left as a future work