Wide Neural Networks Trained with Weight Decay Provably Exhibit Neural Collapse
Authors: Arthur Jacot
Peter Súkeník
Zihan Wang
Marco Mondelli
Venue: ICLR 2025 Oral (2025)
Keywords: neural collapse
gradient descent training
weight decay
balancedness
Primary area: optimization
TL;DR: We consider deep neural networks with at least two final linear layers, and we show that neural collapse provably holds in the end-to-end training of the model with weight decay

Abstract: Deep neural networks (DNNs) at convergence consistently represent the training data in the last layer via a geometric structure referred to as neural collapse. This empirical evidence has spurred a line of theoretical research aimed at proving the emergence of neural collapse, mostly focusing on the unconstrained features model. Here, the features of the penultimate layer are free variables, which makes the model data-agnostic and puts into question its ability to capture DNN training. Our work addresses the issue, moving away from unconstrained features and studying DNNs that end with at least two linear layers. We first prove generic guarantees on neural collapse that assume \emph{(i)} low training error and balancedness of 
linear layers (for within-class variability collapse), and \emph{(ii)} bounded conditioning of the features before the linear part (for orthogonality of class-means, and their alignment with weight matrices). The balancedness refers to the fact that $W_{\ell+1}^\top W_{\ell+1}\approx W_\ell W_\ell ^\top$ for any pair of
consecutive weight matrices
of the linear part, and the bounded conditioning requires a well-behaved ratio between largest and smallest non-zero singular values of the features. We then show that such assumptions hold for gradient descent training with weight decay: \emph{(i)} for networks with a wide first layer, we prove low training error and balancedness, and \emph{(ii)} for solutions that are either nearly optimal or stable under large learning rates, we additionally prove the bounded conditioning. Taken together, our results are the first to show neural collapse in the end-to-end training of DNNs.

