Title: Deep Neural Collapse Is Provably Optimal for the Deep Unconstrained Features Model
Venue: NeurIPS 2023
Accepted: yes
Keywords: neural collapse

unconstrained features model

deep learning
Abstract: 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 work 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.

Deep Neural Collapse Is Provably Optimal for the Deep Unconstrained Features Model

Abstract

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.

1 Introduction

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

104

DNC1

106

DNC2

DNC3

layer_1

layer_1

100

layer_2

layer_2

102

layer_3

104

layer_3

10-1

layer_2_pre_relu

10°

102

10-2

10-2

layer_1

100

10-3

layer_2

10-4

layer_3

10-4

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

1e3

1e3

1e3

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.

2 Related work

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.

3 Preliminaries

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 , . . . , d L > 1 and weight matrices W 1 ∈ R d 2 × d 1 , W 2 ∈ R d 3 × d 2 , . . . , W L ∈ R K × d l . Let H 1 ∈ R d 1 × N , H 2 ∈ R d 2 × N , . . . , H L ∈ R d L × N be the feature matrices before the application of the ReLU activation function, which is denoted by σ , i.e., H 2 = W 1 H 1 and H l = W l -1 σ ( H l -1 ) for l ≥ 3 . When indexing a particular sample, we denote by h ( l ) c,i the i -th sample of the c -th class in the l -th layer, i.e., a particular column of H l . Moreover, let µ ( l ) c = 1 n ∑ n i =1 h ( l ) c,i be the activation class-mean for layer l and class c , and ¯ H l the matrix containing µ ( l ) c 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 ∈ R K × N can be written as I K ⊗ 1 T n , where I K denotes a K × K identity matrix, ⊗

