SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Peter Sukeník  
Institute of Science and Technology  
Klosterneuburg, Austria  
peter.sukenik@ista.ac.at
Marco Mondelli*  
Institute of Science and Technology  
Klosterneuburg, Austria  
marco.mondelli@ista.ac.at
Christoph H. Lampert*  
Institute of Science and Technology  
Klosterneuburg, Austria  
chl@ist.ac.at
Neural collapse (NC) refers to the surprising structure of the last layer of deep neural networks in the terminal phase of gradient descent training. Recently, an increasing amount of experimental evidence has pointed to the propagation of NC to earlier layers of neural networks. However, while the NC in the last layer is well studied theoretically, much less is known about its multi-layered counterpart - deep neural collapse (DNC). In particular, existing work focuses either on linear layers or only on the last two layers at the price of an extra assumption. Our paper fills this gap by generalizing the established analytical framework for NC - the unconstrained features model - to multiple non-linear layers. Our key technical contribution is to show that, in a deep unconstrained features model, the unique global optimum for binary classification exhibits all the properties typical of DNC. This explains the existing experimental evidence of DNC. We also empirically show that  $(i)$  by optimizing deep unconstrained features models via gradient descent, the resulting solution agrees well with our theory, and  $(ii)$  trained networks recover the unconstrained features suitable for the occurrence of DNC, thus supporting the validity of this modeling principle.
In the thought-provoking paper [22], Papyan et al. uncovered a remarkable structure in the last layer of sufficiently expressive and well-trained deep neural networks (DNNs). This phenomenon is dubbed "neural collapse" (NC), and it has been linked to robustness and generalization of DNNs. In particular, in [22] it is shown that all training samples from a single class have the same feature vectors after the penultimate layer (a property called NC1); the globally centered class-means form a simplex equiangular tight frame, the most "spread out" configuration of equally-sized vectors geometrically possible (NC2); the last layer's classifier vectors computing scores for each of the classes align with the centered class-means (NC3); finally, these three properties make the DNN's last layer act as a nearest class-center classifier on the training data (NC4).
The intriguing phenomenon of neural collapse has spurred a flurry of interest aimed at understanding its emergence, and the unconstrained features model (UFM) [21] has emerged as a widely accepted theoretical framework for its analysis. In this model, one assumes that the network possesses such an
Figure 1: Deep neural collapse metrics as a function of training progression. On the left, DNC1 measures within-class variability collapse (the closer to 0, the more collapsed); in the middle, DNC2 measures closeness of feature activations to an orthogonal frame (the closer to 1, the more collapsed); and on the right, DNC3 measures alignment between feature activations and weights (the closer to 0, the more collapsed). The ResNet20 recovers the DNC metrics for DUFM, in accordance with our theory. This validates unconstrained features as a modeling principle for neural collapse.
expressive power that it can represent arbitrary features on the penultimate layer's output. Thus, the feature representations of the training samples on the penultimate layer are treated as free optimization variables, which are optimized for (instead of the previous layers' weights). Since its introduction, the UFM has been understood rather exhaustively, see e.g. [21, 19, 31, 3, 12, 42] and Section 2.
Most recently, researchers started to wonder, whether NC affects solely the last layer. The "cascading neural collapse" is mentioned in [11], where it is pointed out that NC1 might propagate to the layers before the last one. Subsequently, a mathematical model to study this deep version of NC has been proposed in [28], albeit only for two layers; and a deep model has been considered in [2], albeit only for linear activations. On the empirical side, the propagation of NC1 through the layers is demonstrated in [8], and the occurrence of all the NC properties in multiple layers is extensively measured in [26]. All these works agree that neural collapse occurs to some extent on the penultimate and earlier layers, leaving as a key open question the rigorous characterization of its emergence.
To address this question, we generalize the UFM and its extensions in [28, 2] to cover an arbitrary number of non-linear layers. We call this generalization the deep unconstrained features model (DUFM). The features, as in the original UFM, are still treated as free optimization variables, but they are now followed by an arbitrary number of fully-connected layers with ReLU activations between them, which leads to a natural setup to investigate the theoretical underpinnings of deep neural collapse (DNC). Our main contribution is to prove (in Section 4) that DNC is the only globally optimal solution in the DUFM, for a binary classification problem with  $l_{2}$  loss and regularization. This extends the results in [28] and [2] by having more than two layers and non-linearities, respectively. While the statements of [28] are formulated for an arbitrary number of classes, here we restrict ourselves to binary classification. However, we drop an important assumption on the features of the optimal solution made by [28] (see the corresponding Theorem 4.2 and footnote 3), which was only verified experimentally. In a nutshell, we provide the first theoretical validation of the occurrence of DNC in practice: deep neural collapse gives a globally optimal representation of the data under the unconstrained features assumption, for any arbitrary depth.
In addition to our theoretical contribution, we numerically validate in Section 4.2 the occurrence of DNC when training the DUFM with gradient descent. The loss achieved by the DUFM after training agrees well with the theoretically computed optimum, and the corresponding solution clearly exhibits the properties described by our main result. Finally, to check whether the main assumption of the DUFM – the unconstrained features – is relevant in practical training, we show that DNNs trained on CIFAR-10 can indeed represent features suitable for the occurrence of DNC. In particular,  $(i)$  we pre-train a small DUFM to full collapse,  $(ii)$  attach the fully-connected layers (DUFM without the unconstrained features) to the end of a randomly initialized CNN, and finally  $(iii)$  train the CNN with the DUFM's weights fixed. The plots in Figure 1 show that the backbone recovers the unconstrained features leading to DNC. This is quite remarkable, as the collapse, being data-dependent, completely breaks after exchanging the trained unconstrained features with the output of the randomly initialized backbone. Additional details can be found in Section 4.2.
Since the original paper [22], neural collapse has been intensively studied, with [21, 3] introducing the now widely accepted unconstrained features model (UFM), also called layer peeled model. In particular, the work on NC can be loosely categorized into  $(i)$  the study of the mathematical grounds behind NC [31, 19, 40, 3, 42, 12, 21, 7, 18, 29, 36, 37],  $(ii)$  the development of a connection between NC and DNN performance [32, 5, 6, 1, 34, 15, 39, 17, 16] and  $(iii)$  empirical observations concerning NC [35, 20]. The first line of work can be further divided into  $(i - a)$  a static analysis, that studies the global optimality of NC on suitable objectives [31, 19, 40, 28, 42, 12], and  $(i - b)$  a dynamic analysis, that characterizes the emergence of NC in training with gradient-based methods [21, 7, 12, 30]. The vast majority of works uses the UFM to derive their results, but there are attempts to derive NC without it, see e.g. [25, 33, 23, 24, 14, 38]. For additional references, see also the recent survey [13].
More specifically, [31, 19] show global optimality of NC under UFM, when using the cross-entropy (CE) loss. Similar optimality results are obtained in [40, 28] for the MSE loss. For CE loss, a class-imbalanced setting is considered in [3], where the danger of minority classes collapsing to a single representation (thus being indistinguishable by the network) is pointed out. This analysis is further refined in [27], where a generalization of the ETF geometry under class imbalance is discovered. On the edge between static and dynamic analysis are results which characterize the optimization landscape of UFM. They aim to show that all the stationary points are either global optima (thus exhibiting NC) or local maxima/strict saddle points with indefinite Hessian, which enables optimization methods capable of escaping strict saddles to reach the global optimum. For the CE loss, this property is proved in [42] and [12], while in the MSE setup this insight is provided by [40]. A unified analysis for CE, MSE and some other losses is done in [41]. The convergence of the GD dynamics to NC is considered in [21] under MSE loss for near-zero initialization. The more refined analysis of [7] reveals that the last layer's weights of UFM are close to being conditionally optimal given the features and, assuming their exact optimality, the emergence of NC is demonstrated with renormalized gradient flow. For CE loss, [12] shows that gradient flow converges in direction to the KKT point of the max-margin optimization problem, even without any explicit regularization (unlike in the MSE loss setting). As the optimal solutions of this max-margin problem exhibit NC, this yields the convergence of gradient-based methods to NC as well. Quantitative results on the speed of convergence for both the MSE and CE losses are given in [30].
Most relevant to our paper, the emergence of NC on earlier layers of DNNs is empirically described in [8, 26, 11, 4]. In particular, it is demonstrated in [8] that, at the end of training, the NC1 metric decreases log-linearly with depth, thus showing a clear propagation of neural collapse. Most recently, the propagation of all the NC properties beyond the last layer is studied in [26], for a setting that includes biases. On the theoretical side, the UFM with multiple linear layers is considered in [2], where it is shown that NC is globally optimal for different formulations of the problem. The first result for a non-linear model with two layers (separated by a ReLU activation) is given by [28]. There, it is shown that NC, as formulated in the bias-free setting, is globally optimal. However, an additional assumption on the optimum is required, see the detailed discussion in Section 4.
Notation. Let  $N = Kn$  be the total number of training samples, where  $K$  is the number of classes and  $n$  the number of samples per class. We are interested in the last  $L$  layers of a DNN. The layers are counted from the beginning of indexing, i.e., the  $L$ -th layer denotes the last one. The layers are fully connected and bias-free with numbers of input neurons  $d_{1}, \ldots, d_{L} > 1$  and weight matrices  $W_{1} \in \mathbb{R}^{d_{2} \times d_{1}}$ ,  $W_{2} \in \mathbb{R}^{d_{3} \times d_{2}}$ , ...,  $W_{L} \in \mathbb{R}^{K \times d_{l}}$ . Let  $H_{1} \in \mathbb{R}^{d_{1} \times N}$ ,  $H_{2} \in \mathbb{R}^{d_{2} \times N}$ , ...,  $H_{L} \in \mathbb{R}^{d_{L} \times N}$  be the feature matrices before the application of the ReLU activation function, which is denoted by  $\sigma$ , i.e.,  $H_{2} = W_{1}H_{1}$  and  $H_{l} = W_{l-1}\sigma(H_{l-1})$  for  $l \geq 3$ . When indexing a particular sample, we denote by  $h_{c,i}^{(l)}$  the  $i$ -th sample of the  $c$ -th class in the  $l$ -th layer, i.e., a particular column of  $H_{l}$ . Moreover, let  $\mu_{c}^{(l)} = \frac{1}{n}\sum_{i=1}^{n}h_{c,i}^{(l)}$  be the activation class-mean for layer  $l$  and class  $c$ , and  $\bar{H}_{l}$  the matrix containing  $\mu_{c}^{(l)}$  as the  $c$ -th column. For convenience, the training samples are arranged class-by-class (first all the samples of first class, then of the second, etc.). Thus, the label matrix containing the one-hot encodings  $Y \in \mathbb{R}^{K \times N}$  can be written as  $I_{K} \otimes \mathbf{1}_{n}^{T}$ , where  $I_{K}$  denotes a  $K \times K$  identity matrix,  $\otimes$
denotes the Kronecker product and  $\mathbf{1}_n$  a row vector of all-ones of length  $n$ . With this, we can write  $\bar{H}_l = \frac{1}{n} H_l(I_K \otimes \mathbf{1}_n)$  and we further denote  $\overline{\sigma(H_l)} := \frac{1}{n}\sigma(H_l)(I_K \otimes \mathbf{1}_n)$ .
Deep neural collapse. Since we assume the last  $L$  layers to be bias-free, we formulate a version of the deep neural collapse (DNC) without biases (for a formulation including biases, see [26]).
Definition 1. The deep neural collapse at layer  $l$  is described by the following three properties:
DNC1: The within-class variability after  $l - 1$  layers is 0. This property can be stated for the features either before or after the application of the activation function  $\sigma$ . In the former case, the condition requires  $h_{c,i}^{(l)} = h_{c,j}^{(l)}$  for all  $i, j \in [n]$  (or, in matrix notation,  $H_l = \bar{H}_l \otimes \mathbf{1}_n^T$ ); in the latter,  $\sigma(h_{c,i}^{(l)}) = \sigma(h_{c,j}^{(l)})$  for all  $i, j \in [n]$  (or, in matrix notation,  $\sigma(H_l) = \overline{\sigma(H_l)} \otimes \mathbf{1}_n^T$ ).
DNC2: The feature representations of the class-means after  $l - 1$  layers form an orthogonal matrix. As for DNC1, this property can be stated for features either before or after  $\sigma$ . In the former case, the condition requires  $\bar{H}_l^T\bar{H}_l\propto I_K$ ; in the latter,  $\overline{\sigma(H_l)^T}\overline{\sigma(H_l)}\propto I_K$ .
DNC3: The rows of the weight matrix  $W_{l}$  are either 0 or collinear with one of the columns of the class-means matrix  $\overline{\sigma(H_l)}$ .
Note that, in practice, DNNs do not achieve such properties exactly, but they approach them as the training progresses.
Deep unconstrained features model. Next, we generalize the standard unconstrained features model and its extensions [28, 2] to its deep, non-linear counterpart.
Definition 2. The  $L$ -layer deep unconstrained features model (L-DUFM) denotes the following optimization problem:
$$
\min  _ {H _ {1}, W _ {1}, \dots , W _ {L}} \frac {1}{2 N} \| W _ {L} \sigma \left(W _ {L - 1} \sigma \left(\dots W _ {2} \sigma \left(W _ {1} H _ {1}\right)\right)\right) - Y \| _ {F} ^ {2} + \sum_ {l = 1} ^ {L} \frac {\lambda_ {W _ {l}}}{2} \| W _ {l} \| _ {F} ^ {2} + \frac {\lambda_ {H _ {1}}}{2} \| H _ {1} \| _ {F} ^ {2}, \tag {1}
$$
where  $\| \cdot \| _F$  denotes the Frobenius norm and  $\lambda_{H_1},\lambda_{W_1},\ldots ,\lambda_{W_L} > 0$  are regularization parameters.
4 Optimality of deep neural collapse
We are now state our main result (with proof sketch in Section 4.1, full proof in Appendix A).
Theorem 3. The optimal solutions of the  $L$ -DUFM (1) for binary classification  $(K = 2)$  are s.t.
$$
\left(H _ {1} ^ {*}, W _ {1} ^ {*}, \dots , W _ {L} ^ {*}\right) \quad e x h i b i t s D N C, \quad \text {i f} n \lambda_ {H _ {1}} \lambda_ {W _ {1}} \dots \lambda_ {W _ {L}} <   \frac {(L - 1) ^ {L - 1}}{2 ^ {L + 1} L ^ {2 L}}. \tag {2}
$$
More precisely, DNC1 is present on all layers; DNC2 is present on  $\sigma(H_{l}^{*})$  for  $l \geq 2$  and on  $H_{l}^{*}$  for  $l \geq 3$ ; and DNC3 is present for  $l \geq 2$ . The optimal  $W_{1}^{*}, H_{1}^{*}$  do not necessarily exhibit DNC2 or DNC3. If the inequality in (2) holds with the opposite strict sign, then the optimal solution is only  $(H_{1}^{*}, W_{1}^{*}, \ldots, W_{L}^{*}) = (0, 0, \ldots, 0)$ .
In words, Theorem 3 shows that, unless the regularization is so strong that the optimal solution is all-zeros, the optimal solutions of the DUFM exhibit all the DNC properties. Notice that (2) requires the per-layer regularization to be at most of order  $L^{-1}n^{1 / (L + 1)}$ , which is a very mild restriction. If the condition in (2) holds with equality, the set of optimal solutions includes (i) solutions exhibiting DNC, (ii) the all-zero solution, and (iii) additional solutions that do not exhibit DNC.
We highlight that Theorem 3 holds for arbitrarily deep non-linear unconstrained features models. This is a significant improvement over both [28], which considers  $L = 2$ , and [2], which considers the linear version of the DUFM. Furthermore, Theorem 4.2 of [28] has a similar statement for two layers
${}^{2}\mathrm{{NC}}$  often involves an additional fourth property - the last layer acting as a nearest class classifier. This property is well defined only for networks including biases, hence we will not consider it here.
(although without considering the DNC1 property on  $H_{1}, H_{2}$ ), but it requires the extra assumption  $\| \sigma(H_{2}^{*}) \|_{*} = \| H_{2}^{*} \|_{*}$ , where  $\| \cdot \|_{*}$  denotes the nuclear norm. In particular, the argument of [28] crucially relies on the fact that  $\| \sigma(H_{2}^{*}) \|_{*} \leq \| H_{2}^{*} \|_{*}$ . We remark that the inequality  $\| \sigma(M) \|_{*} \leq \| M \|_{*}$  does not hold for a generic matrix  $M$  (see Appendix B for a counterexample) and, at the optimum, it is only observed empirically in [28]. In summary, our result is the first to theoretically address the emerging evidence [26, 8, 4] that neural collapse is not only a story of a single layer, but rather affects multiple layers of a DNN.
4.1 Proof sketch for Theorem 3
First, denote by  $S(\cdot)$  an operator which takes any matrix  $A \in \mathbb{R}^{e \times f}$  as an input and returns the set of  $\min\{e, f\}$  singular values of  $A$ . The singular values are denoted by  $s_i$  and ordered non-increasingly. We specifically label  $s_{l,j}$  the  $j$ -th singular value of  $\sigma(H_l)$  for  $l \geq 2$ .
The argument consists of three steps. First, we assume  $n = 1$ , and prove that the optimal solutions of the  $L$ -DUFM problem exhibit DNC2-3 (as stated in Theorem 3). Second, we consider  $n > 1$  and show that DNC1 holds. Third, we conclude that DNC2-3 hold for  $n > 1$ .
Step 1: DNC2-3 for  $n = 1$ . The idea is to split the optimization problem (1) into a series of simpler sub-problems, where in each sub-problem we solve for a fraction of the free variables and condition on all the rest. Crucially, the optimal solutions of all the sub-problems only depend on the singular values of  $\sigma(H_l)$ , for  $l \geq 2$ . Thus, the final objective can be reformulated in terms of  $S(\sigma(H_l))$ , for  $l \geq 2$ . By noticing that the objective is symmetric and separable w.r.t. the first and second singular values of the matrices, we can solve separately for  $\{s_{l,1}\}_{2 \leq l \leq L}$  and  $\{s_{l,2}\}_{2 \leq l \leq L}$ . At this point, we show that these separated problems have a unique global optimum, which is either 0 or element-wise positive (depending on the values of the regularization parameters  $\lambda_{H_1}, \lambda_{W_1}, \ldots, \lambda_{W_L}$ ). Hence, the optimal  $s_{l,j}$  will be equal for  $j = 1$  and  $j = 2$ , which in turn yields that  $\sigma(H_l)$  is orthogonal. Finally, using the intermediate steps of this procedure, we are able to reverse-deduce the remaining properties.
We now sketch the step of splitting the problem into  $L$  sub-problems. As the  $l$ -th sub-problem, we optimize over matrices with index  $L - l + 1$  and condition on all the other matrices, as well as singular values of  $\{\sigma(H_l)\}_{l \geq 2}$  (keeping them fixed). We start by optimizing over  $W_L$  and then optimize over  $W_l$  for  $2 \leq l \leq L - 1$  (for all  $l$  the optimization problem has the same form and thus this part is compressed into a single lemma). If  $L = 2$ , we do not have any  $l$  for which  $2 \leq l \leq L - 1$  and, hence, we skip this part. Finally, we optimize jointly over  $W_1, H_1$ . We start by formulating the sub-problem for  $W_L$ .
Lemma 4. Fix  $H_{1},W_{1},\ldots ,W_{L - 1}$  . Then, the optimal solution of the optimization problem
$$
\min  _ {W _ {L}} \frac {1}{4} \| W _ {L} \sigma (W _ {L - 1} (\dots W _ {2} \sigma (W _ {1} H _ {1})) - I _ {2} \| _ {F} ^ {2} + \frac {\lambda_ {W _ {L}}}{2} \| W _ {L} \| _ {F} ^ {2}
$$
is achieved at  $W_{L} = \sigma (H_{L})^{T}(\sigma (H_{L})\sigma (H_{L})^{T} + 2\lambda_{W_{L}}I_{d_{L}})^{-1}$  and equals
$$
\frac {\lambda_ {W _ {L}}}{2 \left(s _ {L , 1} ^ {2} + 2 \lambda_ {W _ {L}}\right)} + \frac {\lambda_ {W _ {L}}}{2 \left(s _ {L , 2} ^ {2} + 2 \lambda_ {W _ {L}}\right)}. \tag {3}
$$
The proof follows from the convexity of the problem in  $W_{L}$  and a simple form of the derivative, see the corresponding lemma in Appendix A for details.
We remark that the solution (3) resulting from the optimization over  $W_{L}$  depends on  $\sigma(H_{L})$  only through the singular values  $s_{L,1}, s_{L,2}$ . Thus, we can fix  $s_{L,1}, s_{L,2}$  and  $H_{1}, W_{1}, \ldots, W_{L-2}$ , and optimize over  $W_{L-1}$  (provided  $L - 1 > 1$ ). It turns out that the resulting optimal value of the objective depends only on  $s_{L,1}, s_{L,2}, s_{L-1,1}, s_{L-1,2}$ . Thus, we can now fix  $s_{L-1,1}, s_{L-1,2}$  and  $H_{1}, W_{1}, \ldots, W_{L-3}$ , and optimize over  $W_{L-2}$  (provided  $L - 2 > 1$ ). The optimal solution will, yet again, only depend on the singular values of  $\sigma(H_{L-1}), \sigma(H_{L-2})$ . Thus, we can fix  $s_{L-2,1}, s_{L-2,2}$  and  $H_{1}, W_{1}, \ldots, W_{L-4}$  and repeat the procedure (provided  $L - 3 > 1$ ). In this way, we optimize over  $W_{l}$  for all  $l \geq 2$ .
We formulate one step of this nested optimization chain in the following key lemma. Here,  $X$  abstractly represents  $H_{l-1}$  and  $\{s_1, s_2\}$  represent  $\{s_{l,1}, s_{l,2}\}$ . A proof sketch is provided below and the complete argument is deferred to Appendix A.
Lemma 5. Let  $X$  be any fixed two-column entry-wise non-negative matrix with at least two rows, and let its singular values be given by  $\{\tilde{s}_1, \tilde{s}_2\} = S(X)$ . Then, for a given pair of singular values  $\{s_1, s_2\}$  and a given dimension  $d \geq 2$ , the optimization problem
$$
\min  _ {W} \| W \| _ {F} ^ {2}, \quad \text {s . t .} \{s _ {1}, s _ {2} \} = \mathcal {S} (\sigma (W X)),
$$
further constrained so that the number of rows of  $WX$  is  $d$ , has optimal value equal to
$$
\frac {s _ {1}}{\tilde {s} _ {1}} + \frac {s _ {2}}{\tilde {s} _ {2}},
$$
if  $X$  is full rank. If  $X$  is rank 1, then the optimal value is  $\frac{s_1}{s_1}$  as long as  $s_2 = 0$ , otherwise the problem is not feasible. If  $X = 0$ , then necessarily  $s_1 = s_2 = 0$  with optimal solution 0, otherwise the problem is not feasible.
Proof sketch of Lemma 5. Notice that, for any fixed value of the matrix  $A \coloneqq \sigma(WX)$ , there is at least one  $W$  which achieves such a value of  $A$ , unless the rank of  $A$  is bigger than the rank of  $X$ , in which case the problem is not feasible (as mentioned in the statement). Therefore, we can split the optimization into two nested sub-problems. The inner sub-problem fixes an output  $A$  for which  $\{s_1, s_2\} = S(A)$  and optimizes only over all  $W$  for which  $A = \sigma(WX)$ . The outer problem then takes the optimal value of this inner problem and optimizes over all the feasible outputs  $A$  satisfying the constraint on the singular values. We phrase the nested sub-problems as two separate lemmas.
Lemma 6. Let  $A$  with columns  $a, b$  and  $X$  with columns  $x, y$  be fixed two-column matrices with at least two rows, where  $X$  is entry-wise non-negative. Then, the optimal value of  $\| W \|_F^2$  subject to the constraint  $A = \sigma(WX)$  is
$$
\frac {a ^ {T} a \cdot y ^ {T} y - 2 a ^ {T} b \cdot x ^ {T} y + b ^ {T} b \cdot x ^ {T} x}{x ^ {T} x \cdot y ^ {T} y - (x ^ {T} y) ^ {2}}, \tag {4}
$$