Published as a conference paper at ICLR 2025
WIDE NEURAL NETWORKS TRAINED WITH WEIGHT DECAY PROVABLY EXHIBIT NEURAL COLLAPSE
Arthur Jacot, Peter Súkeník} Zihan Wangi and Marco Mondelli
ABSTRACT
Deep neural networks (DNNs) at convergence consistently represent the training data in the last layer via geometric structure referred to neural collapse. This empirical evidence has spurred a line of theoretical research aimed at proving the emergence of neural collapse, mostly focusing on the unconstrained features model. Here, the features of the penultimate layer are free variables, which makes the model data-agnostic and puts into question its ability to capture DNN train- ing. Our work addresses the issue, moving away from unconstrained features and studying DNNs that end with least two linear layers. We first prove generic guarantees on neural collapse that assume (i) low training error and balancedness of linear layers (for within-class variability collapse), and (ii) bounded condition- ing of the features before the linear part (for orthogonality of class-means, and their alignment with weight matrices). The balancedness refers to the fact that We+1 W.W, T for any pair of consecutive weight matrices of the linear part, and the bounded conditioning requires a well-behaved ratio between largest and smallest non-zero singular values of the features. We then show that such as- sumptions hold for gradient descent training with weight decay: (i) for networks with wide first layer, we prove low training error and balancedness, and (ii) for solutions that are either nearly optimal or stable under large learning rates, we ad- ditionally prove the bounded conditioning. Taken together, our results the first to show neural collapse in the end-to-end training of DNNs.
INTRODUCTION
Among the many possible interpolators that deep neural network (DNN) can find, Papyan et al. (2200) showed a strong bias of gradient-based training towards representations with a highly sym- metric structure in the penultimate layer, which was dubbed neural collapse (NC). In particular, the feature vectors of the training data in the penultimate layer collapse to a single vector per class (NC1); these vectors form orthogonal or simplex equiangular tight frames (NC2), and they are aligned with the last layer's 's row weight vectors (NC3). The question of why and how neural col- lapse emerges has been considered by a popular line of research, see e.g. [u & Steinerberger (2022); E Wojtowytsch (202) and the discussion in Section 2 Many of these works focus on a simplified mathematical framework: the unconstrained features model (UFM) (Mixon et 2020 Han al., 2022 Zhou et al. 2022a), corresponding to the joint optimization over the last layer's weights and the penultimate layer's feature representations, which are treated as free variables. To account for the existence of the training data and of all the layers before the penultimate (i.c., the backbone of the network), some form of regularization on the free features is usually added. A number of papers has proved the optimality of NC in this model (Lu & Steinerberger 2022 E & Wojtowytsch 2022). its emergence with gradient-based methods (Mixon et 22,00 Han et al. (2022) and a benign loss landscape (Zhou et al.) 2022a] Zhuu et al., (2021). However, the major drawback of the UFM lies in its data-agnostic nature: it only acknowledges the presence of training data and backbone through simple form of regularization (e.g.. Frobenius norm sphere constraint), which is far from being equivalent to end-to-ent training. Moving beyond UFM, existing results are either only applicable to
rather shallow networks (of at most three layers) (Kothapalli & Tirer 2024 HHng & Ling 2024) or
hold under strong assumptions, such as symmetric quasi-interpolltion (Xu et al. 2023 ( Rangamani & Banburski-Fahey 2022). block-structured empirical NTK throughout training (Seleznova et al.)
2023), or geodesic structure of the features across all layers (Wang et al. 2024).
"Courant Institute of Mathematical Sciences, NYU. Email: arthur. @nyu.ceeey.edd
I Institute of Science and Technology Austria. Email: peter. wukenikkista. ac. 1 Courant Institute of Mathematical Sciences, NYU. Email: 2x3503@nyu.eds edu
6 Institute of Science and Technology Austria. Email: marco
marco. mondellieeeeee
ac.at

---

Published as a conference paper at ICLR 2025
In this paper, provide the first end-to-end proof of within-class variability collapse (NC1) for class of networks that end with at least two linear layers. Furthermore, we give rather weak sufficient conditions either near-optimality or stability under large learning rates for solutions to exhibit the orthogonality of class-means (NC2) and the alignment of class means with the last weight matrix (NC3). More precisely, our contributions can be summarized as follows:
First, we show that within-class variability collapse (i.e., NC1) occurs, as long as the training error is low and the linear layers are approximately balanced, i.e., |W We+1 - W,W, is small, where We+1. Wy are two consecutive weight matrices of the linear head. If, additionally, the conditioning of the linear head mapping is bounded, we bound the conditioning of the matrix of class means in the last layer, well that of the last weight matrix. This implies that, as the number of linear layers grows, class means become orthogonal and, furthermore, they align with the last layer's row vectors, which proves both NC2 and NC3.
Next, we show that the sufficient conditions above for NCI are satisfied by class of deep net- works with a wide first layer, smooth activations and pyramidal topology, after gradient training with weight decay. This provides the first guarantee of the emergence of NC1 for a deep network trained end-to-end via gradient descent.
We present two sufficient conditions under which the linear head is well-conditioned, hence NC1, NC2 and NC3 hold: either the network approaches a global optimum of the fz-regularized square loss, or it nearly interpolates the data while being stable under large learning rates.
Our numerical experiments on various architectures (fully connected, ResNet) and datasets (MNIST, CIFAR) confirm the insights coming from the theory: (i) NC2 is more prominent as the depth of the linear head increases, and (ii) the final linear layers are balanced at con- vergence, Furthermore, we show that, as the non-linear part of the network gets deeper, the non-negative layers become less non-linear and more balanced.
RELATED WORK
Neural collapse. Since its introduction by Papyan et al. (2020), neural collapse has been inten-
sively studied, both from theoretical and practical viewpoint. Practitioners use the NC for a num-
ber of applications, including transfer learning. OOD detection and generalization bounds (Galanti
al. 2022aa] Haas et al. 2022 Ben-Shaul & Dekel 2022 Li et al. 2023a bl Zhang al., 2024). On the theoretical front, the most widely adopted framework to study the emergence of NC is the un-
constrained features model (UFM) (Mixon al.
2020 Fang et al.] (02T). Under the UFM, the NC
has been proved to be optimal with cross-entropy loss (E & Wojtowytsch] 2022 Lu & Steinerberger
2022 Kunin et al. 2022), MSE loss (Zhou et al. 2022a) and other losses (Zhou al. [0022b). Op-
timality guarantees for the generalization of NC to the class-imbalanced setting have been provided
by Fang et al.] (202T): [Thampoulidis et al.] (2022); Hong & Ling] (2023); Dang al. (2024). Be- sides global optimality, a benign loss landscape around NC solutions has been proved in Zhu et al.)
(202T);Ji et al. (2022); ; Zhou et al. (2022a), and the emergence of NC with gradient-based optimiza- tion under UFM has been studied by |Mixon et al. (2020); Han et al. (2022); Ji al. (2022); Wang cet al. (2022). diang et al. (2023) extend the analysis to large number of classes, Kothapalli et al. (2023) to graph neural networks, Tirer et al. (2023) generalize UFM with perturbation to account for its imperfections, while (Andriopollos et al. (2024) generalize UFM and NC to regression prob- lems. A deep neural collapse theoretically analyzed with deep UFM for the linear case in Dang ct al. (2023): Garrod & Keating (2024), for linear case with data-specific assumptions in Wang et al. (023). ), for the non-linear case with two layers in Tirer & Bruna (2022), and for the deep non-linear case in Sukenik et al. (2023 2024), where, notably, the latter work provides the first negative result on deep NC with deep UFM
A line of recent work aims at circumventing the data-agnostic nature of UFM, showing the emer- gence of neural collapse in settings closer to practice. Specifically, [Seleznova et al. (2023) assume block structure in the empirical NTK matrix. Kernels are used by Kothapalli & Tirer (2024) to
analyze NC in wide two-layer networks, showing mostly negative results in the NTK regime. [ea-) glehole et al. (2024) prove the emergence of deep neural collapse using kernel-based layer-wise wise training, and they show that NC is an optimal solution of adaptive kernel ridge regression in an over-parametrized regime. Sufficient (but rather strong) conditions for the emergence of NC beyond UFM are provided by Pan & Cao (2023). Hong & Ling (2024) analyze two and three layer net- works end-to-end, but only provide conditions under which the UFM optimal solutions are feasible.
2

---

Published as a conference paper at ICLR 2025
Wang et al. (2024) consider the NC formation in residual networks, and the model is similar to the perturbed UFM in Tirer et al. (2023). However, the results crucially assume that the features lie on geodesic in euclidean space, and proving that this is the case after training the ResNet is an open problem. |Rangamani & Banburski-Fahey (2022); ; Xu et al. (2023) focus on homogeneous networks trained via gradient-based methods, making the strong assumption of symmetric quasi-interpolation. Then, Rangamani & Banburski-Faiie (202) do not prove this assumption, and the argument of of ct al. )023) requires a regularization different from the one used in practice as well as interpolators with a given norm (whose existence is an open question). We also note that the property does not hold in practice exactly, which points to the need of a perturbation analysis.
Implicit bias. Our approach leverages the (approximate) balancedness of the weights, which plays a central role in the analysis of the training dynamics of linear networks (Arora et al.) 20188) and leads bias towards low-rank matrices (Arora et al. 2019 Tu et al. (024). In the presence of weight decay, the low-rank bias can be made even more explicit (Dai et al. (02T), and it is rein- forced by stochastic gradient descent <Wang & Jacot (024). Moving towards nonlinear models, shallow networks with weight decay exhibit low-rank bias, as described by the variational norm (Bach, 2017), or Barron norm (E et al. 2019), and this bias provably emerges under (modified) GD dynamics (Abbe et al. 2022 Bietti et al. 2022: Lee et al. 2024). However, a single hidden layer appears to be insufficient to exhibit NC for most datasets. Moving towards deep nonlinear models, networks with weight decay are also known to exhibit low-rank bias (Galanti al. 2022b; Jacot et al, 2022), which can be related to neural collapse (Zangrando et al., 2024). The dimensionality varies between layers, exhibiting bottleneck structure (Ja-) cot [2023a b. Wen & Jacol 2024). We remark that existing results apply to the global minima of the (g-regularized loss. In contrast, our paper provides rather general sufficient conditions that are then provably satisfied by GD, thus showing that the training algorithm is responsible for neural collapse.
and rank of the weights/representatii
3 BALANCEDNESS AND INTERPOLATION IMPLY NEURAL COLLAPSE
Notation and problem setup. Given matrix A of rank k we denote by As: its i-th row, by A. its
i-th column, by 81 (A) 2 ... 2 8 (A) its singular values in non-increasing order, and by K( () the
ratio (1(4) between its largest and smallest non-zero singular values. We denote by || F || lop
[A] and miin (A) its Frobenius norm, its operator norm and its smallest singular value, respectively.
We consider neural network with L, non-linear layers with activation function σ R - R fol- lowed by L2 linear layers. Let L == Li + L2 be the total number of layers, We E f the weight matrix at layer E, X E RdxN the training data and Y E K the labels (for consistency, we set no = d and n = K), where the output dimension K corresponds to the number of classes and N is the number of samples. We consider one-hot encoding, i.e., the rows of Y are elements of the canonical basis. Let Z₂ E TR
N be the output of layer f, given by
X
{= 8=00 0,
Ze
o (W,
E [.
W
EE (L1 + +,...,, L2).
where the activation function a is applied componentwise and, given an integer n, we use the shorthand [n] = [1....,m). We write 0 for the vector obtained by concatenating all param- eters (Wiheet) and Wmet for the partial products of the weight matrices Wm. ... W,, so that Zm = W Wme&+122 Ze for all m 2 É E (LI I..... L1 + L2). We index individual samples in a feature matrix Z as Zes, meaning the i-th sample of the c-th class, sometimes adding upper-inder index to denote the layer or matrix to which the sample belongs. Let me denote the mean of all samples from class e and μa the global mean. Let 7 be the matrix of class-means stacked into columns. The NCI metric on a feature matrix Z is given by tr(Ew)
T((n)) where Σw = * Eng(Ees - Me)( - Me0) (Me (M) -Mo) The NC2 metric on a feature matrix Z is defined as k< k(2), i... the conditioning number of the class-mean matrix of Z. The NC3 metrie on feature matrix Z and a weight matrix W is defined as N Ee, cos(zes, Wa), i.e., the average cosine similarity of features and weight vectors corresponding to the features' class.
and Σg = K IK
At this point, we state result giving set of sufficient conditions for NC1, NC2 and NC3. Theorem 3.1. If the network satisfies approximate interpolation, i.e., 17,-- Yll S E
3

---

Published as a conference paper at ICLR 2025
approximate balancedness, i.e., We+1 - WoW, 11 I c2 (L1+ 1..... 1},
bounded representations and weights, i.e, 1|Zr-2llop 11 1loeelllllWWelloy S rr 1,...,L}:
{Li
then i S min (xx(Y),
[K- -NNN 4K
NCI(21.11) = O ((E1 + v=2)²)-
(2)
If additionally assume that the linear part of the network is not ill-conditioned, i.e., R 1al+4) 1) S C3 then
<(LL) S c3 +((2),
NC2(ZL-1)
S c
1
+0( (V2),
(4)
Cg
+ 6z
a1 (=2 V -1) -1)
(C3 Wt) 2
(5)
2
interested in the regime which 61 is small. In words, ② shows that, when 62 R (i.e., the network approximately interpolates the data in a balanced way), the within- class variability (which captures NCI) vanishes. addition the depth of the linear part of the network grows, the RHS of (3) approaches 1, i.e., the last weight matrix Wr is close to orthogonal. This implies that (i) Z is also close to orthogonal (which captures NC2), and (ii) the weights the last layer align with L-1 (which captures NC3). In fact, when é1 62 x 0 and x(WL) x 1. the RHS of (4) and ( is close 1. Below give a proof sketch. The complete argument is deferred to Appendix
The restrictions 61 8K xY and 61 S (K-1)N are mild, and
B where we state a more precise version of the result (see Theorem B.1 tracking the dependence of the bounds on N, K and 8 (Y).
Proof sketch. We start with NC1. If Z already well-collapsed (which guaranteed by the approximate interpolation) and Wz is well-conditioned, then the only source of within-class vari- ability in Zı--1 is within the null space of W,. However, if W, and W,.1 are balanced, the image of Zr.-1 must be approximately in subspace of the row space of W, and, hence, ZL-1 has little freedom within the kernel of WL. More formally, consider first the case of perfect balancedness (i.e.. 62 = 0), and denote by W the pseudo-inverse of WL. Then, (m(1) Im(,,) and
W W2 WEYll, -| WEY sK(WL)
As 8x **.) can be lower bounded by using the assumptions on approximate interpolation and boundedness of representations, the RHS of (6) is small and, therefore, ZL-1 close matrix with zero within-class variability. Moving to the case E2 * O, we need to show that W Z is close to ZL-1. As WtW,, projects onto the row-space of Wz, only the part of Z in the kernel of WL has be considered. This part is controlled after writing Zμ-1
Wi-1Zt-2. using the boundedness of Zi.-2 2 and the approximate balancedness between Wr, and WT-1- Finally, as Zr.-1 close to WEYY, a direct computation yields the bound on NC1.
Next, bound s:WL), we notice that (W,WT) 1 (WL.L,+ WT L.L,+1) has small operator norm, since the weights of linear layers are approximately balanced. This allows to upper bound x <WW, (2L 2L in terms of s(Wr..,, (1),,,+++).. (plus small perturbation), which gives (3) To lower bound the NC3 metric, we rescale ZL-1 and WL to Z and Wi, so that their columns and rows, respectively, have roughly equal size. Then, we reformulate the problem to proving that (Zi_1, WLY) is close to its theoretical maximum. We proceed to show this by arguing that, this scaling and given that Wz is sufficiently cell-condiiioned. WV can be replaced by +. As Zz-1 is close to W Y, obtain (5). Finally, the bound on NC2 in (4) follows by combining (3) with the closeness between 71-1 and W Y already obtained in the proof of NC1.
4 GRADIENT DESCENT LEADS TO No WITHIN-CLAS VARIABILITY (NC1)
In this section, we show that NC1 holds for networks with one wide layer followed by a pyramidal topology, considered in Nguyen & Mondelli 2020). To do so, we show that the balancedness and
4

---

Published as a conference paper at ICLR 2025
interpolation conditions of Theorem 3.1 hold. We consider the neural network in ①I) and minimize the A-regularized square loss (x
=1 - v + 3 |0ll where 2r. and y obtained by vectorizing Zz and Y, and collects all the network parameters. To do so, we consider the gradient descent (GD) update 0K+1 02 - (0x), where 1) the step size and 02- = (W }= contains all parameters at step k. We also denote by Z the output of layer é after k steps of GD. We make the following assumption on the pyramidal topology of the network, noting that this requirement is
also common in prior work on the loss landscape (Nguyen & Hein 2017 2018
Assumption 4.1. (Pyramidal network topology) Let n 2 N and n2 2 ng 2 2 R.L
We make the following assumptions the activation function of the non-linear layers. Assumption 4.2. (Activation function) Fix y E (0, 1) and B 2 1. Let a satisfy that: (i) (x) e (7, 1], (ii) (x) S [x] for every E I. and (iti) of is 3-Lipschitz.
This includes smooth leaky ReL.Us, the same assumption was employed by Nguyen & Mondelli
(220); Frei et al.] (2022) and similar one by Chatterjee| (2022). In principle, can change at all layers, as long as it satisfies the above requirement. Next, let us introduce some notation for the singular values of the weight matrices at initialization Og (WD)
A mmin (W"), d || + mineets......) de, A
II E=Í
Ae,
II
de.
(7)
t=i
We also define Ap= layer at initialization. Finally, make the following assumption on the initialization. Assumption 4.3. (Initial conditions)
= min (o (WPX) X)) as the smallest singular value of the output of the first hidden
Ards-L min(AF, minre 1.) de) 2 8T8 (2) L
Co((o).
(8)
We note that 8 can be satisfied by choosing a sufficiently small initialization for the second layer and a sufficiently large one for the remaining layers. In fact, the LHS of (8) depends on all the layer weights except the second, so this quantity can be made arbitrarily large. Next, by taking sufficiently small second layer, the term 22))(0) = IZ Y| can be upper bounded by ||||
As consequence, the RHS of (8) is at most || 4/2 As the LHS of (8) can be arbitrarily large, the inequality holds for a suitable initialization.
Theorem 4.4. Let the network satisfy Assumption
4.1 a satisfy Assumption 4.2 and the initial
conditions satisfy Assumption 4.3 Fix 0 < 61 S
(K-1)N
K 62 > 0, and run l steps of X-
2
regularized GD with step size n, where > ((?)). O(22) and l: Ω I log 1 Then,
c322
NCI(Z = 0 ((EL + )=
In words, if regularization and learning rate are small enough and we run GD for sufficiently long, then the within-class variability vanishes. Below we provide a proof sketch. The complete argument is deferred to Appendix B where we state a more precise version of the result (see Theorem [B.2] tracking the dependence of A, k and (C on N. K, the data X, the network architecture and the initial conditions.
Proof sketch. We show that the network trained via X-regularized GD fulfills the three sufficient conditions for NC1 given by Theorem i.c., approximate interpolation, approximate balancedness and bounded representations/weiihns To do so, we distinguish two phases in the training dynamics.
The first phase lasts for logarithmic time in 1/A (or, equivalently. 1/ei) and, here, the loss decreases exponentially fast to value of at most 2Am, < e2. As the learning rate is small enough, the loss cannot increase during the GD dynamics, which already gives approximate interpolation. To show the exponential convergence, we proceed in two steps. First, Lemma 4.1 in (Nguyen Mondelli 2020) gives that the unregularized loss (0) satisfies the Lojaak--ojasiewicz (PI.) inequality
(0)
a' 2
CO)))
(10)
"oo avoid confusion, we that this notation is different from the one used in Nguyen & Mondelli 2020
5

---

Published as a conference paper at ICLR 2025
for all 0 in ball centered at initialization Ou and with sufficiently large radius (captured by ro). Next, we show that, if (d) satisfies the a-PL inequality in (TO), then the regularized loss (x) Co(O) + 2 |00li satisfies shifted a-PL inequality, which implies exponential convergence. This second step is formalized by the proposition below proved in Appendix B
Proposition 4.5. Let Co(f) satisfy the a-PI. inequality (T0) in the ball B(0o, ro). Then, in the same ball, (0) satisfies the inequality
|VX200 (Cx(0) -m,))
(11)
where my (1 + v 42/a)² (1601l2 + rr) Furthermore, assume that CVCx(b))/a and VCo(f) is 31 -Liischhh in B(fu, ro). Then, for any 1l < 1/(231), there exists
log Ex Amy
Amn
k <
(12)
log - m7g
such that the ky -th iterate of GD satisfies
C2(00)
(0,, S 2Amx:
110x, -0l12 s 8V
S ro-
(13)
The second phase lasts for linear time in 1/A (or, equivalently in 1/cf) and logarithmic time in 1/e2 and, here, the weight matrices in the linear part of the network become balanced. More precisely, we adapt the analysis of (Du et al. 2018) to show that, if W a weight matrix of the linear part, [|W242 We+1 W. lpop decreases exponentially and the exponent scales with 1/X, which gives approximate balancedness. Finally, the regularization term in the loss is at most et. the operator norm of representations and weight matrices is bounded by as in (29). and the proof is completed by an application of Theorem 3.1
5 ORTHOGONALITY OF CLASS MEANS (NC2) AND ALIGNMENT WITH LAST WEIGHT MATRIX (NC3)
To guarantee the orthogonality of class means and their alignment with the last weight matrix, the crux is to show that the condition number <<W,,, of the last weight matrix W,. is close to one. In fact, Z 2 WtY, y this implies that the last hidden representation ZL-1 is approximately rotation and rescaling of Y, which gives NC2, and a bound on NC3 of the form (5) also follows.
The fact that <(WL) = is a consequence of the presence of many balanced linear layers. Indeed, balancedness implies W.W.. = (WL.L.-WT..,.-ll - which gives that k(Wr) = x< ((WL::L+1) 2/2 Thus, if the conditioning of the product of the linear layers WL:L++1 can be bounded independently of L2, Le + co. As this difficult to obtain in full generality (in particular the assumptions of Theo- rem 3.1 may be sufficient), show that the conditioning be controlled (i) any global minimizer, and (ii) when the parameters are 'stable' under large learning rates.
guarantee that the conditioning of Wz approaches as
5.1 GLOBAL MINIMIZERS
We first show that any set of parameters that approximately interpolate with small norm has bounded condition number. We will then show that with the right choice of widths and ridge, all global minimizers satisfy these two assumptions.
Proposition 5.1. Let satisfy Assumption
4.2
Then, for any network that satisfies
approximate interpolation, i.e., 12, Y S
bounded parameters, i.e., ||0112 < LK +c,
the linear part W WL.1441 satisfies
8x sx ((x))=))
(11 S exp
T/
c+ LiK logK - 2K log
| op
(14)

---

Published as a conference paper at ICLR 2025
Theorem 5.2. Let satisfy Assumption 4.2
Assume there exist parameters of the nonlinear part
i nonttnn (We) such that LL, Y and 1onotin ll2 Then, any global minimizer of the
regularized loss (x( = 21 Zz. 11.
|101 with X 2
ot
we have
KLtc
|| lop
K
exp
(c
K In K K)
8x
(15)
||X lop
L2
<(WL) S
exp
Le LIK + L1 log K))
(Y -11
2L2
This implies that the bounds on NCI, NC2 and NC3 in x)L) upper bounded as above and E2 0.
and 5 respectively, hold with
The assumption that the parameters of the nonlinear part be chosen to fit the labels Y is guar- anteed for large enough width (i) by relying on any traditional approximation result (Hornik et al.] Leshno al.] 1993 Arora et 2 He et al. 2018), or (ii) by taking the infinite time limit of any convergence results (Nguyen & Mondelli (020), or (iii) by taking the limit A >> 0 in Proposition 4.5 A sketch of the arguments below, with full proofs deferred to Appendix
1989
Proof sketch. To prove Proposition 5.1 we write the norm of the linear and nonlinear parts of the network in terms of the conditioning of W W1 for the linear part, this is direct computation: for the nonlinear part, we use Theorem I of Dai et al. (202T), which lower bounds the norm of the parameters terms of the product of the singular values, and manipulate the latter quantity to obtain again the desired conditioning. Next, to prove Theorem 5.2] we pick the parameters of the nonlinear part Ononlin s.t. Zr, = Y and ||0nontim| 12 = c, and set the linear layers to the identity. This leads to a total parameter norm of K KL2 + a and regularized cost of 2) (KT,2 T.2 + e), and it forces the global minimizer to satisfy the assumptions of Proposition 5.1 which gives the claim on <(WL..,++l
Then, as all local minimizers have balanced linear layers, s(W2) = =(WL=L,+1) which gives Finally, the claim NC1, NC2 and NC3 follows from application of Theorem 3.1
15
5.2 LARGE LEARNING RATES
Previous works have observed that the learning rates used in practice typically 'too large', i.c. the loss may not always be strictly decreasing and GD diverges from gradient flow (Cohen et al. 2021). Thankfully, instead of simply diverging, for large 77 (but not too large) the parameters naturally end up at the 'edge of stability": the top eigenvalue of the Hessian HCO is close to 77' i.e., the threshold below which GD is stable (Cohen et al., 2021 Lewkowycz et al. 2020). One can thus interpret GD with learning rate 7i as minimizing the cost C₁ amongst parameters 0 such that ||HCx| op " These observations are supported by strong empirical evidence and have been also proved theoret- ically for simple models (Damian et al.) 2022). although general result remains difficult to prove due to the chaotic behavior of GD for large 1). Specifically, the Hessian has the form
(0) (VoZz) 0nZr + Tr [(Y [1,] + Alp,
(16)
with VoZr E PXNK and P the number of parameters. The first term is the Fisher information
matrix: this is dual to the NTK e = V, ZT. "E &RRNKXNK (Jacot et al. 2018). Therefore, approximately interpolating points, have |/HCCx - 1|01| the interpolation error and X the regularization parameter. We can thus interpret large learning rates as forcing a bound on the operator norm of the NTK. For networks that approximately interpolate the data with bounded NTK and bounded weights, the following proposition (proved in Appendix B guarantees good conditioning of the weights in the linear part (and therefore NC2-3):
O + O(A), where el is
Proposition 5.3. For any network that satisfies
| (T,TT22,A"11
bounded NTK, i.e., [10|| = maxA
& CL2
1.1
approximate interpolation. i.e., 12 Y 1l1 SE bounded weights, i.e., |[Wellop S r,

---

Published as a conference paper at ICLR 2025
for any M S L2, there is E +1,+1,... L1
M3 such that Kc s(Wr.s) S
COLkr
(M(a)(Y)-o1))
Furthermore, any network that satisfies approximate interpolation and bounded weights is such that
(s I.2
|0| op 2
(17)
K2p²
As an example, by choosing M- L12 we guarantee that there is at least one layer & in the first half
2"
of the linear layers s.t. x(WL.)) S -x Now, assuming the 'edge of stability' phenomenon, learning rate of I- implies a bound ||@|.pp op +O()))+()) I 212 +0(E))+)+0(N)) and thus Proposition 5.3 implics that there is linear layer with bounded conditioning. The issue is that the proof of Theorem requires an extremely small learning rate for large depths (17 ~ e see Theorem B.2 in Appendix B whereas Proposition 5.3 is only useful when 17 ~ L-1 Note that the first phase of training (where the dynamics is analogous to the NTK regime) is stable when 7 I,-L; it is the second phase that appears require 7 c-L ensure GD remains close to gradient flow (we have much less control of the dynamics in this second phase and so we cannot rule the possibility that GD approaches highly unstable regions). We are hopeful that more careful analysis could show that learning rate of ~ L would push GD away from ill ill-condttonne and unstable regions (possibly deviating from gradient flow), so that the end of the second phase the network would be interpolating, balanced, and have well-conditioned linear part thus guaranteeing NC1-3.
vaCkr
The
6 NUMERICAL RESULTS
In all experiments, we consider MSE loss and standard weight decay regularization. We train an MLP and a ResNet20 with an added MLP head on standard datasets (MNIST, CIFAR 10), consider- ing as backbone the first two layers for the MLP and the whole architecture before the linear head for the ResNet. We evaluate the following metrics related to neural collapse: for NC1, we compute ((Ew)/tt(Enn) where Ew, En are the within- and between-class variability matrices of the feature matrices, respectively, for NC2, we display the conditioning number of the class-meat Meas matrix; for NC3, we use the average cosine angle between the rows of weight matrix and the columns of the [W Wei1 - W.w: Il min [[] Wees II,, wwaw
preceding class-mean matrix. Finally, for balancedness we use
and
for negativity we use Dell. [ We measure such metrics (and also index the layers) starting
from the output of the backbone. Our findings can be summarized as follows (see Appendix A for
additional complementary experiments).
The deeper the linear head, the more clear NC occurs. We first test whether the models with deep linear heads exhibit NC and that's the case, whether it gets better as we deepen the linear head. In Figure II we show the NC metrics and the gram matrices of the class- mess-mmean matrices of the last layers for the training on CIFARIO of ResNet20 with extra layers of which the first three have a ReLU activation. We weight decay of 0.001 and learning rate of 0.001, training for 5000 epochs (the learning rate drops ten-fold after 80% of the epochs all our experiments). The plot clearly shows that the collapse is reached throughout the training. We also see that the NC2 metric improves progressively with each layer of the linear head, as predicted by our theory. This effect is also clearly visible from the gram matrices of the class-means (bottom row of Figure DD. which rapidly converge towards the identity. We note that these findings are remarkably consistent across wide variety of hyperparameter settings and architectures. Waag et al. (2023) present similar results.
In Figure 2 we plot the dependence of the NC metrics on the number of layers in the linear head. We train an MIP on MNIST with 5 non-linear layers and number of linear layers ranging from to We average over 5 runs per each combination of weight decay (0.001, 0.004) and with learning rate of 0.001. We also train the ResNet20 on CIFARIO with one non-linear layer head and to linear layers on top. We use the same weight decay and learning rate. The relatively high variance of the results is due to averaging over rather strongly different weight decay values used per each depth. The plots clearly show that the NC2 significantly improves in the last layer the depth of the linear head increases, while gets slightly worse in the input layer to the linear head. This is consistent with our theory. The NC1 does not have strong dependence with the number of layers.
8

---

Published as a conference paper at ICLR 2025
DVCC
regativitee
10
belencndrace
10
10-4
301
Byr
10.
Sym
Bx9r33
Syer
10"h
0.49
2.80
0.01
0.02
0.02
0.04
1,00
0.00
0.01
0.82
3.04
0.00
0.0
0.01
0.02
0:53
0.04
0.03
1000
2000
5000
4000
5000
Rub
625
6706
520
025
Figure 1: Last 7 layers of 9-layer MLP trained on MNIST with weight decay 0.0018 and learning rate 0.001. Top: NC1s, NC2s, balancednesses and negativities, from left to right. Results are averaged over 5 runs, and the confidence band at I standard deviation is displayed. Bottom: Class- mean matrices of the last three layers (i.e., the linear head), the first before the last ReLU.
-3.0
-2.55
3.0
335
0
simber layers
number of Tayarr
number of lapers
number of lapers
-2.0
-2.
-20
log
-20
-35
-4.
Figure 2: Upper/Lower row: MLP/ResNet20 with a deep linear head. Left to right: NC1 in the last layer; NC1 in the first layer of the linear head; NC2 in the last layer; NC2 in the first layer of the linear head. All plots are a function the number of layers in the linear head. Results are averaged over 50 runs (5 runs for each of the 10 hyperparameter setups), and the confidence band at standard deviation is displayed.
10
0.20
06
10
6.10
02
0.
number layers
number layyers
Trmber layers
layers
Figure 3: Left to right: Minimum balancedness; mean balancedness: minimum negativity; mean negativity across non-linear layers of the head as function of the number of non non-linear layers. Re- sults are averaged over 10 runs (5 runs for each the 2 hyperparameter setups), and the confidence band at standard deviation is displayed.
Linear layers are increasingly balanced throughout training. Figure I shows that the metric capturing the balancedness exponentially decreases until it plateaus rather small value (due to the large learning rate) and then again it exponentially decreases at a smaller rate after reducing the learning rate. The balancedness of non-linear layers instead plateaus significantly larger value.
Non-linear layers are increasingly balanced and linear, as the depth of the non-linear part increases. In Figure 3 we show the dependence of balancedness and non-linearity of the

---

Published as a conference paper at ICLR 2025
linear layers function of the depth of the non-linear part. We plot minimum balancedness and non-linearity across all layers, as well as mean balancedness and non- non-linearity. The depth of the non- linear part ranges from 4 to 12 (the first two layers are considered as backbone and not measured), the learning rate is either 0.001 or 0.002 (5 runs each), and the weight decay is 0.016 divided by the total number of layers. Balancedness clearly improves with depth, both on average per layer and in the most balanced layer. Similarly, the non-negativity of the layer that least uses the ReLU clearly decreases, and a decrease in the mean negativity also reported (although less pronounced). Thus, these results suggest that the network tends to use non-linearities less as the depth increases and, in addition, it becomes more balanced. We also note that, in order to fit the data, regardless of the depth, the last non-linear layer(s) always exhibit significant negativity, heavily relying on the ReLU.
7 DISCUSSION AND CONCLUDING REMARKS
Our paper presents a framework to provably obtain neural collapse, which leads to the first result for the end-to-end training of deep networks. We now discuss our assumptions, compare with earlier work by Nguyen & Mondelli (2020) and with the NTK regime, concluding a future direction.
Assumptions. Theorem 3.1 provides connection between neural collapse and properties of well- trained networks, i.c., approximate interpolation, approximate balancedness, bounded representa- tions/weights and good conditioning of the features. As such, the result rather general and requires no assumptions beyond the aforementioned properties. Theorem 4.4 then instantiates our framework for proving neural collapse to a class of networks with pyramidal topology (Assumption[4.1] smooth activations (Assumption 4.2) and for a class of initializations (Assumption 4.3) These assumptions are used only for the analysis of the first phase of the training dynamics, where the network achieves approximate interpolation. Thus, they could be replaced by any other set of assumptions guarantee- ing that gradient descent reaches small training loss. Specifically, such guarantees are obtained by Zou & Gu (2019) for deep ReLU networks (with stronger requirements
aver-parameetrizzaton but no assumptions on the topology) and by byBoobari et al.] (2022) for networks with minimum over- parameterization (under a requirement on the topology milder than Assumption 4.1). As concerns Assumption by following the argument in Appendix C of Nguyen
4.3 on the initialization, we discuss in page 5 a setting where it holds. In addition, Mondelli (2020). one readily obtains that Assumption 4.3 also holds for the widely used LeCun's 's initialization, i.e., W.O has i.i.d. Gaussian entries with varancee 1/n,-1 for all 8 [Z]. as long as n, = S2(N).
Comparison with Nguyen & Mondelli (2020). The condition n = 2(N) mentioned above pro- vides a strict improvement upon n S7 N2) required in Section 3.2 of Nguuen & Mondelli 2020) for LeCun's initialization. This is due to Assumption 4.3 being weaker than the corresponding con- dition on initialization in ingugeen&& & Mondelli (2020), see Assumption 3.1 therein. We now elaborate on similarities and differences between the proof of Theorem and the analysis in Nguyen & Mon- delli (2020). The strategy to handle the first phase of the dynamics in Theorem 4.4 is similar: the weights of the network reach an arbitrarily small loss without leaving suitable ball centered at ini- tialization. However, the implementation of this strategy is significantly different and our approach relies on Proposition 4.5) More precisely, the improvement comes from upper bounding the gradi- ent as in (26) in Appendix B which uses again the PL inequality. In contrast, Nguyen & Mondelli (2020) use the loose bound in (18) of their work. We also note that the analysis of the second phase of the dynamics is entirely new, as balancedness was not needed in Nguyen & Mondelli (2020).
NTK regime and beyond. Our analysis of the training dynamics in Theorem 4.4 is split into two phases. While the first phase uses typical NTK tools, the second phase is governed by the effects of weight decay, leading fundamentally different behavior. In fact, at the end of the second phase, we observe low-rank weight matrices and feature learning, which could not appear with pure NTK dynamics. The two-phase behavior arises thanks to a separation of timescales as gets smaller: the NTK phase lasts roughly (1n steps, but the effect of weight decay is only observable over the longer timescale of the second phase, which lasts roughly O(1) 1 steps. This strategy gives us the best of both worlds: interpolation from the NTK regime, and balancedness from weight decay.
One important assumption in our work is that there is a linear head in the network, containing at least two layers. However, experimental evidence suggests that two linear layers are not necessary for collapse to happen. Thus, an exciting future direction is to use the approach developed here, e.g., the NTK analysis and the separation of timescales, to prove collapse without two linear layers.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
M. M. and P. S. are funded by the European Union (ERC, INF2, project number 101161364). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency. Neither the European Union nor the granting authority can be held responsible for them.
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory (COLT), pp. 4782 4887, 2022.
George Andriopoulos, Zixuan Dong, Li Guo, Zifan Zhao, and Keith Ross. The prevalence of neural collapse in neural multivariate regression. arXiv preprint Xiiv:2409.001800, 2024.
Raman Arora, Amitabh Basu, Poorya Mianjy, and Anirbit Mukherjee. Understanding deep neural networks with rectified linear units. In International Conference on Learning Representations (ICLR), 2018a.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In International conference on machine learning (ICML), pp. 244-253, 2018b.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629-681) 2017.
Daniel Beaglehole, Peter Súkeník, Marco Mondelli, and Mikhail Belkin. Average gradient outer product as mechanism for deep neural collapse. arXiv preprini (4t.:2402.137255 2024.
Ido Ben-Shau and Shai Dekel. Nearest class-center simplification through intermediate layers. In Topological, Algebraic and Geometric Learning Workshops, 2022.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems (NeurIPS), pp. 9768-997 9766,9773,42 2022.
Simone Bombari, Mohammad Hossein Amani, and Marco Mondelli. Memorization and optimiza- tion deep neural networks with minimum over-paaaeeterization. Advances in Neural Infor- mation Processing Systems (NeurIPS), 2022.
Souray Chatterjee. Convergence of gradient descent for deep neural networks. arXiv preprint arXiv:2203. 16462, 2022.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability, In International Conference on Learning Representations (ICLR), 2021.
Zhen Dai, Mina Karzand, and Nathan Srebro. Representation costs linear neural networks: Anal- ysis and design. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilizztion:: The implicit bias of gradient descent the edge of stability. arXiv preprint arXiv:2209. 15594, 2022.
Hien Dang, Tan Nguyen, Tho Tran, Hung Tran, and Nhat Ho. Neural collapse in deep linear net- work: From balanced to imbalanced data. In International Conference on Machine Learning (ICML), 2023.
Hien Dang, Tho Tran. Tan Nguyen, and Nhat Ho. Neural collapse for cross-entropy class- imbalanced learning with unconstrained ReLU feature model. arXiv preprint arXiv:2401.02054 2024.
11

---

Published as a conference paper at ICLR 2025
Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. In Advances in Neural Information Processing Sys- tems (NeurIPS), 2018.
Weinan E and Stephan Wojtowytsch. On the emergence of simplex symmetry the final and penul- timate layers of neural network classifiers. In Mathematical and Scientific Machine Learning, 2022.
Weinan E, Chao Ma, and Lei Wu. Barron spaces and the compositional function spaces for neural network models. arXiv preprint arXiv:1906.0809 08.0999 2019.
Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer- peeled model: Minority collapse imbalanced training. In Proceedings of the National Academy of Sciencess (PNAS), volume 118, 2021.
Spencer Frei, Niladri S Chatterji, and Peter Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Conference on Learning Theory, pp. 1668-27053 PMLR, 2022.
Tomer Galanti, András György, and Mareus Hutter. Improved generalization bounds for transfer learning via neural collapse. In First Workshop on Pre-training: Perspectives, Pitfalls, and Paths Forward at ICML, 2022a.
Tomer Galanti, Zachary S Siegel, Aparna Gupte, and Tomaso Poggio. Sgd and weight decay prov- ably induce a low-rank bias in neural networks. arXiv preprint arXlv:2206.05794, 2022b.
Connall Garrod and Jonathan Keating. Unifying low dimensional observations deep learning through the deep linear unconstrained feature model. arXiv preprint arXiv:2404.06106, 2024.
Jarrod Haas, William Yolland, and Bernhard T Rabus. Linking neural collapse and normalization with improved out-of-distributio detection in deep neural networks. Transactions 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. International Conference on Learning Representations (ICLR), 2022.
Juncai He, Lin Li, Jinchao Xu, and Chunyue Zheng. Relu deep neural networks and linear finite elements. arXiv preprint 1rXiv:1807.039/7 2018.
Wanli Hong and Shuyang Ling. Neural collapse for unconstrained feature model under cross- entropy loss with imbalanced data. arXiv preprint arXiv:2309.09725, 2023.
Wanli Hong and Shuyang Ling. Beyond unconstrained features: Neural collapse for shallow neural networks with general data. arXiv preprint arXiv:2409.01832, 2024.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni- versal approximators. Neural Networks, 2(5):359 366, 1989.
Arthur Jacot. Implicit bias of large depth networks: a notion of rank for nonlinear functions. In International Conference Learning Representations (ICLR), 2023a.
Arthur Jacot. Bottleneck structure in learned features: Low-dimension vs regularity tradeoff. In Advances in Neural Information Processing Systems (NeurIPS), volume 36, pp. 23607-23629, 2023b.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gener- alization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
Arthur Jacot, Eugene Golikov, Clément Hongler, and Franck Gabriel. Feature learning in 1.2- regularized dnns: Attraction/repulsionn and sparsity. In Advances in Neural Information Pro- cessing Systems (NeurIPS), 2022.
12

---

Published as a conference paper at ICLR 2025
Wenlong Ji, Yiping Lu, Yiliang Zhang, Zhun Deng, and Weijie Su. An unconstrained layer- peeeep perspective on neural collapse. In International Conference on Learning Representations (ICLR), 2022.
Jiachen Jiang, Jinxin Zhou, Peng Wang, Qing Qu, Dustin G Mixon, Chong You, and Zhihui Zhu. Generalized neural collapse for a large number of classes. In Conference on Parsimony and Learning (Recent Spotlight Track), 2023.
Vignesh Kothapalli and Tom Tirer. Kernel vs. kernel: Exploring how the data structure affects neural collapse. arXiv preprint arXiv:2406.00 02100, 2024.
Vignesh Kothapalli, Tom Tirer, and Joan Bruna. A neural collapse perspective on feature evolution in graph neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum margin bias of quasi- hoageeoomo neural networks. In International Conference on Learning Represen- tutions (ICLR), 2022.
Jason D Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information-theonee theoreeic limit. arXiv preprint arXiv:2406.01581, 2024.
Moshe Leshno, Vladimir Lin, Allan Pinkus, and Shimon Schocken. Multilayer Feedforward Net- works with Non-Polyno Activation Function Can Approximate Any Function. Neural Net- works,
1993.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint 02220, 2020.
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. In Conference on Parsimony and Learning (Recent Spotlight Track), 2023a.
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. In Conference on Computer Vision and Pattern Recognition (CVPR), 2023b.
Janneng Lu and Stefan Steinerberger. Neural collapse under cross-entropy loss. Applied and Com- putational Harmonic Analysis, 59, 2022.
Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. arXiv preprint ar arXiv:2011.11619 2020.
Quynh Nguyen and Matthias Hein. The loss surface of deep and wide neural networks. In Interna- tional Conference on Machine Learning (ICML), 2017.
Quynh Nguyen and Matthias Hein. Optimization landscape and expressivity of deep CNNs. In International Conference on Machine Learning (ICMI.), 2018.
Quynh Nguyen and Marco Mondelli. Global convergence of deep networks with one wide layer fol- lowed by pyramidal topology. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Leyan Pan and Xinyuan Cao. Towards understanding neural collapse: The effects of batch normal- ization and weight decay. arXiv preprint arXiv:2309.0446 04644, 2023.
Vardan Papyan, X. Y. Han, and David I. 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.
Akshay Rangamani and Andrzej Banburski-Faie Fahey...iFF.y...a. Neural collapse in deep homogeneous classifiers and the role of weight decay. In IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2022.
13

---

Published as a conference paper at ICLR 2025
Mariia Seleznova, Dana Weitzner, Raja Giryes, Gitta Kutyniok, and Hung- Hsu Chou. Neural (tan- gent kernel) collapse. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Peter Súkeník, Marco Mondelli, and Christoph Lampert. Deep neural collapse is provably optimal for the deep unconstrained features model. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Peter Súkeník, Marco Mondelli, and Christoph Lampert. Neural collapse versus low-rank bias: Is deep neural collapse really optimal? arXiv preprint arXiv:2405. /4468, 2024.
Christos Thrampoulidis, Ganesh Ramachandra Kini, Vala Vakilian, and Tina Behnia. Imbalance trouble: Revisiting neural-collapse geometry. In Advances in Neural Information Processing Systems (NeurIPS). 2022.
Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural col- lapse. In International Conference on Machine Learning (ICML), 2022.
Tom Titer, Haoxiang Huang, and Jonathan Niles-Wees Perturbation analysis of neural collapse. In International Conference Machine Learning (ICML), 2023.
Zhenfeng Tu, Santiago Aranguri, and Arthur Jacot. Mixed dynamics in linear networks: Unifying the lazy and active regimes. arXiv preprint arXiv:2405.17588 2024.
Peng Wang, Huikang Liu, Can Yaras, Laura Balzano, and Qing Qu. Linear convergence analysis of neural collapse with unconstrained features. In NeurIPS Workshop Optimization for Machine Learning (OPT), 2022.
Peng Wang, Xiao Li, Can Yaras, Zhihui Zhu, Laura Balzano, Wei Hu, and Qing Qu. Understanding deep representation learning via layerwise feature compression and discrimination. arXiv preprint arXiv:2311.02900 2023.
Sicong Wang, Kuo Gai, and Shihua Zhang. Progressive feedforward collapse of resnet training. arXiv preprint 2rXiv:2405.00985, 2024.
Zihan Wang and Arthur Jacot. Implicit bias of SGD in Læ-regularized linear DNNs: One-way jumps from high to low rank. In International Conference on Learning Representations (ICLR), 2024.
Yuxiao Wen and Arthur Jacot. Which frequencies do CNNs need? Emergent bottleneck structure in feature learning, arXiv preprint arXiv:2402.088 20000000000000 2024.
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 gener- alization bounds. In Research, volume 6, 2023.
Emanuele Zangrando, Piero Deidda, Simone Brugiapaglia, Nicola Guglielmi, and Francesco Tud- isco. Neural rank collapse: Weight decay and small within-class variability yield low-rank bias. arXiv preprint arXiv:2402.03991 2024.
Jiawei Zhang, Yufan Chen. Cheng Jin, Lei Zhu, and Yuantao Gu. EPA: neural collapse inspired robust duitof-dist detector. arXiv preprint 44X1v:2401.01710, 2024.
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), 2022a.
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 Advances in Neural Information Process- ing Systems (NeurIPS), 2022b.
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geo- metric analysis of neural collapse with unconstrained features. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Difan Zou and Quanquan Gu. An improved analysis training over per-parameterized deep neural networks. In Advances in Neural Information Processing Systems (NeurlPS). 2019.
14

---

Published as a conference paper at ICLR 2025
DMC
20"
Delercndrace
3.0
0.8
10-1
30
Sym
12"
30h23
Syee
0.2
0.50
2,92
0.04
0.08
0.30
0.00
0.02
0.06
3.68
0.10
6.03
0.02
0.04
0.0
0.00
0.30
2000
4000
6000
1000
19000
60
Figure 4: Last 7 layers of a 9-layer MLP trained on MNIST with weight decay 0.0018 and learning rate 0.001. Top: NCIs. NC2s, balancednesses and negativities, from left to right. Results are averaged over 5 runs, and the confidence band at standard deviation is displayed. Bottom: Class- mean matrices of the last four layers (i.e., the linear head).
2.4
1.6
25-
-5.15
14-
1.2
Figure 5: ResNet20 trained on MNIST with a deep linear head. Left to right: NC1 in the last layer; NC1 in the first layer of the linear head; NC2 in the last layer; NC2 in the first layer of the linear head. All plots are function of the number of layers in the linear head. Results are based on runs (5 runs for each of the 10 hyperparameter setups), and the confidence band at I standard deviation displayed.
A ADDITIONAL EXPERIMENTS
We complement the experiments from Section with additional numerical findings. We start by showing an analog of Figure for MLP trained on MNIST to show that the behavior is robust with respect the backbone and dataset. The results are shown in Figure 4 The architecture is MLL with non-linear layers followed by 4 linear layers. We use weight decay 0.0018 and learning rate of 0.001, training for 10000 epochs. The plot fully agrees with the one in Figure] with ResNet20 training on CIFARIO in every qualitative aspect, and it only differs in the numerical values attained by the NC metrics. Therefore, the conclusion is that the NC attained across different architectures and NC2 progressively improves as we get closer to the last layer of the DNN.
Next, we extend Figure 2 with experiments on ResNet20 trained on MNIST, which are shown in Figure 5 As above, the results match the interpretations discussed in Section 6 which proves the robustness of findings across different achitectures, datasets and hyperparameter settings.
B DEFERRED PROOFS
Theorem B.1 (More precise version of Theorem 3.1. If the network satisfies
approximate interpolation, i.e.. [Z Y <1.
approximate balancedness, i.e. [ 41 We+1 WoW, 1lop S Ser f l (L1 + +,......... L-
15

---

Published as a conference paper at ICLR 2025
bounded representations and weights, i.e, 1| 1,..., L.)
WWellop r. l {L1 +
then i S min (xx(Y), [K 1)N
x2
r),
NCI(21.1)
(18)
N
2
K-1 K
VN
where (c1,, C2, = ( +
V II
we we additionally assume that the linear part of
the network not too ill-conditioned, i.e., x(WL:L.+1 S C3, then
<((,)) S o3 10 (1 + +e e) + c/3
1/
(19)
43 2(22 2)
with == ( 4 12(22-1)e2
Finally,
1x13, 301 181
NC2(Z1-1) x(Wr) +xx(Y)=+tV(e))eegrr)
(20)
1 SK -xx(Y))-r(((1
(VN - -1)2 ((((WW)) - 21x(W2((+GG)
NC3((1-1,,L))
(21)
By performing Taylor expansions and some algebraic manipulations, one readily obtains 2
from 18
19
20
21 respectively.
Proof of Theorem B.
NC1: We start by proving the claim [8 NC1. We have
8x < (Y + (Z, Y)) 8x (2.) s (W, Zr. S sx(WW))
(22)
where the first inequality the application of Weyl's inequality for singular values and the last one uses sy (A)si (B) 2 sı (AB). Denote W the pscudoinverse of WL. Then, we have that W Wz equals the projection /' on the row space of Wz. We can now write
Zμ-1 PZ (I- - W (2 - +)+ PWW-)----
Note that
E1T S sX(Y) -
since H *w(W) c1)r by 22 Furthermore,
[I -Ll-PPW...m (IU-P)W.- - 21 (
P) = P)(WL-IWT- wlWw)) S 2
Putting these together, we have:
176---WYYL Sr (x((((()))) - + VnL-112 = (V((),
(23)
From now on, we will drop the index I - and treat everything without layer label as belonging to that layer (the membership to any other layer will be indexed). We have that
tr(Ew)
1 N
Σ Zes
Mc) - [dOT
1
Σ
-Hell2
N
Ees - (1 2112
112 N
| - -1-1
N
0.3
16

---

Published as a conference paper at ICLR 2025
Furthermore,
k
K
1
tr(En) = tr
Σ
(ue - ua) - uG)T
Me - 1c
K
IIW 1
k
Σ
like -
Σ
llue 12ll
22
Kr
o=l
0=1
Now proceed by lower-bounding the last term:
Σ 11 -
rb1l22
1
lIne t Htlll
K
K
1
Σ
1i - -Hall2
Σ
1
me 11
1
K
where W. 2 are the class and global means of the label matrix Y. A direct computation yields that
1 DK 1I - = /
Next we have:
K
K
1
1 --ol22
Y
2
N
23
N
2
N
2
26
Y
N
Y11 1
VN
4
G2
Finally, for any fixed e we have
1K v/ 112
1
Σ
-
Y
Σ
||
2til222
Zes
n1
n1
Therefore we get:
K - S N
- =dll 1 2z Y1
E1 VN
Upper bounding these terms in the above computation and dividing tr(Ew) ) by tr(EB) yields (8).
Conditioning of Wi: We now prove the claim (19) on the conditioning of Wi. We rely first on Lemma C.2 to relate the singular values of WT to the Lz-th root of those of WL.L++1. We therefore obtain that
$1 (Wr)222 8x (Wr.)222
81 ( (WL.,++1)2
+
4g 2 1)
7
8K (WI:L,+1)*
12, -1)c2
8K SK(W 2 )2
81 (W Wr:,, +1 12 1
8s (W
+
sx (W +1)
SK (W(:,,+++))
13 22,2((221))) 8K (W (1:L,+1)²
We have that
sk())) < 8K (WWWL1++) |Z llop I SK &WK(W..L (WLLL,+1))
| pby
and therefore
8K (Zz)
SK xx
sK (W (WL:L++1) >
1 rhe
|| phe
This gives that
81 (Wr)
$1 (W(L:I,+1))
1 + WWW..WWL.. (WL:L++12
I.2 , 22,2 62.
3K 21.2 2L.
sK (W:L,++1)
(xx(Y)-a))
L3 2 (12211))) 2
2
|X IX 10n
17

---

Published as a conference paper at ICLR 2025
Thus, the following chain of inequalities gives 19
x(WE) = (s(WE.L.+))
(1 +c)
S (W
T (1 + +))
+1 _o2 -1 e)
2he -
I < 1/2 (1 + c) siz <(WL. 12 E,
-2-
using the concavity of the 2Tz-th root.
NC(2+3); We start with the proof of the claim (1 on NC3, as the derivations will be used in the proof of the claim (2) on NC2 later. Since the cosine similarity does not depend on the scale of the involved matrices, we perform the following rescaling: denoting a = ||,,,], we define Wi = Wr.
and - = aZL-1- Similarly to before we denote (z')(4-1) the i-th sample of the c-th class in the matrix - Then, we write
(()) "
,(W))))
NC3(ZZ W2)
1 N
Σ cos((=)) (1-1) (WL)c) (W)))))
Σ
N [00)
11 11(WE)c:ll
0,
Eest (1:) (1) (WW))) 2 (V)
Nx(WL)(1 + + 2N +
12 + ((W) Y1l - -11-111-111- (
Here, the first inequality follows from upper bounding II(W)))=ll trivially by and
(3)) (1-1))) - )
(L) S <x(L)) (1IVe - ((4)
| + +1lvell2) s x(WL)(1+
Now,
12 > 1r 2 (| r 11, - -Y( ~1 (VN
Furthermore, we readily have that [[W]) Ty 2 VV
x(W2) Finally, we have:
1 1Zi-1-(W))TT-) - || - (W,)+ + |[ - (WL) (Wr)TYll,. Tyll..
From 23
we obtain that
1 (W r S <²²
+ VTL- = riv C2,
(24)
(x((V) -1
Finally, we proceed with upper bounding YY (W%) TY
which can be done via a sand- wich bound on the singular values using the conditioning number. This gives
[W (W:) Y S V (((W))
<K(K(W))) 1).
(.W,,.)
Putting all the obtained bounds together, we get the desired bound on NC3.
Finally, to pass from the bound (19) on <WW)) to the bound 20 on NC2, we use the inequality in (24) obtained in the proof of NC3. Note that obtaining a bound on x(Z%,1) is equivalent to
obtaining a bound on n: xx((-1) since multiplying by a scalar does not change the condition number.
Using 24 we get:
3 (W2)+11 = - $ +|| 1 (W2)+
(2,1,,
sY(
As [[WW)) = and 8K (x((W))+ = 1, we conclude that
K( x(W,) + (2
R( 2
1 - (V(c)c))))(xxx(()>-
which gives the desired bound in 200
18

---

Published as a conference paper at ICLR 2025
Proposition 4.5. Let Cu( satisfy the a-PL inequality (10) in the ball ro). Then, in the same ball. (x) satisfies the inequality
(VC10 (Cx(0) -mmx)
(11)
where my (1 /4 (1601l1 + rr) Furthermore, assume that To 2 Cx(x/0))/ and VCo(0) is LipLiphhitz in (fu,, ro). Then, for any 3 < <1/(23)), there exists
log (x(8 -mm)
h <
(12)
log(1 7
such that the ky-th iterate of GD satisfies
C2(Oo)
Cx(Ok,) - 2Amx:
1 00ll2 S 8V
S ro-
(13)
Proof. Let 0 E B(0,, ro)- Then, the following chain of inequalities holds:
(0) 2011 2 (IIVC((0)122 210112)
2
Colo) 100|2
2
at cx(0)
100 2
A
|101
4
2
(25)
CX(0)
A V 4
1101|2
C1 > (Va 4 + vX)
110113
4A a
IV
(x(0) v
1+
(1101l2 ro)²
Here, in the second line we use that V 2 CC(() 2 A |01| which follows from Cx 2 Amy
(otherwise, the claim is trivial); in the fourth line we use that Va Va-b- b > vo vb for 2 b; in the
fifth line we use that (a b)2
122 for all a, b; and in the sixth line we use that @ in the ball
B(0o, ro). As the LHS of 25 equals (0)]l this proves II]
Next, let (0x) be the GD trajectory. Pick ls s.t. 0g B(0, ro). Then,
(x - -C,(4)) - -1 /. (VCx (C2
(0 +1 1'' (VC)(0k - (0.10--SJVCC(00)1 (O.( VCS
[VCI( (Bx)) - VC2(Bx)ll ds
S -- (0)
(0x)))
= (1(1 1181) (0x)1)
< ! (0+)))
2
< " a
(Cx(Ok) - Amx))
Here, in the fourth line we use that the gradient is B1 31-Lipschitz in B(f.. ro); in the sixth line we use that 1 < 1/(23)); ); and in the last line we use [1 Thus, as long as 0z E B(O,, ro) for all j E [b]. by iterating the argument above, we have
(x - Amx S (C)(0)) Amw) (1-1
a
19

---

Published as a conference paper at ICLR 2025
This readily implies that, by letting 11 be the first index & s.t. (x(Ok) the distance |10k, Bull2 is upper bounded by
- 2Amx.
12
holds. Finally,
n >|] = _--
(0x(())
Σ
4
R1
Cx(0x) - (x(0x)))
(CX(0<) Ama
4
Cx(Ok) - Cx(0x+1)
4 Cx(0k-1) C2(0k)
va
VCA(0)) Ama
va (Ci(0k) AmTA
R1 2
Cx(Os) - (3x+1))
02 Ama
5
+ Va (Cx(0k Ama
Va
(CI(H++1 Ama (Ca(0s) dmx
8
(VCx(0)) amy -VC100+11 2mAA
8 + VCx(0n - Amax
Va
k=
8 Va
(2 Ami & Va VC,(0))
(26)
which concludes the proof.
Theorem B.2 (More precise version of Theorem 4.4) Let the network satisfy Assumption
4.1
satisfy Assumption 4.2 and the initial conditions satisfy Assumption 4.3 Fix 61 S
(K
1)N
K
> O, let b 2 1 be s.t. ||x,vl22 S b for all i, and run k steps of >-regllarized GD with step size n7, where
2 2
L-2
(0)(
e
X S min
2
110012
18(1100l12
+ Ap
2p/22/2
1i+b
17 S min
(27)
281
1, (.,%) 267 32/2
2X
2e7
4|
5N 363 max
5
log (1,) -mmx
Ama
Aeg log
k k
log(1 17
log(1 - 172))
with Bi 5N 5N3363 (I1É_ max(1, (1,, (x)) L 1,/2, m. = (1 + VAXaa)² (1100l12 + ro) ro =
min(Ar, minee (3.....) de), and a 2 2- (6- ,2
Then, we have that
(V((1 (2,))))))
NCI(ZZ N
(28)
2// vN
K
with (EI,,E2,)) =r (Fx(F)- - V 12
and
T--2
1-1
2 1IV X
== max
||
|
(29)
Since the terms 31, m,, To. a do not depend on 61, <2 (but just on the network architecture and initialization), taking
+(ef) and π O(e2) (as in the statement of Theorem 4.4p satisfies the first two requirements in (27). also giving that in (29) is of constant order. Finally, as log(1 - x) Ro -r for small z, the quantity 1yk- which quantifies the time of the dynamics, since it is the product of learning rate and number of GD steps is of order log(1/A) + log(1/e2)/A, Thus, after some
manipulations, one obtains that = Ω :.
log
(as in the statement of Theorem 4.4
20

---

Published as a conference paper at ICLR 2025
Proof of Theorem 8.2 By Lemma 4.1 in Nguyen & Mondelli 2020), the loss Co(0) satisfies the
PL inequality with a 47 - 0 min (Z1) I p=3 mmin (Wp). where we have also used that o' is lower bounded by y by Assumption 4.2 Thus, by taking ro min(Ap, mince (3......) de), we have that, for all 0 B(0,, Cu( satisfies the a-PL inequality with 2- (1-33), Apps---
By using Assumption 4.3 and that À is upper bounded by 2Ca(Oa) in (27). one can readily verify 1|%o1]
that ro 2 8y 8/Cx/00)/a. Furthermore, by Lemma C.I. we have that (C.( is 31 3i-Lipschitz for all 0 ((f,, ro). Hence, we can apply Proposition 4.5 which gives that, for some upper bounded in (12
(0s,)) S 2Ama S e
where the last inequality uses again the upper bound on (27)
With gradient flow, the regularized loss would only decrease further after K steps. Since we are working with gradient descent, simply need assume that the learning rate small enough to guarantee decreasing loss. If the gradient (+) is 3e-Lipschitz, then
Cx (0(++11) (0x,) Cx2 T/V (0x)) Cx(0x)
- -- I'' (0C) - VOx))) (Cx(0)))
$ - (VCx(0))
+ +0I1VC>(0=)) 0 VCA VCN(0, s1/VC()))
S - /VCx (1)) + + (0x))
=- n/ (0x)(1
which non-positee non-positive as long as 17 1 1/32. Furthermore, as long the regularized loss is de- creasing, the parameter norm is bounded by E1 2//. Lemma C.I then implics that the gradient is 32 = 3N max (1, c (2) 32/22 -//2- Lips-hips and, therefore, 1 S 1/32 holds by (2)). This
allows us to conclude that, for all & satisfying the lower bound (27). (x S eq hence the network achieves approximate interpolation, i.c.,
12 - Y S 2
(30)
Next, we show approximate balancedness. To do so, for l E {11 + 2,.... 1}. we define Tht = IT ... ...WW)
(W)) -Y ()t " (W
... (W/-
Then, we have
what (W!+1)T = ((1 nn))t -1T -
-11 -L- (W) + +
Similarly, (WE+) Wh+1
= (1 (1-1A)(WWH)) T Wh - -(1 (1-)))(((WWW))) T + (Tht1) W +
Let us define
D (W - (WW))
Since T (WA)T = (WA,)) TT and (T,) = =(T)p1) W we have
D (1 D + ((TA)) T - -TA(TA)
Recall that, for all & lower bounded in 27), Z
YlIF S <e122 and [W/||= S 110*1/2 S e1 S
L,
which also implies that || || - (1 2/X
|| Thus,
1 Map S (1- $ (1 - (11-)))||DA11.... + 7/|127 11
+ +1 (11
+ | 112p)
|Z y
[ + II 11
J
1 1i71-1
S (1 -1->))(/DD 2n
|||
21

---

Published as a conference paper at ICLR 2025
L1+L
By using the upper bounds 1 S 1 and 17 - 25
in
27 we have that:
if [D] lop 2 e2, then | llop S (1 - ||>i| D llop
if |D S ez, then | 1lop S (1 n2 S 62.
Am log (2(0))
This implies that, for all k 2 0 and ki
log(1 73
| + 1lop S <max((1 11 11D llop. (2)
Note that
4e
| S
|| e-1 llop + || 11pp S
X
which allows to conclude that, for all lower bounded in 27
| 2 S 22.
Finally, we have the following bounds on the representations and weights at step k:
2
2llop
5
||
X
L-1
2
|
1lap 5
61
|
X
2 for l {1 +I....,4) L}.
[W
S
E
X'
Hence, an application of Theorem 3.1 proves the desired result.
Proposition 5.1. Let satisfy Assumption 4.2 Then, for any network that satisfies
approximate interpolation, i.e., 12 Y S e1.
bounded parameters, i.e., 11011 < LK +c,
the linear part WLLLT+1 satisfies
SK sY -
x(Wn.1+1) S exp
c+ Li K - 2K log
(14)
2
|X op
Proof. Let us split the parameter norm into the contribution from the nonlinear and linear layers:
10 10 1|= |10 1
We first lower bound both parts in terms of the product of the linear part W1:4++1- This then allows us bound the conditioning of W W.LL+1. We start with the nonlinear part:
L
12r, 1F S | op
II
Well S XI op ( 110ontiii
11
t=1
which implies that
1
|Z lFF
2z, 1lop
1 1nontin 112 2 L1
2 L1
11
|
Since ||Y WIat +1 Zr, S1, ei, we have
Sk (V) (K Zza) + (k(W:t11)) 1z llop +61,
22

---

Published as a conference paper at ICLR 2025
so that || llop 2
*x 4x sx t1)
and thus
sx
nontlnn 112 L1
(31)
|| llpp
Next, for the linear part, we know from Theorem 1 of Dai et al. (2021) that
Rank( (.1)
min We
Well (I- 1n)
Σ
(((A)
so that
(ank(Wrat.+1)
10 112 2 (L L1)
Σ
E 2 (L L1) + 2og WL.L (32)
where we used the fact that x म
2 logas x.
We also have following bound on the condition number no ((W2:1++))
31 (WE.L+)) <(WL:L++1)=
K
8, (WiLt+)))
[WL:L++14
S II iml
(33)
SKi W +:L1+1)
SK (W (LL:++1
(k(WL.L++1)
By combining 31
32 and
33
(after applying the log on both sides of we obtain
sx(Y)
1011 2 L1
(1-L1)K 2 log W1ttt++14
sK( L:L1
(WLL, +1) |
SK (Y) 61 sK (W ||x
2 L1
(L L1)K 2lo ((WL, + 2K log sK (WL:1++1).
The above is lower bounded by the minimum over all possible choices of s WW This minimum would be attained at sæ (WL:LI+1) = K- bn (x thus leading to the lower bound
|1018 2 LK + 2 (WLL+L)))) - L log K + 2K log 8K *K(Y) 6
|| olp
This implies
8K (Y) c1
(ogx(W).L,+1) S 1101/ LK L log
2K log
|
sct LiKlgkk K - 2K log SK (Y)
-1
|X
Theorem 5.2. Let a satisfy Assumption 4.2
Assume there exist parameters of the nonlinear part ||0 112 = c. Then, at any global minimizer of the
Onottin = = such that ZL, = Y and
regularized loss (x) = I| Zz 2|
|101 with X S KT4e we have
2
|| lop
L,K K + To K Kog (logK)) K)).
x(W1a111)5 <
exp
(c
(Y
(15)
||X lop
x(WL) <
exp
(c LIK + Li log K)
(Y -1
2L2
This implies that the bounds on NCI, NC2 and NC3 and respectively, hold with x(Wr) upper bounded above and 62 0.
23

---

Published as a conference paper at ICLR 2025
Proof. By assumption, there parameters of the nonlinear part nontin such that the representation at the end the nonlinear layers already matches the outputs ZL, = Y with finite parameter norm [|nnnntnn 11 = We can now build deeper network by simply setting all the linear layers to be the K K-dimensional identity, that ||0ten K L2, leading
total parameter norm of |1012 = c c L2. Since the outputs matches the labels exactly, we have that the regularized loss bounded by L2 + c)/2. This implies that any global minimizer 0 satisfies
|| Z S MKL2+c))
1|0-12 KL2 +
Choosing A < Kly+c' we obtain that the global minimizer 0 must satisfy the assumptions of which readily gives the desired upper bound on
Proposition 5.1
To finish the proof, we show that all critical points of the regularized loss have balanced linear layers, so that we may relate the conditioning of WT to that of the product WL.L++1- At any critical point, the gradient w.r.f. to any weight matrix amongst the linear layers must be zero, that is
0 = V w. (+) WT.. W - Y)Z W W
W + XWe.
( 2, WT for all linear layers l. Balancedness This implies that We = WT 1241 - then follows directly
W - WLe+1(Z V) WEL++1 = We+1- X
Finally, the balancedness implies that
( |X
122
K (WE) x(W x(W
(sx(Y)) (sx(Y)-a)
exp
(c
I, + Tn log K))
2
Proposition 5.3. For any network that satisfies
VoTi|ZZA'
bounded NTK, iee. ||0|| maxA
S Cl2
approximate interpolation, i.e., 12 Y11F S111
bounded weights, i.e., ||Wellop 5r,
for any M < L2, there is E {L1 + 1, ......,,,, 1,...,, L1 + M) such that x k(W..)) S Furthermore, any network that satisfies approximate interpolation and bounded weights is such that
VCL.Kr (s)
|0| (Yx(Y)- L2.
(17)
K272
Proof. As the NTK is sum over all layers, we can lower bound it by the contribution of the linear layers only. Formally, we have
V 111 2 Ve... Tr/Z,AT115
We-nit A WL
t=
Let U2...., UK N be K orthonormal vectors that span the preimage of Zr (Z₁ is rank K as long 6j < SK (Y)), and Ei....,cg cx RK be the standard basis of K We can then sum over K² possible choices of matrices A = e.0 to obtain
K
k22 121101lop N
Σ
Vole
Σ
Σ
We-1:L,12 Zz, vy By
W WL.e+1
1
1
24

---

Published as a conference paper at ICLR 2025
We know the fact that PimWT. W1-1-1+122, obtain the lower bound
WisWe.WW-1stt224 Wi to
K
K
1 2
Σ
Σ
W 2.,z1.19 Wr.e+1 W
i t=h1+1
= Σ
| Lit Zz 1 Wiee+lll
(=1+1
L
W
(Zr)
Σ
|Wellop
(=b++1
(sk (Y) 9)
Σ
n(WL.e)
p2
8=
b1+1
This then implies that
(sx(Y) E1)
CL2 2 ||0|lop
Σ x(WL,))2
K2p²
t=
Let assume by contradiction that for all the layers l {L1 + 1, 1..... In + M} have ,2, x(W),)) (1) then
(s (Y) (1) Σ x(WL.))² CL2.
K2p²
1=+++
which yields contradiction. Therefore, there must be layer within the layers & {L1 + 1,...,I, L, + M} such that k(WL..)2 & M(xx())-a)²
CL,K²,2
Furthermore, since x(W...) 2 1, we know that
(Yx -1) L E1)
| op
K2T2
C TECHNICAL RESULTS
Lemma C.I. Let b > 1 be s.t. ||X b for all Then, inside the set of parameters with bounded weights 11Wellop I Ti for all E [L] and ra 2 1, the gradient of the loss VCo(0) is
5/3 (T))=155) Lipsshii
Proof. Consider two parameters (We) 0' (W))- and let Zm, Zi be the corresponding outputs at layer m. Then, we obtain the following telescopic sum:
(2m)) (Zi)i
Σ (W Www.ooo (X)
(34)
E=1
(W
W/woo/oooo////oooooo W.))x..),
so that
m
[l(Zm)i: (2/))il12 <
Σ
r3 I|,
Will ||x
t=1
(35)
L-1
L-1
< b
T m
Σ
w
w/12 < b
II
T m 0 0
j=1
t=1
1=1
25

---

Published as a conference paper at ICLR 2025
Now, the gradient equals
N &- D (Z j=
VCo(0)=
II W2Dgs
t=
where D is diagonal matrix with diagonal entries given equal to the vector o' ((2)))) for e [ and it is equal to the identity otherwise. Similarly, we define D j.a as diagonal matrix with diagonal
entries given equal to the vector d' d((())a)) for E [L,] and equal to the identity otherwise. Since o' is B-Liischitz, we have
L-1
1 S 86
II
ryvm 110 0'12-
jm1
Furthermore since o' S 1 have that || s
By summing all terms that need to be
changed, obtain
L
N
(0) - 1l(
Xxlll 1
(ZL)all2
Σ II
r3 1m
W
mith (m,
N
L Σ
X 1
(ZL)ael2
Bym 110 0'112 X
m=1
j=1
2
L-1
ΣΣ
II
ri
VT 0' 1|222
j=1
where the three terms correspond to the effect of changing the We's the De's and the Ze's respec- tively. This can then be simplified to
VC,(0) - VCo(0)112 S No²
1
Tj
ry
3/2
110 -8'1l2
jml
2
L-1
+ N
1 II
Yj
3/2 3//2//30 3110 0'112
j=1
j=1
Nb
1,3/2 110 0'112
rj
L
< 3N
II
1,5/2 110 - 0'1l2
where we use the fact that |I(Z2)xll2 S 511 Ty.
Lemma C.2. the network satisfies approximate balancedness WWE.Wenn - WoW? S fe E 11, L -1}.
bounded weights ||Wellop for E {1 + 1.......1 ,Ld,
then we have (WLWI)* Wi L:
< L (122(12211 op
where L = L. In.
Proof. We denote D3 = WoW. WT, We+1 and
E Σ
(W (WIz) i-1--1 De (WoWT)
26

---

Published as a conference paper at ICLR 2025
We have that
(WawT) = De (WWW.) (--1-1 (- We41We+ (W.W,)) 8-
De (WWW,) 1-11-1 W W We+1 Dr )
+ (Weti (Waw;))--r- t-1-2
1
= -WWT 11 +
Σ Wees) De (WoW7)
-WWet We+1) e-h + Et.
Then,
W.L, +1 WT WEI. +1
WL. ... WL -.- W7,
-WL... WL,-2 2 WL + ELI+1) "hy W,
W. WW..WW,,- WW 2
+ E1,11
= -WLW]) L2
Σ &
i= +1+1
where & = =W ... E. WT. As a result, we have
L-1
[(WW7)) 42 WL W 1 11,p Σ &
Σ
E
11+1
op
Since ||Willop S for all Ao = [1i + 1,... L}, we have | S (E
and
||Eellop ,2(6-0 - L1je22(((-11-1) = (e L1)egr
Thus, we have
-
L3 s (
WWELLIW W L:L1+1
27


Reviews:
[{"review_id": "zPXleeD4ia", "reviewer": "Reviewer_H9Md", "summary": "Dear authors, thank you for submitting your work to ICLR 2025. The paper considers the phenomenon of 'neural collapse' and attempt to extend current rather special case results to brader class of networks, specifically, a deep neural (non-linear) networks with a wide first layer, funnel architecture and several, i.e. 2+, linear layers (head) before the output. After taking several assumptions, paper shows in series of Theorems (3.1, 4.4,5.2) and Propositions (4.5, 5.1,5.3) that GD training with weight decay (under further assumptions) leads to within class variability collapse (NC1). Results are supported by experiments on MNIST and CIFAR and MLP and ResNet + MLP head.", "questions": "See Weaknesses for the most concerning questions. \nAdditionaly:\n\nQ1: Line 188: \"... aproaches 1 ...\" Shouldn't it be \"... aproaches 2\" based on (3)? Is it still sufficient for orthogonality claims (can the proof be adjusted to account for it)?\n\nQ2: Proof sketch of Theorem 4.4., lines \"306\". Why and how is are \"two phases\" guaranteed to happen during GD training?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "S1: Novelty, pushing for (more) general results on timely and attractive topic of 'neural collapse' in DNN training\nS2: Striving for a theoretically solid argumentation stating necessary assumptions (in Theorems and Propositions) and as Assumtions 4.1, 4.2, 4.3 ...\nS3: Well done Introduction positioning paper within existing works (UFM, NTK and other specific results)", "weaknesses": "W1: Accessibility of the paper for a '15 mins' reader. The main results are formulated using (Abstract, l19) 'balancedness', (Abstract, l20) 'bounded conditioning' and other only later defined phrases and makes it hard to asses the attractivity/usefullness of paper unless reading it whole. It is recommended to rework (the Abstract at least, if no Conclusions are present) to make it clear and self-sustained.\n\nW2: Quite a few assumptions are required in general. More over thay are added along the way, e.g., Assumption 4.2, $|\\sigma(x)| \\leq |x|$, etc., (which is ok, but also narrows the applicability of the final result). Some of those are very technical (especialy those in Theorems, such as in Theorem 3.1, (2),(3), (4)) but as well Assumption 4.3, Theorem 4.4. (10) and more) and with paper specific notation to learn, e.g., $\\lambda_{3 \\rightarrow L}$. It would help paper to have an thorough discussion on applicability/limitations of these assumptions. Perhaps at expense of shortening some 'proof sketches' refering them to SM?\n\nW3: 'Discussion and Conclusion' sections are not presented. This is most likely due to space constrained and will be added in case of acceptance (?). Yet, I find it impacts the paper quality negatively. Especially in a light of the previous (W2) comments, the discussion and conclusions could have brought a critical 'back to reality' summary. Idealy it would bring more intuition for results and their application. Yet, I find it very important to have such Discussion reviewed before publishing ...\n\nW4: Some references are rather inaccurately interpreted/generalized too much perhaps. For instance the lines 399-400 \"...Thankfully, instead of simply diverging, for large η (but not too large) the parameters naturally end up at the ‘edge of stability’: the top eigenvalue of the Hessian is close to $2/\\eta$ , i.e., the threshold below which GD is stable...\" from (Cohen et all. 2021). Referenced work provides only experimental evidence for a phonomenon and only approximately, i.e., for certain settings operator norm of Hessian exceeds 'stability threshold' $2/\\eta$, etc. Than approximation used on l410, especially $O(\\epsilon_1)$ is only valid if $\\nabla^2_{\\theta} Z_L$ norm is bounded, which is ok for NTK regime, but not necessarily for large learning rate regime. Or is it?\n\nW5: Following up on W4, Proposition 5.3, and other Theorem combine NTK with large learning rate regime, which sounds dangerous. Also requirement on wide first layer, suggest a NTK limit case is required. Could authors clarify a bit more on this relation?\n\nOverall, I find it to be a solid attractive paper with technically legit reasoning, taking few shortcuts (some noted above) and with missing discussion and conclusions. I suggest authors to work on alleviating weaknesses and discussing limitations to improve contributions of this interesting work significantly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5CJWmi7Ofo", "reviewer": "Reviewer_PRAd", "summary": "This paper presents an interesting theoretical advancement in understanding neural collapse in practical, end-to-end training scenarios, moving beyond the unconstrained features model. The authors provide a rigorous demonstration that neural collapse arises in networks with linear layers appended to a nonlinear backbone, given conditions of interpolation and balancedness. They show that these conditions hold for sufficiently wide networks trained with gradient descent and L2 regularization. The empirical results further support the theoretical findings, showcasing the robustness of neural collapse across various architectures and datasets.", "questions": "In Figure 3 the authors' numerical results show that non-linear layers are increasingly linear, as the depth of the non-linear part increases. Could the authors provide more insights into how this observation relates to their theoretical results and the mechanisms driving this increased linearity?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "This work provides an interesting theoretical insight on the role of training algorithm in the emergence of neural collapse, which I found especially exciting, and I think it opens up new directions for understanding the generalization properties of deep learning models.", "weaknesses": "In my opinion, this is a solid paper and I can not think of a weakness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MrPkc9RkIe", "reviewer": "Reviewer_1hwn", "summary": "This paper provides theoretical evidence for end-to-end training of true deep neural networks. This is contrary to previous works which primarily rely on the unconstrained features model. The paper provides explicit bounds on NC1, NC2 and NC3 for both globally optimal ($\\ell_2$) regularized deep neural networks as well as neural networks trained with gradient descent. The results provide new insights on the role of additional linear layers, weight decay regularization and large learning rates with respect to neural collapse.", "questions": "- What exactly is the role of these additional linear layers?  Are they only required for proving NC2 and NC3?\n- Do these results suggest that neural networks with 1  non-layer and many linear layers can also exhibit neural collapse?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "- The authors provide the first theoretical analysis of neural collapse that does not rely on the unconstrained features model.\n- The use of additional linear layers is novel and interesting technique. To the best of my knowledge such an architecture has not been studied.\n- The results apply to deep neural networks trained with gradient descent (under a particular architecture/nonstandard activation function) as well as networks which are globally optimal for the weight decay regularized objective.", "weaknesses": "- My main concern with the paper is the writing. The theoretical statements are very difficult to parse this is especially the case for eqs 2-5 in Theorem 3.1. Can the authors provide a more interpretable version of this theorem (hide constants, big-O etc.) and delay the details to the appendix?\n- I have the same concern about Theorem 4.4.\n\n### Minor\n- The proof of Theorem 3.1 in the Appendix could also be more clear. For example mentioning Weyl's inequality in the first inequality of (21)\n- In the proof of Theorem 5.2 at the Appendix (line 1264) shouldn't it be\n$$\\kappa(W_L) = \\kappa(W_{L:L_1+1})^{\\frac{1}{L_2}}$$\nnot \n$$\\kappa(W_L) = \\kappa(W_{L:L_1+1})^{\\frac{1}{L_1}}$$\n- Same concern about $L_1$ vs $L_2$ in the statement of Theorem 5.2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SBvBuRvwHz", "reviewer": "Reviewer_ggM6", "summary": "The authors identify a set of conditions under which various neural collapse phenomena provably occur in deep neural nets.  They consider neural nets which have a sequence of nonlinear layers and then a sequence of linear layers.  They find that approximate interpolation, weight balanced-ness, and boundedness suffice for deriving various neural collapse phenomena.  They then show that GD on networks with a \"pyramidal\" overparameterized topology (i.e., first width is >= number samples, remaining widths are decreasing), under suitable initialization and regularization, allow for one of the neural collapse phenomena to hold.  They then identify conditions (near-interpolation and small-norm) which ensure that global minimizers of the loss can satisfy all of the neural collapse phenomena.   Finally, they look at the neural collapse phenomena  from the perspective of the edge of stability via an analysis of the properties of the hessian under EoS assumptions.", "questions": "1. Are there no assumptions on the activation function $\\sigma$ needed for Theorem 3.1?  None are stated in Section 3, although assumptino 4.2 appears later.  It seems odd that a potentially discontinuous and horribly behaved activation function would be permitted (I imagine Lipschitz is required?)\n\n2. Assumption 4.2 seems more like a smooth leaky relu, which has appeared in prior work [Chatterjee arXiv:2203.16462, Frei et al COLT 2022]\n\n3. The authors talk about (NC1)-(NC3), what about (NC4) from the original Papyan paper?\n\n4. Are there any important differences between the proof of Theorem 4.4 and the proofs/results in Nguyen and Mondelli?   I'm not familiar with that work, but it seems like it's not super necessary to have details about optimization (e.g. the PL inequality) in the main section of the paper, it distracts from what I think are the more interesting and stronger results elswhere in it (it would also give additional space to have a proper conclusion and outline what the authors think are future directions).   I am assuming balancedness didn't appear in the prior analysis but does here.  A few sentences describing high-level differences between proofs and ideas would be helpful.  \n\n5. Also, to be clear, the optimization analysis in Sec 4 is in the NTK regime right?  I didn't see this explicitly mentioned but it should be if it is known.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The authors provide a novel characterization of conditions under which neural collapse can provably occur.  I am not an expert on the NC theory literature, but to my knowledge, no prior work identified balancedness and interpolation as being key factors which can allow for this to occur, and this is a pretty strong finding.   Since the balancedness of post-fixed linear layers is the strongest condition (as most common neural nets do not have multiple linear layers at the end, only a single one), the findings in Section 5 about how boundedness and interpolation can suffice for NC2+NC3 are also a nice addition.  The numerical findings nicely complement their theoretical ones.", "weaknesses": "There aren't any serious weaknesses to me.  The strongest assumptions--namely the need for multiple linear layers in order for the NC phenomenon to occur via the main theorem--seem necessary per experiments in Figure 1.  So it seems that these assumptions are strong for a fundamental reason.\n\nThe pyramidal structure assumption is a bit odd/strong, but only seems needed for the optimization result, which I don't think is the central point of the paper.  \n\nThe authors don't provide a conclusion or discussion section, and it would have been useful to have comments from the authors about what they think are the weakest parts of their work and what work should be prioritized in the future.  I think they could get some additional space by removing some of the details about optimization from Sec 4 since I don't think they're super important.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eS32vck9mm", "reviewer": "Reviewer_39ne", "summary": "This paper studies neural collapse phenomenon in training wide neural networks. The first result in this work establishes some general conditions (interpolation, balancedness between layers, well-conditioned-ness of weight matrices) such that NC1, NC2, and NC3 can hold. The second result considers training a wide neural network with gradient descent such that the aforementioned conditions hold after training, which implies neural collapse can happen after training.", "questions": "I am wondering whether the layer balanced-ness property after training can be proved as a direct consequence of [1] in the author's setting.\n\n[1] Du, Simon S., Wei Hu, and Jason D. Lee. \"Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced.\" Advances in neural information processing systems 31 (2018).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work goes beyond the conventional unconstrained feature model (which many previous works have worked on) and proves that neural collapse can happen after training for full-connected neural networks satisfying some architectural conditions such as pyramidal topology and smooth activation. The conditions established in Theorem 3.1 of this work under which neural collapse can happen seem very reasonable. Indeed, later, the authors proved that those conditions can be satisfied by training a neural network via gradient descent which automatically implies that neural collapse can happen after training. \n\nI am not very familiar with the prior works along this line of research such as (Kothapalli & Tirer, 2024), (Hong & Ling, 2024) and many others mentioned in the introduction and related work. Based on my knowledge on neural collapse, I think this work made some interesting contributions towards understanding this phenomenon.", "weaknesses": "1. The analysis critically relies on the fact that the last two layers of the neural network are linear. I can definitely see this condition makes the problem a lot easier to analyze. I am wondering how hard it is to remove such restrictions. \n2. It seems the analysis of Theorem 4.4 relies on the neural network in the NTK regime, as the pyramidal topology assumption has appeared in previous works such as (Nguyen & Mondelli, 2020). I don't regard this as a major weakness even if it turned out to be true that the networks are in the NTK regime given the contribution of this work, however, I do appreciate clarification on this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)