denotes the Kronecker product and 1 n a row vector of all-ones of length n . With this, we can write ¯ H l = 1 n H l ( I K ⊗ 1 n ) and we further denote σ ( H l ) := 1 n σ ( H l )( I K ⊗ 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: 2

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 σ . In the former case, the condition requires h ( l ) c,i = h ( l ) c,j for all i, j ∈ [ n ] (or, in matrix notation, H l = ¯ H l ⊗ 1 T n ); in the latter, σ ( h ( l ) c,i ) = σ ( h ( l ) c,j ) for all i, j ∈ [ n ] (or, in matrix notation, σ ( H l ) = σ ( H l ) ⊗ 1 T n ).

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 σ . In the former case, the condition requires ¯ H T l ¯ H l ∝ I K ; in the latter, σ ( H l ) T σ ( H l ) ∝ 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 σ ( 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:

where ∥·∥ F denotes the Frobenius norm and λ H 1 , λ W 1 , . . . , λ 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.

More precisely, DNC1 is present on all layers; DNC2 is present on σ ( H ∗ l ) for l ≥ 2 and on H ∗ l for l ≥ 3 ; and DNC3 is present for l ≥ 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 , . . . , W ∗ L ) = (0 , 0 , . . . , 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

(although without considering the DNC1 property on H 1 , H 2 ), but it requires the extra assumption ∥ σ ( H ∗ 2 ) ∥ ∗ = ∥ H ∗ 2 ∥ ∗ , where ∥·∥ ∗ denotes the nuclear norm. In particular, the argument of [28] crucially relies on the fact that ∥ σ ( H ∗ 2 ) ∥ ∗ ≤ ∥ H ∗ 2 ∥ ∗ . We remark that the inequality ∥ σ ( M ) ∥ ∗ ≤ ∥ 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 ( · ) an operator which takes any matrix A ∈ R e × 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 σ ( H l ) for l ≥ 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 σ ( H l ) , for l ≥ 2 . Thus, the final objective can be reformulated in terms of S ( σ ( H l )) , for l ≥ 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 ≤ l ≤ L and { s l, 2 } 2 ≤ l ≤ 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 λ H 1 , λ W 1 , . . . , λ W L ). Hence, the optimal s l,j will be equal for j = 1 and j = 2 , which in turn yields that σ ( 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 { σ ( H l ) } l ≥ 2 (keeping them fixed). We start by optimizing over W L and then optimize over W l for 2 ≤ l ≤ 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 ≤ l ≤ 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 , . . . , W L -1 . Then, the optimal solution of the optimization problem

is achieved at W L = σ ( H L ) T ( σ ( H L ) σ ( H L ) T +2 λ W L I d L ) -1 and equals

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 σ ( 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 , . . . , 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 , . . . , 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 σ ( H L -1 ) , σ ( H L -2 ) . Thus, we can fix s L -2 , 1 , s L -2 , 2 and H 1 , W 1 , . . . , W L -4 and repeat the procedure (provided L -3 > 1 ). In this way, we optimize over W l for all l ≥ 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 { ˜ s 1 , ˜ s 2 } = S ( X ) . Then, for a given pair of singular values { s 1 , s 2 } and a given dimension d ≥ 2 , the optimization problem

further constrained so that the number of rows of WX is d , has optimal value equal to

if X is full rank. If X is rank 1, then the optimal value is 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 := σ ( 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 = σ ( 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 ∥ 2 F subject to the constraint A = σ ( WX ) is

unless x and y are aligned. In that case, the optimal value is b T b y T y , and we require that a, b are aligned and ∥ y ∥ / ∥ x ∥ = ∥ b ∥ / ∥ a ∥ for the problem to be feasible. Moreover, the matrix W ∗ X at the optimal W ∗ is non-negative and thus W ∗ X = σ ( W ∗ X ) . If the matrix X is orthogonal, then the rows of the optimal W ∗ are either 0, or aligned with one of the columns of X.

Proof sketch of Lemma 6. The optimization problem is separable with respect to the rows of W , hence we can solve for each row separately. As the problem is convex and has linear constraints, the KKT conditions are necessary and sufficient for optimality. By solving such conditions explicitly and computing the objective value, the result follows. The details are deferred to Appendix A.

Now, we can plug (4) into the outer optimization problem, where we solve for the output matrix A given the constraint on singular values. We only optimize over the numerator of (4), as the denominator does not depend on the variable.

Lemma 7. Let X be a fixed entry-wise non-negative rank-2 matrix (the rank-1 case is treated separately later) with columns x, y and A = σ ( WX ) . The minimum of the following optimization problem over this non-negative matrix A with columns a, b :

Proof sketch of Lemma 7. The optimization problem (5) is non-convex and has non-linear constraints. It is still solvable using the KKT conditions, however some care is required as such conditions are no longer necessary and sufficient. Using the Linear Independence Constraint Qualification (LICQ) rule, we handle the necessity. Then, by solving for KKT conditions and comparing the objective value with solutions for which the LICQ condition does not hold, we obtain the desired description of the optimal solution. The details are deferred to Appendix A.

̸

At this point, we are ready to obtain the result of Lemma 5. If ˜ s 1 , ˜ s 2 = 0 then, by combining the result of Lemma 7 with (4), the proof is complete. If ˜ s 2 = 0 , we can optimize (4), given s 2 = 0 and the extra constraint a T a b T b = x T x y T y , which follows from the requirement on the rank of A and the fact that X has rank 1. The claim then follows from a direct computation of the objective value using the constraint. The case ˜ s 1 = ˜ s 2 = 0 gives that X = 0 , and is trivial.

So far, we have reduced the optimization over { W l } L -1 l =2 to the optimization over the singular values of { σ ( H l ) } L l =2 . We now do the same with the first layer's matrices W 1 , H 1 . For this, we again split the optimization of λ W 1 2 ∥ W 1 ∥ 2 F + λ H 1 2 ∥ H 1 ∥ 2 F given { s 1 , s 2 } = S ( σ ( W 1 H 1 )) into two nested sub-problems: the inner one performs the optimization for a fixed output H 2 := W 1 H 1 ; the outer one optimizes over the output H 2 , given the optimal value of the inner problem.

The inner sub-problem is solved by the following Lemma, which describes a variational form of the nuclear norm. The statement is equivalent to Lemma C.1 of [28].

Lemma 8. The optimization problem

is minimized at value √ λ A λ B ∥ C ∥ ∗ , and the minimizers are of the form A ∗ = γ A U Σ 1 / 2 R T , B ∗ = γ B R Σ 1 / 2 V T . Here, the constants γ A , γ B only depend on λ A , λ B ; U Σ V T is the SVD of C ; and R is an orthogonal matrix.

To be more concrete, we apply Lemma 8 with A = W 1 and B = H 1 (so that C = W 1 H 1 = H 2 ). Thus, we are only left with optimizing ∥ H 2 ∥ ∗ given S ( σ ( H 2 )) = { s 1 , s 2 } . We solve this problem in the following, much more general lemma, which might be of independent interest. Its proof is deferred to Appendix A.

Lemma 9. Let L ≥ 2 be a positive integer. Then, the optimal value of the optimization problem

equals ( s 1 + s 2 ) 2 L . Here ∥·∥ S p denotes the p -Schatten pseudo-norm.

Remark 10. As a by-product of the analysis in the proof of Lemma 9, we obtain a characterization of the minimizers of (7) , which will be useful in Step 2 . Let H ∗ be an optimizer of (7) . Then, σ ( H ∗ ) has orthogonal columns. For L > 2 , its negative entries are set uniquely so that H ∗ has rank 1. For L = 2 , its negative entries are arranged as follows. All the rows of σ ( H ∗ ) which contain exactly one positive entry on the first column will result in rows of H ∗ that are multiples of each other. Similarly, all the rows of σ ( H ∗ ) which contain exactly one positive entry on the second column will result in rows of H ∗ that are multiples of each other. The zero rows of σ ( H ∗ ) remain zero rows in H ∗ . Moreover, the l 2 norm of the negative entries of the rows of H ∗ with positive entry in the first column equals the l 2 norm of the negative entries of the rows of H ∗ where the positive entry is in the second column. Finally, taking any two rows of H ∗ with counterfactual positioning of the signs of the entries, the product of negative entries does not exceed the product of positive entries.

Armed with Lemma 4, 5 and 9, we reformulate the L -DUFM objective (1) only in terms of { s l, 1 , s l, 2 } 2 ≤ l ≤ L . We then split this joint problem into two separate, identical optimizations over { s l, 1 } 2 ≤ l ≤ L and { s l, 2 } 2 ≤ l ≤ L , which are solved via the lemma below (taking x l := s 2 l, 1 and x l := s 2 l, 2 , respectively). Its proof is deferred to Appendix A.

Lemma 11. The optimization problem

is optimized at an entry-wise positive, unique solution if

and at ( x 2 , . . . , x L ) = (0 , . . . , 0) otherwise. When (9) holds with equality, both solutions are optimal. Here, 0 / 0 is defined to be 0 .

Note that (9) is equivalent to (2) when n = 1 . Furthermore, if the optimization problem (8) has a single non-zero solution, then the optimal s ∗ l, 1 and s ∗ l, 2 are forced to coincide for all l . From this, we obtain the orthogonality of σ ( H ∗ l ) for l ≥ 2 . Lemma 6 then gives that the optimal H ∗ l for l ≥ 3 must be orthogonal as well, since it is non-negative and equal to σ ( H ∗ l ) . Finally, the statement about W ∗ l for l ≥ 2 also follows from Lemma 6. This concludes the proof of DNC2-3 for n = 1 . As a by-product of Lemma 8, we additionally obtain a precise characterization of the optimal W ∗ 1 , H ∗ 1 .

Step 2: DNC1 for n > 1 . Now, we assume n > 1 and the conclusion of Step 1 . We show that DNC1 holds for the optimizer of (1). We give a sketch here, deferring the proof to Appendix A.

̸

We proceed by contradiction. Assume there exist ˜ H ∗ 1 , W ∗ 1 , . . . , W ∗ L which achieve the minimum of (1), but ˜ H ∗ 1 does not exhibit DNC1. Note that the objective in (1) is separable w.r.t. the columns of H 1 . Thus, for both classes c ∈ { 1 , 2 } , the partial objectives corresponding to the columns { h (1) c,i } n i =1 of ˜ H ∗ 1 must be equal for all i . Indeed, if that is not the case, we can exchange all the columns of ˜ H ∗ 1 with the ones that achieve the smallest partial objective value and obtain a contradiction. Now, we can construct an alternative ¯ H ∗ 1 such that, for each class c , we pick any h (1) c,i and place it instead of { h (1) c,j } j = i . By construction, ¯ H ∗ 1 exhibits DNC1 and it is still an optimal solution. At this point, let us construct H ∗ 1 from ¯ H ∗ 1 by taking a single sample from both classes. We claim that ( H ∗ 1 , W ∗ 1 , . . . , W ∗ L ) is an optimal solution of (1) with n = 1 and nλ H 1 as the regularization term for H 1 , while keeping d 1 , . . . , d L and λ W 1 , . . . , λ W L the same. To see why this is the case, consider an alternative G ∗ 1 achieving a smaller loss. Then, ( G ∗ 1 ⊗ 1 T n , W ∗ 1 , . . . , W ∗ L ) would achieve a smaller loss than ( ¯ H ∗ 1 , W ∗ 1 , . . . , W ∗ L ) , which contradicts its optimality.

̸

Recall that H ∗ 1 is an optimal solution of (1) with n = 1 . By Step 1 , the columns x, y of σ ( H ∗ 2 ) are such that x, y ≥ 0 , x T y = 0 , and ∥ x ∥ = ∥ y ∥ (here, by the notation x ≥ 0 we mean an entry-wise inequality). Assume σ ( H ∗ 2 ) = 0 . First we show that x, y and thus the whole σ ( H ∗ 2 ) are uniquely determined by ( W ∗ 1 , W ∗ 2 , . . . , W ∗ L ) . For this we use the uniqueness of the scaling obtained in Lemma 11 and the DNC3 property for the ( σ ( H ∗ 2 ) , W ∗ 2 ) pair. Then, by Lemma 9, all the matrices A 2 such that σ ( A 2 ) = σ ( H ∗ 2 ) and ∥ A 2 ∥ ∗ = ∥ H ∗ 2 ∥ ∗ have columns of the form x -ty, y -tx for t ∈ [0 , 1] . Let A t 2 be a matrix of that form for a particular t. Aseries of steps which include computing the SVD of A t 2 and the usage of Lemma 8 reveal that, among 0 ≤ t ≤ 1 , there is a single t ∗ and A ∗ 2 = A t ∗ 2 such that W ∗ 1 solves the optimization problem (6) in Lemma 8 if H 2 = A t 2 . Thus, having fixed A ∗ 2 that can be the output of the first layer, this in turn gives the uniqueness of the unconstrained features H 1 solving (6) while W 1 = W ∗ 1 , H 2 = H ∗ 2 = A ∗ 2 . However, as ˜ H ∗ 1 is not DNC1 collapsed by assumption, there are multiple possible choices for ¯ H ∗ 1 and then H ∗ 1 , which gives a contradiction with the uniqueness of H 1 solving (6).

Step 3: DNC2-3 for n > 1 . Having established DNC1, we obtain DNC2-3 by reducing again to the case n = 1 , which allows us to conclude the proof of the main result. We finally remark that W ∗ 1 and H ∗ 1 do not need to be orthogonal, due to the Σ 1 / 2 multiplication in the statement of Lemma 8. This means that W ∗ 1 and H ∗ 1 may not exhibit DNC2-3 (as stated in Theorem 3).

4.2 Numerical results

DUFM training. To support the validity of our theoretical result and to demonstrate that the global optimum in the DUFM can be found efficiently, we conduct numerical experiments in which we solve (1) with randomly initialized weights using gradient descent. In Figure 2, we plot the training progression of the DNC metrics for two choices of the depth: L = 3 (first row) and L = 6 (second row). In both cases, we use n = 50 training samples and constant layer width of 64 . The learning rate is 0 . 5 , though smaller learning rates produce the same results. We run full gradient descent for 10 5 steps, and the weight decay is set to 5 · 10 -4 for all the layers. The plots are based on 10 experiments and the error bands are computed as one empirical standard deviation to both sides. To measure the DNC occurrence, we plot the following metrics as a function of training. For DNC1, we define Σ l W = 1 N ∑ 2 c =1 ∑ n i =1 ( h ( l ) c,i -µ ( l ) c )( h ( l ) c,i -µ ( l ) c ) T as the within-class variability at layer l and Σ l B = 1 2 ∑ 2 c =1 ( µ ( l ) c -µ ( l ) G )( µ ( l ) c -µ ( l ) G ) T as the between-class variability, where µ ( l ) G is the global feature mean at layer l . Then, the DNC1 metric at layer l is ∥ Σ W Σ † B ∥ 2 F , with † being the pseudo-inverse operator. This is an upper bound on the trace of the same matrix, which is sometimes used in the literature for measuring NC1. An ideal DNC1 would show 0 on this metric. For DNC2, our metric at layer l is given by s l, 1 s l, 2 (recall that these are the singular values of σ ( H l ) ), and we consider the feature matrices both after and before the ReLU activation. An ideal DNC2 would show 1 on this metric. For DNC3, we consider the weighted average sine of the angles between each row of W l and the closest column of σ ( H l ) ; the weights are the l 2 norms of the rows of W l . To avoid

DNC2

DNC1

DNC3

10°

layer_1

101

layer_1

10°

layer_2

10-3

layer_2

layer_3

layer_3

10-2

layer_2_pre_relu

10-4

10-6

10-9

10-6

layer_1

10-12

10-8

layer_2

layer_3

10°

10-15

10-101

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

08

100

step

1e3

step

1e3

step

1e3

DNC2

DNC1

DNC3

layer_1

100

10°

layer_1

layer_2

layer_2

10-3

layer_3

101

layer_3

10-2

layer_4

layer_4

10-4

layer_1

10-6

layer_5

layer_5

layer_2

10-9

layer_6

layer_2_pre_relu

layer_6

10-6

layer_3

layer_4

10-12

10-8

layer_5

10°

layer_6

10-15

10-10

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

step

1e3

step

1e3

step

1e3

numerical issues, we exclude rows of W l whose l 2 norm is smaller than 10 -6 . An ideal DNC3 would show 0 on this metric.

The plots in Figure 2 closely follow the predictions of Theorem 3 and, taken as a whole, show that DUFM training via gradient descent clearly exhibits deep neural collapse, for all the three metrics considered, and for both depths. More specifically, DNC1 is progressively achieved in all layers, with the layers closer to the output being faster. DNC2 is achieved for all σ ( H l ) with l ≥ 2 and H l for l ≥ 3 (though we didn't include this in plots to avoid too many lines) and, as predicted by our theory, it is not achieved on H 2 (the dotted line in the middle plot) and on H 1 (the blue line). Interestingly, for these numbers of layers, SGD does not find solutions which would exhibit DNC2 on H 1 , even though such solutions exist. Hence, gradient descent is not implicitly biased towards them, see Appendix C for a more detailed discussion. DNC3 is achieved on all layers - even on the first, which is not covered by our theory and where the effect is the most pronounced; the deeper layers achieve values of DNC3 of order 10 -2 , which translates into an angle of less than one degree between weight rows and feature columns. Finally, we highlight that the training loss matches the optimal value computed numerically via (17) (the corresponding plot is reported in Appendix C), and hence the solution found by gradient descent is globally optimal. Additional ablation studies concerning the number of layers, weight decay and width of the layers are reported in Appendix C.

ResNet20 training with pre-trained DUFM. We support the validity of DUFM as a modeling principle for deep neural collapse via the following experiments. We first train an L -DUFM for L = 3 , 5 with K = 2 and n = 5000 to full optimality, as in the previous paragraph. Then, we replace H ∗ 1 with the output of a randomly initialized backbone, we fix the weight matrices in the DUFM, and we train end-to-end to convergence on the classes 0 and 1 of CIFAR-10. In Figure 1, we report the DNC metrics of our 3 -DUFM head, as the training of the backbone ResNet20 progresses. While at the start of the training the features H 1 are random and, therefore, the model does not exhibit DNC, the ResNet20 recovers well the DNC properties after training. This suggests that the crucial assumption of DUFM - the unconstrained features, indeed holds in practice. The plots are based on 7 experiments with confidence bands of one empirical standard deviation to both sides. Additional details as well as the experiments with 5 -DUFM and VGG13 and DenseNet121 are discussed in Appendix D.

5 Conclusion

In this work, we propose the deep unconstrained features model (DUFM) - a multi-layered version of the popular unconstrained features model - as a tool to theoretically understand the phenomenon of deep neural collapse (DNC). For binary classification, we show that DNC is globally optimal in the DUFM. This provides the first rigorous evidence of neural collapse occurring for arbitrarily many layers. Our numerical results show that gradient descent efficiently finds a solution in agreement with our theory - one which displays neural collapse in multiple layers. Finally, we provide evidence in favor of DUFM as a modeling principle, by showing that standard DNNs can learn the unconstrained features that exhibit all DNC properties.

Interesting open problems include (i) the generalization to multiple classes ( K > 2 ), (ii) the analysis of the gradient dynamics in the DUFM, and (iii) understanding the impact of biases on the collapsed solution, as well as that of the cross-entropy loss. As for the first point, a crucial issue is to control ∥ W 1 H 1 ∥ ∗ , since the inequality ∥ σ ( M ) ∥ ∗ ≤ ∥ M ∥ ∗ does not hold for general matrices M . In addition, while Lemma 5 still appears to be valid for K > 2 (and it would be an interesting stand-alone result), its extension likely requires a new argument.

Acknowledgments and Disclosure of Funding

M. M. is partially supported by the 2019 Lopez-Loreta Prize. The authors would like to sincerely thank Eugenia Iofinova, Bernd Prach and Simone Bombari for valuable feedback to our manuscript.

References

Ido Ben-Shaul and Shai Dekel. Nearest class-center simplification through intermediate layers. In Topological, Algebraic and Geometric Learning Workshops , 2022.

Hien Dang, Tan Nguyen, Tho Tran, Hung Tran, and Nhat Ho. Neural collapse in deep linear network: From balanced to imbalanced data. arXiv preprint arXiv:2301.00437 , 2023.

Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training. In Proceedings of the National Academy of Sciences (PNAS) , volume 118, 2021.

Tomer Galanti, Liane Galanti, and Ido Ben-Shaul. On the implicit bias towards minimal depth of deep neural networks. arXiv preprint arXiv:2202.09028 , 2022.

Tomer Galanti, András György, and Marcus Hutter. Improved generalization bounds for transfer learning via neural collapse. In First Workshop on Pre-training: Perspectives, Pitfalls, and Paths Forward at ICML , 2022.

Jarrod Haas, William Yolland, and Bernhard T Rabus. Linking neural collapse and l2 normalization with improved out-of-distribution detection in deep neural networks. Transactions on Machine Learning Research (TMLR) , 2022.

X. Y. Han, Vardan Papyan, and David L Donoho. Neural collapse under mse loss: Proximity to and dynamics on the central path. In International Conference on Learning Representations (ICLR) , 2022.

Hangfeng He and Weijie J Su. A law of data separation in deep learning. arXiv preprint arXiv:2210.17020 , 2022.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision (ICCV) , 2015.

Arthur E Hoerl and Robert W Kennard. Ridge regression: Biased estimation for nonorthogonal problems. In Technometrics , volume 12, 1970.

Like Hui, Mikhail Belkin, and Preetum Nakkiran. Limitations of neural collapse for understanding generalization in deep learning. arXiv preprint arXiv:2202.08384 , 2022.

Wenlong Ji, Yiping Lu, Yiliang Zhang, Zhun Deng, and Weijie J Su. An unconstrained layerpeeled perspective on neural collapse. In International Conference on Learning Representations (ICLR) , 2022.

Vignesh Kothapalli. Neural collapse: A review on modelling principles and generalization. In Transactions on Machine Learning Research (TMLR) , 2023.

Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum margin bias of quasi-homogeneous neural networks. arXiv preprint arXiv:2210.03820 , 2022.

Xiao Li, Sheng Liu, Jinxin Zhou, Xinyu Lu, Carlos Fernandez-Granda, Zhihui Zhu, and Qing Qu. Principled and efficient transfer learning of deep models via neural collapse. arXiv preprint arXiv:2212.12206 , 2022.

Zexi Li, Xinyi Shang, Rui He, Tao Lin, and Chao Wu. No fear of classifier biases: Neural collapse inspired federated learning with synthetic and fixed classifier. arXiv preprint arXiv:2303.10058 , 2023.

Tong Liang and Jim Davis. Inducing neural collapse to a fixed hierarchy-aware frame for reducing mistake severity. arXiv preprint arXiv:2303.05689 , 2023.

Weiyang Liu, Longhui Yu, Adrian Weller, and Bernhard Schölkopf. Generalizing and decoupling neural collapse via hyperspherical uniformity gap. In International Conference on Learning Representations (ICLR) , 2023.

Jianfeng Lu and Stefan Steinerberger. Neural collapse under cross-entropy loss. In Applied and Computational Harmonic Analysis , volume 59, 2022.

Jiawei Ma, Chong You, Sashank J Reddi, Sadeep Jayasumana, Himanshu Jain, Felix Yu, ShihFu Chang, and Sanjiv Kumar. Do we need neural collapse? Learning diverse features for fine-grained and long-tail classification. 2023.

Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. arXiv preprint arXiv:2011.11619 , 2020.

Vardan Papyan, X. Y. Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. In Proceedings of the National Academy of Sciences (PNAS) , volume 117, 2020.

Tomaso Poggio and Qianli Liao. Explicit regularization and implicit bias in deep network classifiers trained with the square loss. arXiv preprint arXiv:2101.00072 , 2020.

Tomaso Poggio and Qianli Liao. Implicit dynamic regularization in deep networks. Technical report, Center for Brains, Minds and Machines (CBMM), 2020.

Akshay Rangamani and Andrzej Banburski-Fahey. Neural collapse in deep homogeneous classifiers and the role of weight decay. In IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , 2022.

Akshay Rangamani, Marius Lindegaard, Tomer Galanti, and Tomaso Poggio. Feature learning in deep classifiers through intermediate neural collapse. Technical Report , 2023.

Christos Thrampoulidis, Ganesh Ramachandra Kini, Vala Vakilian, and Tina Behnia. Imbalance trouble: Revisiting neural-collapse geometry. In Conference on Neural Information Processing Systems (NeurIPS) , 2022.

Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural collapse. In International Conference on Machine Learning (ICML) , 2022.

Tom Tirer, Haoxiang Huang, and Jonathan Niles-Weed. Perturbation analysis of neural collapse. arXiv preprint arXiv:2210.16658 , 2022.

Peng Wang, Huikang Liu, Can Yaras, Laura Balzano, and Qing Qu. Linear convergence analysis of neural collapse with unconstrained features. In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop) , 2022.

E Weinan and Stephan Wojtowytsch. On the emergence of simplex symmetry in the final and penultimate layers of neural network classifiers. In Mathematical and Scientific Machine Learning , 2022.

Liang Xie, Yibo Yang, Deng Cai, and Xiaofei He. Neural collapse inspired attraction-repulsionbalanced loss for imbalanced learning. In Neurocomputing , 2023.

Mengjia Xu, Akshay Rangamani, Qianli Liao, Tomer Galanti, and Tomaso Poggio. Dynamics in deep classifiers trained with the square loss: Normalization, low rank, neural collapse, and generalization bounds. In Research , volume 6, 2023.

Yibo Yang, Liang Xie, Shixiang Chen, Xiangtai Li, Zhouchen Lin, and Dacheng Tao. Do we really need a learnable classifier at the end of deep neural network? In Conference on Neural Information Processing Systems (NeurIPS) , 2022.

Yongyi Yang, Jacob Steinhardt, and Wei Hu. Are neurons actually collapsed? on the fine-grained structure in neural representations. In I Can't Believe It's Not Better Workshop: Understanding Deep Learning Through Empirical Falsification (NeurIPS 2022 Workshop) , 2022.

Can Yaras, Peng Wang, Zhihui Zhu, Laura Balzano, and Qing Qu. Neural collapse with normalized features: A geometric analysis over the riemannian manifold. arXiv preprint arXiv:2209.09211 , 2022.

John Zarka, Florentin Guth, and Stéphane Mallat. Separation and concentration in deep networks. In International Conference on Learning Representations (ICLR) , 2021.

Jiayao Zhang, Hua Wang, and Weijie Su. Imitating deep learning dynamics via locally elastic stochastic differential equations. In Conference on Neural Information Processing Systems (NeurIPS) , 2021.

Zhisheng Zhong, Jiequan Cui, Yibo Yang, Xiaoyang Wu, Xiaojuan Qi, Xiangyu Zhang, and Jiaya Jia. Understanding imbalanced semantic segmentation through neural collapse. arXiv preprint arXiv:2301.01100 , 2023.

Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In International Conference on Machine Learning (ICML) , 2022.

Jinxin Zhou, Chong You, Xiao Li, Kangning Liu, Sheng Liu, Qing Qu, and Zhihui Zhu. Are all losses created equal: A neural collapse perspective. In Conference on Neural Information Processing Systems (NeurIPS) , 2022.

Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geometric analysis of neural collapse with unconstrained features. In Conference on Neural Information Processing Systems (NeurIPS) , 2021.

A Proofs

Theorem 3. The optimal solutions of the L -DUFM (1) for binary classification ( K = 2 ) are s.t.

More precisely, DNC1 is present on all layers; DNC2 is present on σ ( H ∗ l ) for l ≥ 2 and on H ∗ l for l ≥ 3 ; and DNC3 is present for l ≥ 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 , . . . , W ∗ L ) = (0 , 0 , . . . , 0) .

Proof. We will split the proof into two distinct cases: n = 1 and n > 1 . The proof for n > 1 will require the statement to hold for n = 1 , so we start with this case. The explicit form of the problem goes as:

We will solve this problem using a sequence of L sub-problems, where in sub-problem l , we only optimize over matrices with index L -l +1 and condition on all the other matrices, as well as singular values of all σ ( H l ) for l ≥ 2 (keeping them fixed). Throughout the process, we will strongly rely on the crucial phenomenon that the optimal solutions of all the sub-problems will only depend on the singular values of the matrices σ ( H l ) . We will solve the problem starting from the last weight matrix W L , through W L -1 and so on, and finishing with the joint optimization over W 1 , H 1 . We divide this procedure into 4 lemmas. The first lemma will only treat W L . The second lemma will treat W l for 2 ≤ l ≤ L -1 . This lemma will only be used if L ≥ 3 . The third lemma will treat W 1 , H 1 jointly and the fourth, final lemma will then treat the only free variables left - the singular values of σ ( H l ) for all l ≥ 2 . Let us denote by S ( · ) an operator which returns a set of singular values of any given matrix (in full SVD, i.e. the SVD with number of singular values equal to the smaller dimension of the matrix). Moreover, let s l,i be the i -th singular value of σ ( H l ) .

Lemma 4. Fix H 1 , W 1 , . . . , W L -1 . Then, the optimal solution of the optimization problem

is achieved at W L = σ ( H L ) T ( σ ( H L ) σ ( H L ) T +2 λ W L I d L ) -1 and equals

Proof. This is a well-known fact from the theory of ridge regression [10]. However, for readers' convenience, we do the proof here too. If we denote the objective of this as L , then we get:

Since the objective is strongly convex in W L , the unique global optimum is found by setting to 0 the derivative. Solving for it, we have

Let us denote σ ( H L ) = U Σ V T the rectangular SVD of σ ( H L ) where U ∈ R d L × d L , Σ ∈ R d L × 2 , V ∈ R 2 × 2 . Writing the optimal W L in terms of SVD of σ ( H L ) we get:

Therefore, we obtain the SVD of W L and that (ΣΣ T +2 λ W L I d L ) -1 is a matrix of singular values of W L , from which we obtain:

Furthermore, the first summand of the objective is

Putting everything together, we get the following optimal value of the objective:

which gives the desired result.

Note that the term we optimized for is the only part of the objective in (10) which depends on W L . Now we are ready to state the crucial lemma (necessary when L ≥ 3 ), which is presented in an abstract way so it can be easily applied outside of the scope of this work.

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 { ˜ s 1 , ˜ s 2 } = S ( X ) . Then, for a given pair of singular values { s 1 , s 2 } and a given dimension d ≥ 2 , the optimization problem

further constrained so that the number of rows of WX is d , has optimal value equal to

if X is full rank. If X is rank 1, then the optimal value is 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. Notice that, for any fixed value of the matrix A := σ ( 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 , however this is a special case explicitly mentioned in the statement of the lemma for which the problem is non-feasible. Therefore, we can split the optimization into solving 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 = σ ( 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 will phrase both of the nested sub-problems as 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 ∥ 2 F subject to the constraint A = σ ( WX ) is

unless x and y are aligned. In that case, the optimal value is b T b y T y , and we require that a, b are aligned and ∥ y ∥ / ∥ x ∥ = ∥ b ∥ / ∥ a ∥ for the problem to be feasible. Moreover, the matrix W ∗ X at the optimal W ∗ is non-negative and thus W ∗ X = σ ( W ∗ X ) . If the matrix X is orthogonal, then the rows of the optimal W ∗ are either 0, or aligned with one of the columns of X.

Proof. Note that this optimization problem is separable in rows of W . Thus it suffices to solve the problem for a single row w j and then sum up over all rows. We will, moreover, split the analysis into two cases.

Case 1: Here we will assume that the row of A we are optimizing for is ( a j , b j ) , where a j , b j > 0 . In this case, we are facing the following (simple) convex optimization problem:

This has a known solution - the pseudoinverse as it is the minimum interpolating solution of a two-datapoint linear regression. However, using the formula would not simplify the expressions in the slightest and so we will solve this problem manually. The Lagrange function for this problem is

where d is the dimension of the columns of X and rows of W. We know that solving the KKT conditions for a convex problem with linear constraints yields sufficient and necessary conditions for the optimality. Let us write:

Together with constraints, these are the only KKT conditions for this problem. Solving for w ji we get:

Plugging in the constraints we get:

We can solve for v from the first equation and plug in to the second equation. Since y T y must be non-zero, otherwise the constraint w T j y = b j > 0 could not be satsfied, we know v and u are well-defined. We get:

unless x = αy , which we will deal with later. Now we can plug u in the expression on v and after some computations we get

Combining, we get the optimal w j :

If x = αy , then it is clear from (12) that b j x T y = a j y T y, from which we can deduce α = a j /b j . From the expression (11) for w ji , it is clear that the only solution in this case is w j = b j y T y y and the same must hold for all other j ∈ { 1 , . . . , ˜ d } , where ˜ d is the number of rows of W .

Case 2: Here we will assume that the row of A we are optimizing for is, w.l.o.g. ( a j , 0) , where a j > 0 . In this case, we are facing the following (simple) convex optimization problem:

Again, we can write the Lagrange function for this problem and again, solving for KKT conditions will provide full description of the optimal solution:

We obtain the same condition for w ji as in case 1. However, since one of our constraints is now an inequality, we will deal with complementarity conditions. Let us first assume that u = 0 . In that case w ji = -vx i / 2 and from the first constraint

Thus, w j = a j x T x x . However judging from this, to satisfy the first constraint, we must have x i y i = 0 for all i (remember that ∑ w ji y i must be non-positive and that x, y have non-negative entries). Therefore necessarily x T y = 0 . However, looking at the solution of case 1, we see that after plugging in x T y = 0 and b j = 0 , we get the same expression. Thus, this is a special case of case 1. The second option for complementarity condition is that ∑ w ji y i = 0 . Then, we can proceed in the same way as in case 1 to obtain the same formula after plugging in b = 0 . Then again, this is a special case of the formula in case 1.

All in all, we have fully solved the optimization problem in question. Now it remains to compute the optimal value. In the special case that x = αy , we must have bx T y = ay T y , otherwise there is no feasible solution. In this case this is equivalent to α = a/b, w j = b j / ( y T y ) y and w T j w j = b 2 j / ( y T y ) . Therefore, summing over all rows of W (and keeping in mind that once x = αy , the same characterization holds for every row) we get ∥ W ∥ 2 F = b T b y T y .

In the general case, let us compute:

Summing up over rows of W , we get the desired expression:

The rest of the statement of the lemma clearly follows from the formula on optimal rows of W.

Now, as promised, we will optimize over the output matrix A given that we know what the partial optimal value for each A is.

Lemma 7. Let X be a fixed entry-wise non-negative rank-2 matrix (the rank-1 case is treated separately later) with columns x, y and A = σ ( WX ) . The minimum of the following optimization problem over this non-negative matrix A with columns a, b :

̸

Proof. First note that since X is assumed to be rank 2, x = αy and so the expression in the objective is exactly the numerator of the optimal value of the minimization problem over W as stated in Lemma 6 (the denominator is constant within the scope of this lemma). Now, using the well-known formulas for the sum and product of squared singular values, we can easily rephrase our constraint in the following way:

First, noticing that both the constraints as well as the objective only depend on a T a, b T b, a T b , we can reduce the problem to one, where we treat these quantities as free variables that we are optimizing for. The only thing we need to remember will be implicit constraints that a T a, b T b ≥ 0 and ( a T b ) 2 ≤ a T ab T b , which is necessary from Cauchy-Schwarz inequality. We will solve this optimization problem using, again, KKT conditions. This time, we will use a stronger rule for deciding on the necessity of KKT conditions - the linear independence constraint qualification (LICQ) rule. According to this rule, any optimal solution has to fulfill the KKT conditions if the gradients of the equality constraints (in this case only two) are linearly independent. In our case, if we order the variables as ( a T a, b T b, a T b ) and writing the gradients as: (1 , 1 , 0) and ( b T b, a T a, -2 a T b ) we see that they can only be dependent if a T b = 0 and a T a = b T b . This is only a feasible solution if s 1 = s 2 . We will return to this case later, showing that it achieves the same objective value as the one coming from KKT conditions.

Now, let us write down the Lagrange function to this problem:

Computing the partial derivatives we get:

First consider what would happen if v = 0 . In this case, x T y = 0 from the third equation and x T x = y T y from the first two. Then, however, the objective function is constant on the feasible set and equals 1 / 2( s 2 1 + s 2 2 )(˜ s 2 1 + ˜ s 2 2 ) , which agrees with the statement of the Lemma.

̸

Next, assume v = 0 . Then we can express all the variables as follows:

Plugging this back into the constraints, we can solve for v, u . We get:

Plugging this into the second constraint together with the expressions for a T a, b T b and then simplifying we get:

Recall that the case s 1 = s 2 is treated separately, so we do not divide by zero. Now we can rewrite this expression w.r.t. { ˜ s 1 , ˜ s 2 } :

We can now simply solve for this to obtain the roots:

which after some simplifications yields:

The formula for v can now be obtained:

Having this, we can now compute the objective value using these quantities. We write it down for u + , v + :

which is the expression we wanted to prove. Had we done the same computation, but for u -, v -, we would get

which is at least as big as the previous expression, so we can drop this one. However, we still need to check, whether this solution of the KKT system is indeed feasible, because we imposed implicit constraints into the problem. Thus, we need to check: a T a, b T b ≥ 0 and a T ab T b ≥ ( a T b ) 2 . The first two are analogous. We have: a T a = -x T x + u v , but since the optimal v + is strictly negative, we can only analyze the sign of 2 x T x +2 u which is:

However, using the same formulas as the constraints in this optimization problem on ˜ s 1 , ˜ s 2 instead, one can, using very simple computations, check that this indeed holds. For the next condition, consider that a T ab T b -( a T b ) 2 must have the same sign as ( x T x + u )( y T y + u ) -( x T y ) 2 = x T xy T y -( x T y ) 2 + u 2 + u ( x T x + y T y ) . This resembles the quadratic function in (13). u + is a root of that quadratic function. Both this new quadratic function, as well as the one in (13) have the same coefficients in front of u 2 and u . Therefore, to test for the sign of the expression above, it only suffices to check, whether this alternative quadratic function has at least as big absolute constant:

The last inequality trivially holds. Therefore, the solution of KKT conditions is feasible.

Now, let us consider the special case when a T b = 0 , a T a = b T b and s 1 = s 2 for which the LICQ criterion is not satisfied. A very simple computation then yields that any feasible solution achieves the same objective value as the one above. This means that if the problem does have a global solution, it must be this one, otherwise the optimal solution would neither satisfy a KKT condition, nor be in the region where LICQ does not hold. Therefore, we must, finally, reason that this is indeed a globally optimal solution, which is not guaranteed by KKT conditions themselves. For this, it is only necessary to note that the feasible set is obviously compact. Since the objective value is continuous, we can use the well-known fact that such an optimization problem indeed has a global solution.

Now we are ready to finish the proof of Lemma 5. Let us first still assume ˜ s 1 , ˜ s 2 > 0 . Then, combining Lemma 6 with Lemma 7 we get:

equals

Now consider the case ˜ s 2 = 0 . According to Lemma 6 this implies s 2 = 0 . When optimizing over the matrix A we are facing the following optimization problem:

where the second constraint is clear from the considerations done in Lemma 6. Combining these two constraints, one simply gets

and plugging this in the objective value we have s 2 1 ˜ s 2 1 , which is consistent with the statement of Lemma 7. The case when X is 0 is trivial. This proves Lemma 5.

Notice that, for all indices but l = 1 , we managed to reduce the optimization over full matrices W l to the optimization over only the singular values of σ ( H k ) matrices. We will do the same with the first layer's matrices W 1 , H 1 . For this, we again split the optimization of λ W 1 2 ∥ W 1 ∥ 2 F + λ H 1 2 ∥ H 1 ∥ 2 F given { s 1 , s 2 } = S ( σ ( W 1 H 1 )) into two nested sub-problems, the inner one optimizing the quantity only for a fixed output H 2 := W 1 H 1 and the outer one optimizing over the output H 2 given the optimal value of the inner problem. We can solve the inner subproblem by a direct use of the following lemma, which describes a variational form of a nuclear norm.

Lemma 8. The optimization problem

is minimized at value √ λ A λ B ∥ C ∥ ∗ , and the minimizers are of the form A ∗ = γ A U Σ 1 / 2 R T , B ∗ = γ B R Σ 1 / 2 V T . Here, the constants γ A , γ B only depend on λ A , λ B ; U Σ V T is the SVD of C ; and R is an orthogonal matrix.

Using this Lemma 8, we know that the optimal value of the inner optimization problem is √ λ W 1 λ H 1 ∥ H 2 ∥ ∗ . Therefore we only need to solve the outer problem:

To solve this problem, we will state a much more general statement, which can be of independent interest, in the following lemma.

Lemma 9. Let L ≥ 2 be a positive integer. Then, the optimal value of the optimization problem

equals ( s 1 + s 2 ) 2 L . Here ∥·∥ S p denotes the p -Schatten pseudo-norm.

Proof. First note that this optimization problem has two free components: (i) σ ( H ) so that it satisfies the constraint, and (ii) H with σ ( H ) being fixed. This enables us to split the problem into inner and outer minimization. However, the complexity of the problem will prevent us from doing this conditioning in a straightforward way. Further, let us also implicitly note that within the proof we implicitly assume the matrix in question has at least three rows. The case with two rows is very simple and follows the proof with more rows.

We start by explicitly computing the objective value as a function of important statistics of H . Let σ ( H ) have columns r, s . Then,

Using the well-known formula on the eigenvalues of such matrix, we know λ 1 + λ 2 = r T r + s T s and λ 1 λ 2 = r T rs T s -( r T s ) 2 . Solving this system we get:

Let us, for now, fix this σ ( H ) . For a while we will only optimize over H s.t. σ ( H ) is fixed. We will think about the optimization problem in the following way: consider any δ ≥ 0 . Assume we are allowed to 'inject' negative values in the place of zero entries of a fixed σ ( H ) so that the sum of squares of those negative entries is exactly δ . What is the minimal value of the objective value given this δ and what is the optimal injection? Let us denote any quantity which was subject to such injection (or perturbation) with an extra lower-index P .

To answer this question, first have a closer look at (14). It is clear that after injecting the negative noise of total squared size δ , the r T P r P + s T P s P in a perturbed matrix will increase by exactly δ compared to r T r + s T s . Therefore, if we condition on δ , the only degree of freedom lies in the -4( r T P r P s T P s P -( r T P s P ) 2 ) in the perturbed matrix. Crucially, since the objective value is λ 1 /L 1 ,P + λ 1 /L 2 ,P , by concavity of 1 /L -th power, for any fixed sum of λ 1 ,P + λ 2 ,P = r T r + s T s + δ we want the λ i,P to be as far apart as possible. This is obtainable by maximizing the quantity -4( r T P r P s T P s P -( r T P s P ) 2 ) which is equivalent to minimizing r T P r P s T P s P -( r T P s P ) 2 . Importantly, as it will be clear from the forthcoming reasoning, the minimal value of λ 2 ,P is decreasing with

δ . Therefore, it only makes sense to consider all the δ values for which the minimal λ 2 ,P is still bigger than 0 up to the δ value for which the minimal λ 2 ,P hits 0 for the first time. From that δ on, every optimal λ 2 ,P is 0, but the λ 1 ,P increases and thus becomes sub-optimal. Now let us split the δ into δ 1 and δ 2 , denoting the total squared size of perturbations per column. If we also fix δ 1 and δ 2 and condition on them, even the quantity r T P r P s T P s P is fixed and we are left with maximizing ( r T P s P ) 2 . Denote p = r T s . To make the next argument clear, let us split the matrix σ ( H ) into four blocks of rows, where in each block there is one of the four following types of rows: (+ , +) , (+ , 0) , (0 , +) , (0 , 0) . We will now define variables a, b, c, d as sum of squared entries corresponding to one column of one block. Specifically:

The value we are trying to minimize equals, for an unperturbed σ ( H ) written in this way, ( a + c )( b + d ) -p 2 . Of course, σ ( H ) might not have all the types of blocks, in which case the empty sum is just defined as 0. First, we argue that whatever is the optimal assignment of the negative values given δ 1 , δ 2 , they do not appear in both the (0 , 0) block and (+ , 0) , (0 , +) blocks. This is because the only variable we are optimizing for in this conditioning is ( r T P s P ) 2 . Assume that, given the perturbation, the quantity r T P s P is negative. Then, any negative values present in the (0 , 0) block can be moved to (+ , 0) , (0 , +) blocks strictly decreasing r T P s P and thus increasing ( r T P s P ) 2 . On the other hand, if r T P s P happens to be positive, then any negative values present in the (+ , 0) , (0 , +) blocks can be moved to the (0 , 0) block, strictly increasing the r T P s P value and thus ( r T P s P ) 2 too. Therefore, we only need to consider cases when we inject the negative values to either the (0 , 0) block or to (+ , 0) , (0 , +) blocks, but not both at once. Let us start with the second case and analyze it. Clearly, if σ ( H ) did not have either the (0 , 0) block or (+ , 0) , (0 , +) blocks, then we do not even need this reasoning anymore. Also it is clear, as we will see later, that a matrix σ ( H ) which would completely lack 0 entries is suboptimal and thus there is always space to inject negative entries.

Case 1: Injecting negative values to (+ , 0) , (0 , +) blocks. For the same definition of a, b, c, d as above, we further define the squared sizes of total injected negative values in the perturbed H P per block as x, y , i.e.,

and thus x = δ 2 , y = δ 1 , x + y = δ. Trivially, such a construction only makes sense if it is possible to make r T P s P negative, otherwise it is strictly better to construct an alternative matrix σ ( H ) with a + c = a + c ; b + d = b + d ; p = p , but which includes the (0 , 0) block and put all the negative value there (this is always possible for matrices with at least 3 rows, since we only need two rows to construct σ ( H ) which satisfies the outlined constraints). This alternative matrix together with negative values in the (0 , 0) block would be better, because obviously ( r P T s P ) 2 > ( r T P s P ) 2 .

Once r T P s P is negative, though, it is clear that it is optimal, given the resources x, y , to produce as negative value as possible so as to maximize ( r T P s P ) 2 . By the use of Cauchy-Schwartz inequality, this is possible exactly when the injected negative vectors in the blocks are aligned with their positive counterparts of that block. Then, the inner product between the negative and positive part of a block will equal exactly √ cx or √ yd . Furthermore, if it is possible to construct σ ( H ) for which a + b < a + b ; a + c = a + c ; b + d = b + d ; p = p (which is only possible if there are at least two rows in the (+ , +) block and if p 2 < ab - thanks to the fact that then the angle between columns in the (+ , +) block is a free variable and we can decrease it while decreasing a, b to keep p unchanged using the law of cosine), then our objective value strictly decreases, since √ cx + √ yd > √ cx + √ dy . Therefore, it is always optimal for the first and second column of the (+ , +) block to be aligned, which also yields ab = p 2 . This value of p is necessary also if the (+ , +) block has only one row. Thus, we have that the original value of r T rs T s -( r T s ) 2 is ( a + c )( b + d ) -ab . The value after perturbation, using all the knowledge we acquired so far is ( a + c + y )( b + d + x ) -( √ ab - √ cx - √ dy ) 2 . Subtracting, we get a difference: by + ax + xy +2 √ abc √ x +2 √ abd √ y -2 √ cd √ xy . Consider now an alternative

.

see that either x = δ/ 2 or cd = x ( δ -x ) , which after solving yields x 1 , 2 = δ ± √ δ 2 -4 cd 2 . If δ 2 ≤ 4 cd then the only stationary point is δ/ 2 and this is clearly an optimal solution. For δ 2 = 4 cd , the optimal solution yields a rank 1 matrix so that λ 2 ,P = 0 . This means we do not need to care about δ 2 > 4 cd , as this would not achieve smaller value of λ 2 ,P , but would increase the value of λ 1 ,P . So, the solution where x = y = δ/ 2 is optimal and the objective value is δ 2 / 4 -s 1 s 2 δ. This is clearly smaller than 0 so we actually achieved a decrease in objective value.

matrix H P such that x, y remain unchanged and moreover a = b = 0 , c + d = a + b + c + d and cd = ( a + c )( b + d ) -ab , meaning that the σ ( H P ) has the same singular values as σ ( H P ) However, computing the same difference of the same quantities one gets: xy -2 √ cd √ xy. This value is strictly smaller than the previous one for any fixed pair of x, y . Therefore, a matrix σ ( H ) which does not contain the (+ , +) block is the only possible optimal solution for the injection which goes into (+ , 0) , (0 , +) blocks. Thus, we can reduce the previous problem into considering a matrix H P with a = b = 0 and minimizing the difference in objective values xy -2 √ cd √ xy, given x + y = δ . After plugging in y = δ -x , differentiating for x and putting equal to 0 , we get δ -2 x -√ cd √ x ( δ -x ) ( δ -2 x ) = 0 , which is equivalent to ( δ -2 x )(1 -√ cd √ x ( δ -x ) ) = 0 . From here, we

Case 2: Injecting negative values to (0 , 0) block. For this, it is not necessary to distinguish (+ , +) , (+ , 0) , (0 , +)) blocks, so we treat them as one. We can define:

The original objective value is ef -p 2 , while the perturbed one is ( e + x )( f + y ) -( p + √ xy ) 2 and so the difference is xf + ye -2 p √ xy. However, using Cauchy-Schwartz inequality, this is bigger or equal to xf + ye -2 √ fexy ≥ 0 . Thus, it is not possible to improve the objective value in this case, unlike the previous case where it was. Therefore, this case is clearly sub-optimal and we do not need to deal with it.

The only thing that could potentially go wrong with this proof is that there is not enough rows to build our necessary blocks. This can only happen if σ ( H ) is a 2 × 2 matrix. This case is, however, very simple and we will not explicitly treat it, but the statement and the proof path are the same for it.

We solved our optimization problem for every fixed allowed perturbation size δ . Now we just need to optimize over 0 ≤ δ ≤ 2 s 1 s 2 . So we face the following optimization problem:

Denote λ δ, 1 the value of the first summand of the above objective value (without 1 /L -th power) and similarly λ δ, 2 the value of the second summand (so those are the eigenvalues of optimal H T P H P given exactly δ allowed negative injection). The derivative of the objective w.r.t. δ equals:

where D δ := √ ( s 2 1 + s 2 2 + δ ) 2 -4( s 2 1 s 2 2 + δ 2 / 4 -s 1 s 2 δ denotes the square root part of the λ δ ;1 , 2 expression. We will prove that this derivative is either 0 for L = 2 or strictly negative on [0 , 2 s 1 s 2 ] everywhere except a set of measure 0 for L ≥ 3 . This guarantees for L = 2 the optimal value being equal to the unperturbed matrix while also defines the set of optimal solutions. For L ≥ 3 the optimal value is attained at δ = 2 s 1 s 2 , the objective value is easily expressible and the optimal solution is characterized by a single perturbation size. To test whether this expression is smaller than 0, nothing happens if we multiply it with 2 Lλ L -1 L δ, 1 λ 2 L -2 L δ, 2 , so we get:

We can now compute

to get

This expression resembles a form for which the use of Jensen inequality would be convenient. However, there are two issues. First, the coefficients, though summing up to 1 are not a convex combination since ( s 1 + s 2 ) 2 D δ ≥ 1 . So this is a combination outside of the convex hull of the inputs. However, Jensen inequality outside the convex hull is applicable with opposite sign. The bigger issue is that the function to use the Jensen inequality on, f ( x ) = x 2 L -2 L is only easily defined on R + . To make sure that we can safely use Jensen inequality, we first need to check that the linear combination, though outside of the convex hull of the inputs, still lies in R + . We are therefore asking whether

Now, we can multiply out and simplify this expression (after multiplying by 2):

but this holds even with equality. So we proved that the combination of inputs we are considering here is exactly 0 for which x 2 L -2 L is well-defined. So we can freely use Jensen inequality:

Note that if L = 2 , the function we applied Jensen inequality on was, in fact, linear and thus we retained equality. If L ≥ 3 , then since the function in question is strictly convex, the equality can only hold if λ δ, 2 = 2 s 1 s 2 -δ. After writing down the expression on λ δ, 2 it becomes clear that this is equivalent to a quadratic equation on δ, which can have at most two solutions. So we proved that the derivative is smaller than 0 on the whole interval of interest except a set of measure at most 0, where it is 0. Therefore, the only optimum is at δ = 2 s 1 s 2 . Plugging this in the objective value one gets ( s 1 + s 2 ) 2 L and from the proof we also get the precise description of the optimizers for both L = 2 and L > 2 cases.

Now we are ready to formally make the reduction of the optimization problem in (10) into the optimization over only singular values of the matrices σ ( H l ) for 2 ≤ l ≤ L. We formulate this reduction in a separate lemma for clarity.

Lemma 12. The optimal value of the following optimization problem:

with H 1 having two columns and all the matrices having at least two rows is

Now we are only left with optimizing the objective over all the s l,i . Crucially, the objective function is separable and symmetric in the i index. Therefore we can optimize for each index separately. Importantly, if there is only a single solution to this problem, then the solution must be the same for both indices i , which will yield orthogonality of the σ ( H l ) matrices. The ultimate goal is to show that this is indeed the case modulo some special circumstances. The optimal value of this reduced optimization problem is presented in the following, final lemma with abstract notation of the free variables.

Lemma 11. The optimization problem

is optimized at an entry-wise positive, unique solution if

and at ( x 2 , . . . , x L ) = (0 , . . . , 0) otherwise. When (9) holds with equality, both solutions are optimal. Here, 0 / 0 is defined to be 0 .

Proof. Here we will need to distinguish the cases L = 2 and L > 2 . We start with the case L = 2 . Let us reparametrize the problem so that we will work with y = √ x 2 . Weget the following expression to optimize the y over:

This function is simple enough so that we understand its shape. It has precisely one inflection point, below which it is concave, later it is convex. The derivative at 0 is strictly positive. Judging from this, the function either has 0 as a single global optimum, or it has two global optima, one at 0 and one at a strictly positive local minimum of the strictly convex part of the function, or the global minimum is achieved uniquely at the strictly convex part of the function. To find out, we only need to solve f ( y ) = f (0) for y, where f stands for the optimized function and determine the number of solutions. The quadratic equation this is equivalent to after excluding the point y = 0 as a solution is 4 y 2 √ λ W 1 λ H 1 -y +8 λ W 2 √ λ W 1 λ H 1 = 0 . This has a unique solution if and only if λ W 2 λ W 1 λ H 1 = 1 / 256 , which is precisely the threshold from the statement of the lemma. The cases outside of this threshold are easily mapped to the system having the global solution 0 or some positive point.

Let us now consider the case L > 2 . First note that, once for any l 0 the optimal x ∗ l 0 = 0 , then at this particular solution all the x 's must be 0. If l 0 = L, then this claim is apparent backward-inductively optimizing always over x with one index lower. If 2 ≤ l 0 < L , then the claim is trivial backwardinductively for all l ≤ l 0 , while for all l > l 0 this solution is forced by feasibility issues. From now on we will thus work with an implicit strict positiveness assumption. Taking derivatives of this objective w.r.t. x l for different l and putting them equal 0 we arrive at three types of equations depending on l :

Here we can notice two things. First, we have L -1 equations for L -1 variables, so we should expect a single solution in most of the cases. Second, the equations are rather simple and recursive.

Having either a value of x 2 or x L , all other values are uniquely determined. Let us denote q := x L x L -1 . Based on this, we can arrive at the following formula:

for all 1 ≤ k ≤ L -2 . An empty product is defined as 1. To derive this, for L = 3 the statement is only about x 2 and is trivial. For L > 3 , the k = 1 case is trivial, then one can proceed by induction using the recurrent formula x 2 l = λ Wl λ Wl -1 x l +1 x l -1 . From formula (16), one can derive another useful equation:

Combining this with (15), we also get x 2 = λ 2 WL -1 λ W 1 λ H 1 q 2 . Combining now this equation with formula (16) for x 2 from before one gets:

Notice that we managed to express all the terms present in the main objective as simple functions of q and regularization parameters. Plugging all of these expressions back in the objective, we get:

Since we only care about relative values of this expression, the optimization will be equivalent after multiplying it by 2 /λ W L -1 so we get:

Differentiating this function twice we can see that the second derivative is negative from 0 up to a single threshold at which it is 0, then it is positive. Therefore, based on the value of the first derivative at this threshold, and taking into account that the first derivative at 0 is strictly positive, we can easily see that the function has either a single global minimum at 0 or it has two global minima at 0 and at some positive q or it has a single global optimum at some positive q . To find out which of the cases we are at, we can simply solve for

After excluding the trivial solution q = 0 , this is equivalent to:

The case at which we stand w.r.t. the global solutions of the above objective function is equivalent to the number of solutions of this equation. Since the LHS of this equation is a very simple function, it is apparent that the number of solutions to this equation is equivalent to the sign of the LHS evaluated at the single zero-derivative positive q . The derivative is 2 λ W L -1 L 2 q L -1 -( L -1) q L -2 which after solving for q > 0 yields q = L -1 2 λ WL -1 L 2 . Plugging this back into the LHS above, we get the expression

Thresholding this for 0 and simplifying, we get the final

If the LHS is smaller than this threshold, the function in (18) has a single non-zero global optimum. If the LHS is at the threshold, the optimization problem has both 0 and some single non-zero point for a global solution. If the LHS is above the threshold, the optimization problem has only 0 as a single solution.

With all this, Theorem 3 follows very easily. We only need to conclude that if the optimization problem in Lemma 11 has a single non-zero solution, then both the optimal s ∗ l, 1 and s ∗ l, 2 are forced to be the same for each l. From this, we obtain the orthogonality of all the σ ( H ∗ l ) for l ≥ 2 . However, from Lemma 6 we know that even the optimal H ∗ l for l ≥ 3 must be orthogonal, since it is nonnegative and equal to σ ( H ∗ l ) . From the same lemma, we conclude the statement about W ∗ l for l ≥ 2 . The optimal W ∗ 1 , H ∗ 1 follow easily from Lemma 8.

By now, we only considered the case n = 1 . Now, we proceed to the general n case, where we need to also prove DNC1. We will proceed by contradiction. Assume there exist ˜ H ∗ 1 , W ∗ 1 , . . . , W ∗ L which optimize L -DUFM, but ˜ H ∗ 1 is not DNC1-collapsed. Now, since the objective function in 10 is separable in the columns of H 1 , we know

where h (1) c,i , h (1) c,j are columns of ˜ H ∗ 1 , for all c, i, j applicable and e c is the c -th standard basis vector. Otherwise, we could have just exchanged all sub-optimal h (1) c,i with the optimal one and achieve better objective value. However, if the last equality holds, then we can construct an alternative ¯ H ∗ 1 where for each class c , we pick any h (1) c,i and place it instead of all other h (1) c,j within the class c . In this way, the ¯ H ∗ 1 will be DNC1-collapsed, while still an optimal solution. However, for ¯ H ∗ 1 , we now know that if H ∗ 1 is constructed by taking just a single sample from both classes (thus forcing n = 1 ), this must be an optimal solution of the n = 1 L -DUFM with nλ H 1 as the regularization term for H 1 and same feature vector dimensions and other regularization terms. This is because if there was a better solution for the corresponding problem, then we could construct a counterexample to optimality in the original L -DUFM problem by just simply taking this alternative, better solution and using the feature vectors of this solution for each sample of each individual class.

̸

Therefore, since H ∗ 1 must be optimal in the new n = 1 problem (with a slight abuse of notation, we re-label the nλ H 1 back to λ H 1 ), we know that σ ( H ∗ 2 ) is collapsed in the way described by the theorem statement for n = 1 . By non-negativity of σ ( H ∗ 2 ) , this means that σ ( H ∗ 2 ) has two columns, let's call them x, y for which: x, y ≥ 0; x T y = 0 , ∥ x ∥ = ∥ y ∥ (by the notation x ≥ 0 we mean entry-wise inequality). This also implies that there is no such i that x i , y i > 0 . Moreover, let us assume σ ( H ∗ 2 ) = 0 . We can do that, as otherwise simple arguments would lead us into the degenerate solution where all W ∗ l = 0 , for which obviously all the columns in the originally optimal ˜ H ∗ 1 would need to be 0 implying simple DNC1 and reducing easily to the case of too big regularization explicitly stated in Theorem 3. We can now omit this case and assume non-triviality of the solution. We will first show that for a given set of optimal weight matrices ( W ∗ 1 , W ∗ 2 , . . . , W ∗ L ) , the matrix σ ( H ∗ 2 ) is uniquely determined regardless of what H ∗ 1 is. First, by the uniqueness of the solutions in Lemma 11 and by the fact that s 2 i depends monotonically on ∥ x ∥ , it follows that for any candidate direction z, ∥ z ∥ = 1 , there is a unique multiplier α > 0 for which αz can be a column of σ ( H ∗ 2 ) , i.e. there is a uniquely determined norm under which any candidate direction z must enter σ ( H ∗ 2 ) . Now, however, by DNC3 for l = 2 we know that rows of W 2 are either zero or aligned with the columns of σ ( H ∗ 2 ) . Moreover, since H ∗ 3 (which in case L = 2 shall be considered the output of the network) is orthogonal and thus rank 2, the weight matrix W ∗ 2 must contain both the columns of σ ( H ∗ 2 ) as its rows in direction. However, it cannot contain more than 2 different directions, otherwise for any

fixed optimal σ ( H ∗ 2 ) , there would be an extra direction which would contradict DNC3. Thus, W ∗ 2 contains exactly two linearly independent (orthogonal) directions. But then, in turn, only these two directions can be columns of σ ( H ∗ 2 ) . So σ ( H ∗ 2 ) is unique, possibly up to a switching of its columns x, y . However, only one permutation of these columns is optimal, because we know that the output of the matrix, H L +1 is orthogonal and only one of its permutations is better aligned with I 2 . Therefore σ ( H ∗ 2 ) is fully unique.

Invoking Lemma 9, we know that the whole range of matrices A 2 are such that σ ( A 2 ) = σ ( H ∗ 2 ) , while ∥ A 2 ∥ ∗ = ∥ H ∗ 2 ∥ ∗ . Namely, by Lemma 9, all the matrices A 2 which satisfy that condition are represented by having columns of the form x -ty, y -tx for t ∈ [0 , 1] . This can be derived as follows. According to the lemma and the description of σ ( H ∗ 2 ) , this matrix is a good candidate for the non-negative part of the optimal solution of the corresponding optimization problem, and so the optimal solution (in the case of L = 2 a solution for which the nuclear norms of it and its non-negative part are equal) whose non-negative part equals σ ( H ∗ 2 ) exists. The negative values, according to Lemma 9, must follow a strict, yet still non-unique pattern. The negative values of the (+ , 0) rows of σ ( H ∗ 2 ) must be arranged so that they are aligned with the positive values in the corresponding rows, thus being multiple of x . The negative values of the (0 , +) rows must be, on the other hand, aligned with y . The coefficients of homothety between x and the negative part in the (+ , 0) rows and between y and the negative part of the (0 , +) rows must be equal due to the condition on equal l 2 norms of the negative parts corresponding to (+ , 0) , (0 , +) rows and the fact that ∥ x ∥ = ∥ y ∥ . Finally, the coefficients of homothety must be non-positive (yielding non-negative t ) to retain the negativeness and no smaller than -1 , due to the final condition of Lemma 9, stating that the product of any two negative values of counterfactually signed rows must be at most as big as the product of the corresponding positive entries and due to the already established alignment of positive and negative parts. In this way, the only matrices A 2 which satisfy these conditions are represented by having columns of the form x -ty, y -tx for t ∈ [0 , 1] , as stated above (note that -ty is aligned with x on (+ , 0) rows of σ ( H ∗ 2 ) and similarly for -tx ).

̸

Let us denote by A t 2 a matrix of that form for a particular t . Now, the SVD of A t 2 is as follows: U t is composed of ( x -y ) / ( √ 2 ∥ x ∥ ) and ( x + y ) / ( √ 2 ∥ x ∥ ) , Σ t is a diagonal matrix with diagonal entries equal to ∥ x ∥ (1 + t ) and ∥ x ∥ (1 -t ) , and V t is composed of the vectors (1 / √ 2 , -1 / √ 2) T and (1 / √ 2 , 1 / √ 2) T . For t = 0 , this is the unique SVD up to ± multiplications. For t = 0 , A 0 2 is orthogonal and the SVD is non-unique. However, as we will see, this will not be relevant in the following argument. Invoking Lemma 8 we see that, for each t , the set of possible W t 1 which can be the optimal solution is S t = { c λ U t (Σ t ) 1 / 2 R T } for all possible orthogonal R of suitable dimensions. From this representation, we see that the non-uniqueness of SVD is 'wrapped' in the freedom of R T . If t = 0 , Σ 0 is diagonal and commutative and U 0 R T obviously cover all possible representations of U in the SVD of A 0 2 . If t > 0 , we can easily cover the freedom in signs by R too. Thus, we can assume a fixed U and V for all t ∈ [0 , 1] . The question is, whether for 0 ≤ t 1 < t 2 ≤ 1 , S t 1 ∩ S t 2 = ∅ . To find out, consider some R t 1 and R t 2 orthogonal. We want to know whether it is possible that c λ U t 1 (Σ t 1 ) 1 / 2 ( R t 1 ) T = c λ U t 2 (Σ t 2 ) 1 / 2 ( R t 2 ) T . First, this obviously does not hold if t 2 = 1 , because of the rank inequality. Assume t 2 < 1 and do the computation:

where in the first equivalence we used the established knowledge that we can use the unique representative U for all t . However, the diagonal elements of (Σ t 1 ) 1 / 2 (Σ t 2 ) -1 / 2 are √ (1 + t 1 ) / (1 + t 2 ) and √ (1 -t 1 ) / (1 -t 2 ) . The first one is strictly greater than 1. On the other hand, all the entries of ( R t 2 ) T ( R t 1 ) are inner products between columns of orthogonal matrices and thus at most 1 in absolute value. This proves S t 1 ∩ S t 2 = ∅ for 0 ≤ t 1 < t 2 ≤ 1 .

Thus, there is just a single t ∗ for which W ∗ 1 is optimal and, moreover, there is a single A ∗ 2 = A t ∗ 2 that can be the output of the first layer. Then, if we analyze the two systems of linear equations A ∗ 2 = W ∗ 1 X , where X is the d × 2 matrix variable (each system for each column of X ), we know that there is a unique minl 2 norm solution per system. The minl 2 norm solution must be the only optimal solution, because H 1 is regularized using the Frobenius norm. Thus, there is a single H 1

which can be the optimal solution. However, this is a contradiction with the assumption that we had a freedom while constructing ¯ H ∗ 1 from ˜ H ∗ 1 - a possibility of choosing different columns from some classes as the representative of that class. Therefore, the ˜ H ∗ 1 must have had a single feature vector per class in the first place, which contradicts the assumption that it is not DNC1 collapsed.

Now that we have the DNC1 collapse, we can obtain the DNC2-3 collapse easily by just again going from the full ˜ H ∗ to the reduced n = 1 version of it. Moreover, we know what are the forms of W ∗ 1 and H ∗ 1 - orthogonal transformations of the SVD of H ∗ 2 . However, it must be said that these optimal matrices do not need to be orthogonal themselves, due to the (Σ t ) 1 / 2 multiplication, which makes the matrix non-orthogonal (consider for instance just taking the identity matrix extended to rectangularity for R ). This concludes the proof of Theorem 3.

B Extension to multiple classes

Here, we comment shortly on why it is challenging to formulate an equivalent of Theorem 3 for K > 2 . First of all, it is already unclear how to carry out the proof of Lemma 5 for K > 2 and even whether it holds. However, the bigger issue lies in Lemma 9. Note that ∥ H 2 ∥ ∗ ≥ ∥ σ ( H 2 ) ∥ ∗ only holds for matrices with two rows or columns. A simple counterexample with 3 × 3 matrix goes as follows:

for which ∥ A ∥ ∗ = 3 . 464 and ∥ σ ( A ) ∥ ∗ = 3 . 494 . However, this does not mean that Theorem 3 does not hold for K > 2 . For instance, at orthogonal matrices, which are collapsed, one can show that this equation holds by computing the sub-gradient of the nuclear norm and using convexity. It is also hard to come up with counterexamples to the inequality for matrices which are close to orthogonal. Therefore, it is our belief that the cost of losing orthogonality in other loss terms makes the compensation of achieving slightly lower nuclear norm insufficient. However, to make this argument formal, one needs a sophisticated control over this quantity together with all the other results.

C Further numerical experiments on DUFM and ablations

In this section, we further investigate the effect of significant parameters on the training dynamics of DUFM. We focus on the effect of depth, regularization strength and width on the outcome of DUFM optimization. All the plots (except the one in Appendix C.5) are averaged over 10 runs with one standard deviation to both sides displayed as a confidence band. First, however, we supplement the numerical results from Section 4.2 with the promised plots of the losses.

C.1 Extra material to the numerical results in Section 4.2

In Figure 3, we plot the loss functions and theoretical optimum, together with the disentangled loss functions computed separately on the fit term and the regularization terms for H 1 , W l , with 1 ≤ l ≤ L. As already discussed in Section 4.2, the DUFM finds the global optimum perfectly and exhibits superb DNC.

C.2 Effect of depth on DUFM training

As already seen by showing 3- and 6-layered DUFM, the depth does not play a significant role on the training of DUFM. However, an implicit bias of GD can be observed for very deep DUFMs. To demonstrate this, we plot the 10-DUFM training results in Figure 4.

Though not well visible on the plots, the DNC2 is sometimes achieved also on H 1 , H 2 , which in the shallower DUFMs was an exception. This shows that with increasing depth, GD is more and more implicitly biased towards finding DNC2 representations for these feature matrices too.

loss

loss

100

100

total_loss

total_loss

10-1

fit_loss

10-1

fit_loss

theoretical_optimum

theoretical_optimum

10-2

10-2

10-3

10-3

10-4

10-4

10-5

10-5

10-6

10-6

0

20

40

60

80

100

0

20

40

60

80

100

step

1e3

step

1e3

loss

DNC1

100

total_loss

layer_1

fit_loss

10°

layer_2

10-1,

......

theoretical_optimum

layer_3

10-2

layer_4

10-2.

10-4

layer_5

10-6.

layer_6

10-3.

layer_7

10-8

layer_8

10-4.

10-10

layer_9

layer_10

10-5

10-12

10-14,

10-6.

0

20

40

60

80

100

0

20

40

60

80

100

step

1e3

step

1e3

DNC2

DNC3

layer_1

10°

layer_2

layer_3

layer_4

10-2

layer_1

101

layer_5

layer_2

layer_6

10-4

layer_3

layer_7

layer_4

layer_8

10-6

layer_5

layer_9

layer_6

layer_10

layer_7

layer_2_pre_relu

10-8

layer_8

layer_9

100

layer_10

0

20

40

60

80

100

0

20

40

60

80

100

1e3

1e3

step step

C.3 Effect of width on DUFM training

Unlike depth, width has a significant effect on the DUFM training, especially at small widths. The reason is that the model parameters often get stuck in sub-optimal configurations due to the difficult optimization landscape. To compare, we ran 4-DUFM training with the same hyperparameters except the width. We repeated the experiments for widths of 2, 5, 10, 20, 64, 1000. We plot all the results in Figure 5.

The results suggest that the risk of DUFM not reaching the global optimum and thus DNC drastically decreases with the width. Already for width 20, the DUFM did not fail to reach the global optimum once in ten runs. On the other hand, for width of only 2, the DUFM did not reach the global optimum in any run. Interestingly, the smaller widths yielded less expected behaviour. For instance for width 10, we encountered two runs both reaching a perfect optimum and DNC, of which one exhibited DNC2 also for H 1 , H 2 , while the other did not. Noticeably, for width 1000 we see higher variance in some DNC metrics. This is because some of these specific metrics converged to very low values for some runs and for some other stayed comparably higher. This effect appears to be an implicit bias of SGD.

C.4 The effect of weight decay on DUFM training

As expected for the MSE loss setup, the weight decay is the main driving force for the occurrence of DNC. Therefore, predictably, the weight decay has a direct influence on the speed of convergence towards both DNC and global optimality. This is clearly demonstrated in Figure 6, where we plot the training dynamics of 4-DUFM with width 64 and weight decays of sizes 0 . 0001 , 0 . 0005 , 0 . 0009 , 0 . 0013 , 0 . 0017 . By increasing the value of the weight decay, the loss as well as DNC metrics converge much faster. However, there is no qualitative change in the result of the training for different values of weight decay.

C.5 Interconnection of DNC and optimality of DUFM

As an interesting bonus, we show a single, cherry-picked run for the 4-DUFM with width 10 (although such runs occur quite often at this width). This training first got stuck at an almost-optimal solution, then after a long stay in the saddle point, it jumped to the optimum. The DNC metrics showed a clear decrease at exactly that point of the training, as shown in Figure 7. This run clearly shows how interconnected the deep neural collapse is with optimality.

D Pre-trained DUFM experiments on CIFAR-10

Here, we specify the training conditions from Section 4.2 in more detail, and add results of further runs to provide more robust evidence. We use three different architectures - ResNet20, DenseNet121 and VGG13. We only impose one change to these - the standard backbone is extended with one extra fully-connected layer with output dimension 64 to enable free features and reasonable size of the DUFM. The weights are initialized with He initialization [9] with scaling 1. The DUFMs in question are 3- and 5-layered of constant width 64. We train with weight decay 0 . 00001 and with learning rate 0 . 001 on 1000000 full GD steps. The backbone was then trained using stochastic gradient descent with batch size 128 , learning rate 0 . 00001 and with 10000 epochs on classes 0 and 1 of CIFAR-10. While in Figure 1 we plot results of a single training of the 3 -DUFM model followed by 7 independend trainings of independent initializations of the ResNet20, here we add two more independent trainings of the DUFM averaged over further 7 independent runs of ResNet20 training. The results are qualitatively the same as in the main body, which suggests that the ability of ResNet20 to make unconstrained features suitable for the emergence of DNC is robust w.r.t. the well-trained DUFM head. The results of these two runs are depicted in Figure 8.

As a side note, we remark that unlike end-to-end training or DUFM training, the ResNet20 training with pre-trained DUFM head is more prone to undesired behaviours. First of all, we need a rather small learning rate to prevent divergence and even with that, a minority of runs never jumps from 50% training accuracy. Therefore, in our results, we only report those runs, which converged to near 100% training accuracy.

loss

DNC1

1016

DNC2

DNC3

100

layer_1

layer_1

100

layer_1

10-4

102

layer_3

layer_2

1012

layer_2

layer_3

10-2

layer_2

layer_3

10-8

101

Uu

layer_4

108

layer_2_pre_relu

layer_4

10-6

10-4

layer_4

10-12

total_loss

fit_loss

100

104

10-8

theoretical_optimum

10°

10-10

10-16

0

20

40

step

60

80

1e3

100

10-1

0

20

40

step

60

80

100

1e3

0

20

step

40

60

1e3

0

20

40

step

60

80

100

1e3

100

loss

102

DNC1

102

layer_1

DNC2

100

DNC3

10-1

101

101.5

layer_2

10-1

10-2

10-1

100

101

layer_3

layer_4

10-2

10-3

10-2

10-3

layer_1

100.5

layer_2_pre_relu

10-3

layer_1

10-4

total_loss

10-4

layer_2

10°

10-4

layer_2

10-5

fit_loss

10-5

layer_3

10-5

layer_3

theoretical_optimum

10-6

layer_4

layer_4

10-6

10-7

10-6

0

20

40

step

60

80

100

1e3

0

20

40

step

60

80

100

1e3

0

20

40

step

60

80

100

1e3

20

40

step

60

80

1e3

100

10-2

10-1

10-3

10°

loss

10-2

10-6

10-4

100

DNC1

layer_3

layer_2

layer_1

layer_4

101

DNC2

layer_2

layer_3

layer_4

layer_2_pre_relu

layer_1

10-2

10-3

10-1

10-4

100

DNC3

10-4

total_loss

10-10

10-8

10-5

layer_2

layer_1

10-5

fit_loss

10-12

10-6

layer_3

10-6

theoretical_optimum

10-14

100

10-7

layer_4

10-8

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

step

1e3

step

1e3

step

1e3

step

1e3

loss

DNC1

DNC2

DNC3

10-3

10-1

10-2

100

theoretical_optimum

total_loss

fit_loss

10-2

10-4

10-6

100

layer_1

layer_2

layer_4

layer_3

101

layer_3

layer_4

layer_1

layer_2

layer_2_pre_relu

10-1

10-2

10-3

10-4

10-5

100

10-4

10-10

10-8

10-7

10-6

layer_1

layer_2

10-5

10-12

10-8

layer_3

10-14

10°

layer_4

10-6

0

20

40

60

80

100

0

20

40

60

80

100

0

20

4060

80

100

0

20

40

60

80

100

step

1e3

step

1e3

step

1e3

step

1e3

10-2

10-3

10-1

100

total_loss

fit_loss

theoretical_optimum

loss

10-2

10-4

10-6

100

DNC1

layer_2

layer_3

layer_1

layer_4

101

DNC2

layer_3

layer_1

layer_2

layer_4

layer_2_pre_relu

10-2

10-3

10-4

10-2

10-5

100

DNC3

10-8

10-6

layer_1

10-4

10-10

10-7

layer_2

10-5

10-12

10-8

layer_3

10-14

10°

layer_4

10-6

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

step

1e3

step

1e3

step

1e3

step

1e3

loss

DNC1

DNC2

DNC3

100

total_loss

fit_loss

10-2

10°

layer_2

layer_1

101

layer_2

layer_1

10-1

100

10-1

10-2

theoretical_optimum

10-4

10-6

layer_4

layer_3

layer_4

layer_3

10-2

10-3

10-10

10-8

layer_2_pre_relu

10-4

10-3

layer_1

10-4

layer_2

10-12

10-5

layer_3

10-5

10-14

100

10-6

layer_4

10-6

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

step

1e3

step

1e3

step

1e3

step

1e3

10-2

10-1

100

fit_loss

total_loss

theoretical_optimum

loss

10-2

100

DNC1

101

DNC2

layer_2

layer_1

layer_3

10-0.5

100

DNC3

layer_1

layer_2

layer_3

10-3

10-4

layer_4

layer_2_pre_relu

10-1

layer_4

10-4

10-6

layer_1

10-5

10-8

layer_2

101.5

10-6

10-7

10-10

layer_3

layer_4

10°

10-2

0

20

40

60

80

100

。

20

40

09

80

100

0

20

4060

80

100

0

20

40

60

80

100

step

1e3

step

1e3

step

1e3

step

1e3

10-3

10-1

10-2

100

total_loss

fit_loss

theoretical_optimum

loss

10-2

10-6

10-4

100

DNC1

layer_3

layer_2

layer_1

layer_4

101

DNC2

layer_2

layer_1

layer_3

layer_4

layer_2_pre_relu

10-4

10-2

100

DNC3

10-8

10-6

layer_1

10-4

10-10

layer_2

10-5

10-12

10-8

layer_3

10-6

0

20

40

step

60

80

100

1e3

10-14

20

40

step

60

80

100

1e3

10°

0

20

40

step

60

80

100

1e3

0

layer_4

20

40

step

60

80

100

1e3

loss

DNC1

DNC2

DNC3

10-1

10-2↓

10°

fit_loss

total_loss

theoretical_optimum

10-2

10-4

10-6

100

layer_1

layer_2

layer_3

layer_4

101

layer_1

layer_3

layer_4

layer_2

layer_2_pre_relu

10-2

10-4

100

10-3

10-8

10-6

layer_1

10-4

10-10

10-12

10-8

layer_2

layer_3

10-5

10-14

10°

layer_4

0

20

40

60

80

100

0

20

40

60

80

100

0

20

4060

80

100

0

20

40

60

80

100

step

1e3

step

1e3

step

1e3

step

1e3

loss

DNC1

DNC2

DNC3

10-1

10-2

10°

total_loss

fit_loss

theoretical_optimum

10-2

10-4

10-6

100

layer_1

layer_2

layer_4

layer_3

101

layer_2

layer_4

layer_1

layer_3

layer_2_pre_relu

10-2

10-4

FoOt

10-3

10-8

10-6

layer_1

10-10

layer_2

10-4

10-12

10-8

10-5

0

20

40

60

80

100

10-14

0

20

60

80

100

10°

0

20

40

60

80

100

0

layer_4

20

40

60

80

100

step

1e3

step

1e3

step

1e3

step

1e3

loss

DNC1

DNC2

DNC3

10-2

10-1

100

total_loss

theoretical_optimum

10-4

10-6

10-2

100

layer_2

layer_1

layer_3

layer_4

101

layer_2

layer_3

layer_1

layer_4

layer_2_pre_relu

10-2

10-4

100

10-3

10-8

10-6

layer_1

10-10]

layer_2

10-4

10-12

10-8

layer_3

10-5

10-14

10°

layer_4

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

1e3

1e3

1e3

1e3

step step

step step

loss

100

DNC1

total_loss

layer_1

10-1,

fit_loss

10°

layer_2

theoretical_optimum

10-2

layer_3

10-2

10-4,

layer_4

10-3.

10-6.

10-8

10-4

10-10

10-5

10-12.

10-14

10-6.

0

20

40

60

80

100

0

20

40

60

80

100

step

1e3

step

1e3

DNC2

DNC3

layer_1

10°

layer_1

101

layer_2

layer_2

layer_3

layer_3

layer_4

10-1,

layer_4

layer_2_pre_relu

10-2

10°

10-3

0

20

40

60

80

100

0

20

40

60

80

100

step

1e3

step

1e3

loss

DNC1

DNC2

100

DNC3

100

total_loss

fit_loss

101

layer_1

layer_2

104

layer_2

layer_1

reg_loss_Ho

10°

layer_3

103

layer_3

10-1

10-1

10-1

102

layer_2_pre_relu

10-2↓

10-2

10-2

10-3

10-3

10-4

100

101

10-3

layer_1

layer_2

layer_3

10-4

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

1e3

1e3

step

1e3

step

1e3

loss

DNC1

DNC2

100

DNC3

100

total_loss

fit_loss

101

layer_1

layer_2

104

layer_2

layer_1

reg_loss_Ho

100

layer_3

103

layer_3

10-1

10-1

10-1↓

layer_2_pre_relu

102

10-2

10-2

10-2

10-3

101

10-3

layer_1

10-3

10-4

10°

layer_2

layer_3

10-4

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

1e3

1e3

step

1e3

1e3

step

loss

DNC1

DNC2

DNC3

100

layer_1

layer_1

100

layer_1

10-1

102

layer_2

layer_2

layer_2

101

layer_3

layer_3

10-1

layer_3

10-2

10°

layer_2_pre_relu

10-3

total_loss

fit_loss

10-2

10-1

10-2

10-4

10-3

10-3

10-5

10-4

101

105

100

10-4

10-6

0

1

2

3

4

5

10-6

0

1

2

3

4

5

0

1

2

3

4

5

0

1

2

3

4

5

le3

1e3

1e3

1e3

100

loss

total_loss

DNC1

layer_1

DNC2

layer_1

10°

DNC3

fit_loss

102

layer_2

layer_2

10-1

101.

layer_3

layer_3

10-1

layer_2_pre_relu

100

10-2

10-1

101

10-2

layer_1

10-2

10-3

layer_2

10-3

10-3

10°

layer_3

0.0

0.5

1.0

1.5

0.0

0.5

1.0

1.5

0.0

0.5

1.0

1.5

0.0

0.5

1.0

1.5

1e3

1e3

1e3

1e3

100

loss

total_loss

102

DNC1

layer_1

101

DNC2

100

DNC3

layer_1

fit_loss

101

layer_2

layer_2

10-1

layer_3

layer_3

100

layer_4

10-1

layer_4

layer_5

layer_5

102

10-1

100

layer_1

layer_2

10-2

layer_3

10-2

10-3

layer_4

10-3

layer_5

layer_2_pre_relu

10-3

10-4

0

1

2

3

4

5

0

1

2

3

4

5

0

1

2

3

4

5

0

1

2

3

4

5

1e3

1e3

1e3

1e3

We further present the results of the same experiment but with 5 -DUFM attached to ResNet20 as well as 3 -DUFM attached to DenseNet121 and VGG13, to even further validate the robustness of our results. These results are presented in Figures 9 and 10, and they are in full accordance with the ResNet20 experiments.

Peer reviews:
comments: SUMMARY:
The paper introduces an extension to the standard unconstrained features model in neural collapse theory by incorporating multiple layers above the unconstrained features. The study demonstrates that, under certain conditions and at optimality, all of the top layers exhibit all properties of neural collapse.

STRENGTHS:
-- The paper exhibits excellent writing and effectively surveys the existing literature.
-- The paper contributes to a promising direction in the theory of neural collapse by presenting a significant extension, and it conducts a non-trivial theoretical analysis that offers compelling evidence for the occurrence of neural collapse at intermediate layers. This extension aligns with previous empirical findings, verifying the theoretical predictions made in earlier studies.

WEAKNESSES:
-- Like other papers in the field, the rationale behind asserting that SGD converges to a solution satisfying neural collapse solely based on optimality is not entirely clear. It would greatly benefit the paper if the authors could provide further clarification and expand upon this aspect.
-- While the paper focuses on the binary classification case, which is acceptable, it would be valuable to discuss the feasibility and potential challenges of extending the analysis to a multiclass setting.
-- I believe that expanding on the implications and interpretations of the results would significantly improve the paper. It would be beneficial for the authors to discuss how these findings relate to aspects like generalization and what implications they have for our understanding of deep learning in general.

QUESTIONS:
See weaknesses.

LIMITATIONS:
There are no negative societal issues with this paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studied Deep Neural Collapse (DNC) which extends the structure of Neural Collapse of the last layer to multi-layers in deep neural networks. Theoretically, this paper generalizes the established unconstrained features model to the deep unconstrained features model (DUFM) of multi-layer non-linear models for binary classification. Based on the deep unconstrained features model, they show the global optimum exhibits several extended properties. Empirical experiments demonstrate that the results of the optimized deep unconstrained features model and trained networks agree with the theory.\n\nSTRENGTHS:\n1. The paper generalizes the Neural Collapse to Deep Neural Collapse which describes the structure of the feature other than the last layer. The proof based on the deep unconstrained features model for binary classification is clear and easy to follow. It gives a framework to find and prove the feature structure of other layers in DUFM.\n\n2. The extensive experiments and numerical results shows that the theory successfully predicts the behavior of the deep unconstrained features model(DUFM) which exhibits proposed DNC properties. \n\n3. The paper is well written, it is smooth to follow the whole logic and the idea of the method.\n\nWEAKNESSES:\n1. Though the model and proof consider the Deep Neural Collapse which works for not only features of the last layer, the model and the proof are restrictive which works for binary classification and no bias case. Many proofs and derivations rely on binary classification like Lemma 5 - Lemma 7 which is the key to giving the conditional optimum of previous layers. \n\n2. In addition, the DUFM adds a regularization term for the first layer feature which is not fit with the real training of neural networks to some extent. \n\n3. Compared to the numerical results of DUFM, the experiments for real training of neural networks are limited to the 3-layer Resnet20. If there are more results for the training of other neural networks it will make the experiments more convincible.\n\n4. There are some typos in the proof, like L517 in the appendix where I think the $\\sigma_1^2$ should be $s_{L, i}^2$.\n\nQUESTIONS:\n1. Most previous research[1,2, 3] focus on models with bias while in [3] they also study the unconstrained features model without bias and show the features exhibit OF rather than STF under UFM with bias. So I am just wondering if it is possible to extend the DUFM to the bias case and show the STF structure in DNC2?\n\n2. There are some experiments results for 3-layer Resnet-20 in the paper which also show DNC, but in [4] the results seem to show that the Neural Collapse process gradually across the layers of the neural network like the measure of NC1 decreases exponentially. I think this may be caused by the gap between the DUFM and real trained networks. Are there more experiments on real trained networks and whether they exhibit DNCs? If so, why is there such a difference between the results in this paper and the results in [4]. If not, is that mean there is a gap and needs other explanations for such a phenomenon? \n\n3. All experiments based on DUFM in this paper add the regularization term of the first layer feature. By the derivation of the DUFM, it seems that possible to add a regularization term on the intermediate layer feature, for example, add $||H_5||_2$ and optimize $W_l, l \\ge 5$, are there any experiment results can demonstrate in this case how the DNC performs?\n\n\nReference:\n[1] Zhu, Z., Ding, T., Zhou, J., Li, X., You, C., Sulam, J., & Qu, Q. (2021). A geometric analysis of neural collapse with unconstrained features. Advances in Neural Information Processing Systems, 34, 29820-29834.\n\n[2] Fang, C., He, H., Long, Q., & Su, W. J. (2021). Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences, 118(43), e2103091118.\n\n[3] Tirer, T., & Bruna, J. (2022, June). Extended unconstrained features model for exploring deep neural collapse. In International Conference on Machine Learning (pp. 21478-21505). PMLR.\n\n[4]He, H., & Su, W. J. (2022). A law of data separation in deep learning. arXiv preprint arXiv:2210.17020.\n\nLIMITATIONS:\nThe framework and the proof relies on the binary classification and no bias case which make the DUFM not so correlated to the real trained networks.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a novel approach to investigating deep neural collapse in deep neural networks by presenting a generalization of the analytical framework for neural collapse (NC) to multiple non-linear layers. The paper introduces the deep unconstrained features model to demonstrate that the unique global optimum of the MSE loss for binary classification exhibits all the properties typical of deep neural collapse. However, the paper only considers the MSE loss without a bias term for binary classification, which may limit the applicability of deep neural collapse for other losses. Specifically, the theoretical conclusions may only be provably optimal in some cases.\n\nSTRENGTHS:\n- The paper presents a  generalization of the unconstrained features model  to multiple non-linear layers, filling a gap in the existing literature.\n- The paper contributes to a better understanding of deep neural collapse in binary classifcation deep neural networks, which has important implications for the design and optimization of deep learning models.\n\nWEAKNESSES:\n- The theoretical analysis only involves binary classification, and it is unclear whether the proposed approach can be extended to multi-class classification problems, especially for a large number of classes that would change the final form (DNC2) of the resulting deep neural collapse.\n- The paper only considers the MSE loss without a (last-layer) bias term, which may limit the applicability of deep neural collapse for other losses. Specifically, benifiting from the MSE loss without the bias term, the unique global solution can exhibit close-formed results, which may further allow the derivation of the proof to be greatly simplified.\n-  while neural collapse induces a simplex equiangular tight frame [1], which may not hold for other loss functions  or if a bias term is included [2]. \n- The paper lacks extensive experiments to validate deep neural collapse. Additional experiments on a variety of datasets (especially for a large number of classes) and deep neural networks would help to better understand and confirm the prevalence of deep neural collapse.\n- The paper could benefit from a more thorough discussion of the practical implications of the proposed approach for the design and optimization of deep learning models.\n\n[1] Vardan Papyan, X. Y. Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. In *Proceedings of the National Academy of Sciences (PNAS)*, volume 117, 2020.\n\n[2] Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In *International Conference on Machine Learning (ICML)*, 2022.\n\nQUESTIONS:\nsee weaknesses.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the concept of deep neural collapse (DNC), extending the understanding of neural collapse to earlier layers of deep neural networks. It proposes the deep unconstrained features model (DUFM) as a theoretical framework to analyze DNC. The authors demonstrate that for binary classification, DNC is a globally optimal solution in the DUFM, providing rigorous evidence of its occurrence in multiple layers. Numerical experiments confirm that gradient descent efficiently finds solutions in line with the theory, exhibiting neural collapse in multiple layers. The paper concludes by highlighting open problems for future research, such as generalizing the analysis to multiple classes, exploring gradient dynamics in the DUFM, and understanding the impact of biases and loss functions on the collapsed solutions.\n\nSTRENGTHS:\n$\\bf (1)\\ \\textbf{Strong Technical Analysis}:$ A key strength of this paper is its strong and rigorous technical analysis. The authors provide a solid theoretical framework, specifically the deep unconstrained features model (DUFM), to analyze and understand the phenomenon of deep neural collapse (DNC). Theorems and empirical evidence are presented to support their findings, enhancing the credibility and reliability of the results.\n\n$\\bf (2)\\ \\textbf{Extension to Multiple Non-linear Layers:}$  This paper addresses a crucial gap in the existing literature by extending the analysis of neural collapse beyond the last layer to multiple non-linear layers. By generalizing the unconstrained features model, the authors provide insights into deep neural collapse (DNC) occurring in deeper layers of neural networks. This extension contributes to a more comprehensive understanding of the behavior of deep networks during training.\n\n$\\bf (3)\\ \\textbf{Empirical Validation:}$ The paper provides empirical evidence to validate the theoretical framework. The authors demonstrate that by optimizing deep unconstrained feature models using gradient descent, the resulting solutions align well with the theoretical predictions. This empirical validation reinforces the credibility of the proposed modeling principles and their applicability in real-world scenarios.\n\n$\\bf (4)\\ \\textbf{Relevance and Contribution to the Field:}$ The paper addresses a significant topic of interest in the field of deep learning by focusing on the phenomenon of neural collapse in deep neural networks. By providing theoretical insights and empirical evidence, the paper contributes to the existing body of knowledge and expands our understanding of the behavior of deep networks during training.\n\nWEAKNESSES:\n$\\bf (1)\\ \\textbf{Theorem 3 cannot explain the progressive neural collapse:}$ In Ref [8], the authors observe that each layer roughly improves a certain measure of data separation by an equal multiplicative factor. A similar observation is depicted in Figure 1 of the current paper. However, Theorem 3 contradicts this observation by demonstrating that for all layers $l \\ge 2$, the within-class variability of features becomes zero, and the features form an orthogonal matrix. This discrepancy suggests that there may be additional factors or dynamics at play in the behavior of the layers beyond what is captured by the initial observation. Further investigation is necessary to reconcile these contrasting findings and gain a more comprehensive understanding of the behavior and properties of features in deep neural networks.\n\n$\\bf (2) \\textbf{Role of activations:}$ One weakness of the paper is the limited discussion or exploration of the specific role and impact of activations in the context of deep neural collapse (DNC). While the paper acknowledges that DNC1 and DNC2 hold for layer l ≥ 2 before and after activations, it does not provide an in-depth analysis or explanation of how the activations influence or interact with DNC.\n\n$\\bf (3) \\textbf{Limited Scope:}$  The paper primarily focuses on deep neural collapse (DNC) in the context of binary classification. While this specific focus allows for in-depth analysis, it may limit the generalizability of the findings to other problem domains or network architectures. Further exploration of DNC in more diverse settings could strengthen the paper's applicability.\n\n$\\bf Typos:$ (1) At line 122, it should be $W_L \\in \\mathbb{R}^{K\\times d_L}$.\n\nQUESTIONS:\n$\\bf 1.$ The results in this paper are based on the L-layer deep unconstrained features model. This omits all the structures in the data input. A natural question is whether the analysis can be extended to the case where the data input has some special structures, such as the case where the data is whitened, i.e., assuming $H_1H_1^T=I$ in Definition 1.  \n\n$\\bf 2.$ What is the main technical difficulty in analyzing multiple-class cases?\n\n$\\bf 3.$ What is the role of the activations in deep neural networks? Why do DNC1 and DNC2 hold for layer l ≥ 2 before and after activations?\n\nLIMITATIONS:\nNot applicable", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}