Neural Collapse Beyond the Unconstrained Features Model:  Landscape, Dynamics, and Generalization in the Mean-Field Regime
Authors: Diyuan Wu
Marco Mondelli
Venue: ICML 2025 spotlightposter (2025)
Keywords: neural collapse
mean-field analysis
gradient flow
generalization error
loss landscape
Primary area: deep_learning->theory
TL;DR: We prove that NC1 (vanishing within-class variability) holds when training a class of 3-layer networks via gradient flow, due to loss landscape properties; we further show co-occurrence of NC1 and small test error for certain data distributions.

Abstract: Neural Collapse is a phenomenon where the last-layer representations of a well-trained neural network converge to a highly structured geometry. In this paper, we focus on its first (and most basic) property, known as NC1: the within-class variability vanishes. 
While prior theoretical studies establish the occurrence of NC1 via the data-agnostic unconstrained features model, our work adopts a data-specific perspective, analyzing NC1 in a three-layer neural network, with the first two layers operating in the mean-field regime and followed by a linear layer. In particular, we establish a fundamental connection between NC1 and the loss landscape: we prove that points with small empirical loss and gradient norm (thus, close to being stationary) approximately satisfy NC1, and the closeness to NC1 is 
controlled by the residual loss and gradient norm. We then show that (i) gradient flow on the mean squared error converges to NC1 solutions with small empirical loss, and (ii) for well-separated data distributions, both NC1 and vanishing test loss are achieved simultaneously. This aligns with the empirical observation that NC1 emerges during training while models attain near-zero test error.  Overall, our results demonstrate that NC1 arises from gradient training due to the properties of the loss landscape, and they show the co-occurrence of NC1 and small test error for certain data distributions.

Neural Collapse beyond the Unconstrained Features Model: Landscape, Dynamics, and Generalization in the Mcan-Field Regime
Diyuan Wu
Marco Mondelli
Abstract
stages of training where: (i) the penultimate-layer features converge to their respective class means (NC1), (ii) these class means form an equiangular tight frame (ETF) or an orthogonal frame (NC2), and (iti) the columns of the final layer's classifier matrix similarly form an ETF or orthogonal frame, implementing nearest class-meat decision rule on the penultimate-layer features (NC3). A popular line of the- oretical research has investigated the occurrence of NC via the unconstrained features model (UFM), see (Fang et al., 2021a: Han et al., 2022; Mixon et al., 2022) and the discus- sion in Section 2. In this framework, the penultimate- layer features are treated as free optimization variables, leading to benign loss landscape for the resulting optimization prob- lem. The primary justification for adopting the UFM is that the complex feature- layers encountered in practice are approximated by universal learner. While the UFM provides an intriguing theoretical perspective on NC. it has notable limitations. In particular, it neglects the dependence on the data distribution, rendering it unsuitable for theoreti- cally analyzing the relationship between NC during training and the test error (Hui al., 2022). Furthermore, the train- ing dynamics under the UFM framework is not equivalent to the actual training dynamics of neural networks, which makes it challenging to investigate the occurrence of NC from dynamical perspective.
Neural Collapse is phenomenon where the last- layer representations of well-trained neural net- work converge to highly structured geometry. In this paper, we focus on its first (and most ba- sic) property, known as NCI: the within-class variability vanishes. While prior theoretical stud- ies establish the occurrence of NC1 via the data- agnostic unconstrained features model, our work adopts data-specific perspective, analyzing NC1 in a three-layer neural network, with the first two layers operating in the mean-field regime and fol- lowed by a linear layer. In particular, we establish fundamental connection between NC1 and the loss landscape: we prove that points with small empirical loss and gradient norm (thus, close being stationary) approximately satisfy NCI, and the closeness to NC1 is controlled by the residual loss and gradient norm. We then show that (i) gra- dient flow on the mean squared error converges to NCI solutions with small empirical loss, and (ii) for well-separated data distributions, both NC1 and vanishing test loss are achieved simultane- ously. This aligns with the empirical observation that NC1 emerges during training while models attain near-zero test error. Overall, our results demonstrate that NCI arises from gradient train- ing due to the properties of the loss landscape, and they show the co-occurrence of NC1 and small test error for certain data distributions.
To address the limitations of UFM, we consider training three-layer network via gradient flow on the standard mean squared error (MSE) loss. Specifically, we employ two-layer neural network in the mean-fiiel regime (Mei et al., 2018) as the feature-learning component, and then concatenate it with a linear layer as the final predictor. Our main results both (i) establish sufficient conditions on the loss landscape for the first - and most basic - property of neural collapse, i.e., NC1, to hold, and (äi) show that such conditions are fact satisfied by training the architecture above. This differentiates our paper from recent studies aiming to theoretically explain the NC phenomenon beyond unconstrained features, as existing work either provides only sufficient conditions for NC to occur (Seleznova et al., 2024), focuses on the NTK regime (Jacot et al., 2024), relies on specific training algorithms (Beaglehole al., 2024) or on specific regularization (Hong Ling, 2024a), see Section 2 for discussion of related work. Specifically,
1. Introduction
Neural Collapse (NC), first identified by Papyan et al. (2020), describes phenomenon observed during the final
'Institute of Science and Technology Austria Correspondence to: Marco Mondelli
(ISTA), Klosterneuburg. Austria. Diyuan Wu stlyyan.wwQist.ac...o <marro.mondelieissaa.cco
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Neural Collapse beyond the Unconstrainted Features Model
contributions are summarized below:
two-layer neural networks in both the NNGP and the NTK limit, proving neural collapse for 1-dimensional Gaus- sian data. Hong & Ling (2024a) study NC for shallow and deep neural networks, also characterizing the generalization error. However, they regularize the loss by the 12-norm of the features rather than the weights, which is different from the weight decay used in practice, Jacot et al. (2024) establish the occurrence of NC for deep neural networks with multiple linear layers, given balancedness assump- tion on all the linear layers; in addition, they also prove that balancedness is achieved via gradient descent training using NTK tools. Compared to (Jacot et al., 2024), our proof does not rely on any balancedness condition, and it only requires the gradient norm to be small, which is naturally achievable via gradient flow. In fact, the stationary points which our results apply may not be balanced, see the discussions at the end of Section 4.1 and 4.2.
First, we connect the emergence of NC1, i.e., the fact that the within-class variability vanishes, with properties of the loss landscape: we show that all approximately stationary points with small empirical loss are roughly NC1 solutions, and the degree to which the within-class variability vanishes is controlled by gradient norm and loss. This implies the prevalence of NC1 during training, as practical training procedures typically converge to such points with small gradient and loss.
Next, we prove that gradient flow a three-layer net- work operating in the weean-field regime satisfies the two conditions above (small gradient norm and small loss) and, therefore, it converges to an NCI solution. While achieving approximately stationary points is expected, the primary challenge lies in controlling the empirical loss due to the model's 's non-convex nature.
Mean-field analysis for networks with more than two lay- ers. While the properties of the loss landscape and training dynamics of two-layer neural networks in the mean- field regime have been extensively studied (Mei et al., 2018; Chen et al., 2020; Javanmard et al., 2020; Shevchenko et al., 2022; Hu et al., 2021; Suzuki et al., 2024a; Takakura & Suzuki, 2024), networks with more than two layers still prove to be challenging to analyze. Prior works (Lu et al., 2020; Araújo et al., 2019; Shevchenko & Mondelli, 2020; Fang ct al., 2021b; Pham & Nguyen, 2021; Nguyen & Pham, 2023) have investigated the mean-field regime for deep neu- ral networks, where the widths of all layers tend to infinity. In contrast, we let only the width of the first layer tend to infinity, while the width of the second layer remains of con- stant order. A closely related paper is by Kim & Suzuki (2024), which studies the in-context loss landscape of a two- layer linear transformer with a formulation similar to ours. However, the global convergence results in (Kim & Suzuki, 2024) rely on assumptions such as absence of weight decay, taking a two time-scale limit, and using a birth-death pro- cess (rather than the widely-used gradient flow), which are not applicable to our setting.
Finally, we show that, for certain well-separated data distributions, it is possible to achieve NC1 during train- ing as well as vanishing test error, which corroborates the empirical finding that NC1 and strong generalization occur simultaneously.
2. Related Work
Neural collapse: UFM and beyond. The introduction of the UFM in (Mixon et al., 2022; Fang al., 202 la) has prompted line of work studying the emergence of neural collapse for that model. Specifically, Zhou et al. (2022) focus on the two-layer UFM model, showing that all its stationary points satisfy neural collapse. Han et al. (2022) prove convergence of gradient flow on UFM NC solutions. Tirer& & Bruna (2022) demonstrate that the global minimizers also satisfy neural collapse when the UFM has multiple linear layers or it incorporates the ReLU activation. Súkeník et al. (2023) extend the results to a deep UFM model for binary classification. Súkeník et al. (2024) then show that, for the deep UFM and multi-class classification, all the global optima still satisfy NC1, but not NC2 and NC3, due to the low-rank bias of the model. We also refer to (Kothapalli, 2023) for a rather recent and detailed review.
3. Problem Setting
Going beyond the UFM, Seleznova et al. (2024) study the connection between NC and the neural tangent kernel (NTK), showing NC under certain block structure assump- tions on the NTK matrix. However, the occurrence of such a block structure during training is unclear. Beaglchole et al. (2024) establish NC both empirically and theoretically for Deep Recursive Feature Machine training a method that constructs a neural network by iteratively mapping the data through the average gradient outer product and then apply- ing an untrained random feature map. Pan & Cao (2023) consider classification with cross-entropy loss, providing a quantitative bound for NC. Kothapalli & Tirer (2024) focus
Notation. Given an integer n, we use the shorthand [<] == (I....,n) n). Given a vector v E R&d let v[i] be its i-th entry and Diag(v) E Rdddd the diagonal matrix with e on the diagonal. Let 1, E Rdd be the all-one vector of dimension d. Given matrix A, let [A] be its (i. element. We denote by || ||r, || |lop the Frobenius and op- erator norms of a matrix, and by (A, B) Tr B} the Frobenius inner product. Let @ be the Kronecker product and vec(-) the vectorization of the matrix obtained by stack- ing columns. Given vector valued function f : Rdd - Rd we denote by V f : Rdd - R its divergence. Given a real valued function g, denote by || = sup [g| its infinity
2

---

Neural Collapse beyond the Unconstrainted Features Model
norm. Let (2() be the space of probability measures on R. with finite second moment, and Wi ();,, W2(:.-) the Wasserstein-1 and -2 metrics, respectively.
We also denote X [.1..... Fn. E R We consider regularized problem with 1,2 and entropy regularization, denoting the regularized loss and the free energy
Three-layer fully connected neural networks. We start by defining the following wndinite-width neural network as feature-learning layer:
de E (11011), (4) (x W) = Lo(e, dw IW
2
(l) where a E IR, I E Rd, and p = Law(a, u) is the joint measure of a, The network is parameterized by a prob- ability distribution p E (R(R+t)), and it represents the mean-field limit of the finite-width two-layer network below (Mei et al., 2018):
hp(x)= = Eplao(u [u]],
E) - - W) + Aw
|W1
(5)
E 11 + و B Blg pl,
where the entropic regularization is explicitly added to per- form noisy gradient flow, see (20). While we focus on balanced classification for technical clarity and brevity, our results extend to unbalanced classification (as considered e.g. in (Thrampoulidis et al., 2022; Hong & Ling, 2024b)) and regression (as considered e.g. in (Andriopoulos et al., 2024)) with minimal modifications.
N
Σ aajolu x), 8g (a,, uz) Ltd. P. (2)
hy (2) = N 2=1
For technical convenience, throughout this paper we directly consider the width network (1). In fact, its differ- ence with the finite-width counterpart (2) can be readily bounded using results from (Mei et al., 2018; 2019).
Neural collapse metric. We focus on the first property of neural collapse and, given a feature matrix H E R we consider the following metrie of NCI as the ratio between in-class variance and total variance:
Next, we cascade linear layer, obtaining a three-layer neural network as follows:
f(z p. W) W hp(z), TWThp(z),
(3)
NCI = Thf( -M) (H -H)
where W E ERX4 and y E TR is a (constant) multiplicative factor. Throughout the paper, we will refer to f in (3) as the predictor. Neural networks with two linear layers in the end are also studied by Jacot al. (2024), and adding multiplicative factor y is proposed by Chen et al. (2020) to guarantee the global convergence of the dynamics.
Trf
where H E 12xx" is the matrix of centered features and Me E R the matrix of in-class means, defined as
The motivation for considering this model is to explore how training data affects the emergence of neural collapse. While previous studies on UFM offer insights into neural collapse, their key limitation lies in the disregard for the influence of training data. The primary justification for using the UFM is that it functions universal learner, thus emulating the complex feature learning layers encountered in practice. The three-layer network in the mean-field regime defined in (3) not only performs feature learning by taking into account the training data, but the feature layer in (1) is also recognized as a universal learner (Ma et al., 2022).
M, = 1 11 n: - Mg. M. = - Hy n
m
In words, if .NC(HH) is small, the within-class variability negligible compared to the overall variability across classes, capturing the closeness of features to respective class means.
4. Within-class Variability Collapse during Training
4.1. Sufficient Conditions for NC1
q-class balanced classification. We consider a q-class balanced classification problem, with each class having m data points. We denote by qm the total number of training samples and assume that p 2 q. The empirical loss function is given by
Throughout the paper, we make the following assumptions that are mild and standard in the related literature, see e.g. (Mei et al., 2018; Chen et al., 2020; Suzuki et al., 2024a).
Assumption 1. (A1) Regularity of the initialization: We initialize the training algorithm with Wo such that Wo Wa Ly and Po = NO Ipta).
1
Co(p, W)
- 1 Hp -0--Y 1|?
2n
(A2) Boundedness of the data: for all i, |ævll2 S
where
H, = (1)(2))... 1p (x() E Roxn
(A3) Regularity of the activation function: I|o()))lo [=o((=)]loo [=))())]ooo [|a" (x) S Ci for some universal constant C1.
Y = (e....ey..... ......eg E Roxn
mn columns
3

---

Neural Collapse beyond the Unconstrainted Features Model
We remark that the ReLU activation does not satisfy the regularity conditions above, but we still expect our results to hold by taking the limit of sequence of approximations of ReLU.
where, with an abuse of notation, the term (es) indicates that the norm of the vector on the LHS is at most of order
eg. By plugging (11) into Hp = -Bp jao (u' X)], we obtain
K K,(X, X) +0(es). (12)
Hp p -p- >>WWTWWWWTTT YY)
We now define an eg-stationary point of the free energy.
n
Definition 4.1. We say that (p, W') is an etsstaionnayy point of (p, W) w.r.t p if the following holds:
Note that (12) is a linear equation in Hp. Thus, by solving it explicitly and tracking the error in es, we obtain (6). Finally, the crux of the argument for (9) is to use again stationarity to show that (up to an error of order es)
E
Vo To En ()
S
(K' & W') (K' & (W' W') + >>Ing)) -1 vec(Y)
(13)
Here, we recall that, given functional G: 9) + R. its first variation at p is the function ())(()) .R. - IR
- + (p,
such that, for all p' E (R)))
where K" := KX(XXX))/m and
and W := W.
[ - 5p G(O) - -p)0 p)d0= lim
G((1 c)p + cp' G(())
Note that Y decomposition of Hp in (9), satisfies NC1. Indeed, Y is the one-hot vector of labels and, thus, for two data point Tt, Tj in the same class k; we have yi yj. which implies that (WWTW) w) = W W(W "y. The second term Eg in the decomposition (9) is small, as long as cs and C. (p, W) are small. Hence, the key question is whether we can achieve a nearly stationary point having small loss (p, W) via a certain training dynamics, which is addressed in the next sections.
i.e., the first term in the
6
The result below (proved in Appendix B.1) characterizes the feature Hp at any es-stationary point.
Theorem 4.2. Under Assumption 1, for any cg-stationary point (p. W), we have the following characterization of the learned feature:
ve)(H,) X)
@ W
n
X)
@ (W (WTW) W) + XpInn) ) vec(Y)
As the result in (9) requires W not to be too ill-conditioned, we now prove that this is the case, as long as the regulariza- tion terms Aw, Ap 3-1 and the regularized loss Zx," are sufficiently small.
+ Ep,
(6) Lemma 4.3. Let Ap = 195-1, Aw 2 where Xo Xby are universal constants. Fix any N 0, 0 < Eo S
where
[Ei(s Api = 4n o (W)) and the kernel K K2(X, X) E R.n induced by p is
1/2, and assume that
(7)
B 2 max e
4a log
e ea log((2))) (2C3nB(29)-1)
X) TX X)).
(8)
4q
6A((gB32
As a consequence, if W is non-singular, we have
n
(14)
Hp --=-~~10 11----WW''" - + Pelcs W), (9)
for some constant B that doesn't I depend on B. Suppose further that (p, W) is any point such that
where
[E2sc V.P. W)ll (A C mmaa
max (W) + xp²)
Ginez+ + 2n "L, W) + 2x,²² " CW)2CPne3
Cx W) S B3-
(15)
mmin (2)2
Then, we have that
(10)
Omin (W) 2 g-
Proof sketch. As (p, W) is cs-stationary, the following ex- pression for a holds almost surely w.r.t. the measure p:
m 22) < | < ( 3)
71 W(NW H,- - Y)((XTu)
The proof is by contradiction. Suppose that W has a small singular value, then the projection of H, in the correspond- ing left singular space of W needs to be large, since the
a+ +
(11)
n
+ 1-18-1 Va = O(Es),
4

---

Neural Collapse beyond the Unconstrainted Features Model
regularized loss is small and Y is isotropic. However, large component of H, in subspace will in turn lead to large regularized loss due to the second-moment moment regularization term. The complete argument is deferred to Appendix B.2.
A 1 and Xw N -1 for universal constants
Ap Ap. Then, as WW is of rank < p, (19) gives that, for any constant ||Eplaa]] cww llop 2 (A2)-1
We complement the theoretical result in Lemma 4.5 with numerical simulations, discussed at the end of Section 4.2, showing that for there are settings such that gradient-basee training over standard datasets (MNIST, CIFAR-100) the neural network achieves NC1 without converging bal- anced solution.
Next, we compute the NC1 metric induced by Theorem 4.2. Corollary 4.4. Consider the setting of Theorem 4.2 and assume that
(9-1)1 1)n
E2
(16)
80 maxe (W)
q
where Eg_ == Ee 3 W) is bounded as in (10).
4.2. Achieving NC1 via Gradient-based Training
Then, for any e stattitinnryy point (p, W) with non-singular W, we have
From Theorem 4.2 and Corollary 4.4, we know that NC1 is achieved at any es-stationary point w.r.t. p having small empirical loss. We now consider training and W with gradient flow, i.e..
16 Es (g- 1)m
(CI(H,)) S
(17)
E
2 manx
(W)
dW, = W WW Wi)dt; 5 d0, = Ve
The proof of Corollary 4.4 is direct calculation, and it is provided in Appendix B.3. Note that, in the setting of Lemma 4.3, we have that
(20)
8p Cr,ale (.))d, + 228
1
(q 1)m- 2((1og (0olog/))-0) a)
and we show that, having trained long enough, one ensures that both 6ş and the empirical loss are sufficiently small.
8 () max
Eg = 0(3 1000 2 3)a + 140200
200 (log B)th ESS)
The convergence to an eg-stationary point with arbitrary small es is a direct consequence of the fact that, under gradient flow, the gradient norm vanishes.
(18)
Now, let us pick a sufficiently small 8-1 (corresponding to small regularization) and then a sufficiently small 6g
Lemma 4.6. Under Assumption fix 3,7 > 0 and consider an initialization (po, Wo) with finite free energy. For t 2 O, let
(corresponding to reaching a stationary point). Then, (18) implies that (16) holds and the upper bound on the NC1 metric in (17) vanishes.
Vo 5 &n (p. (e
=pp
6p
Imbalancedness of stationary point. The recent work by Jacot al. (2024) shows that, for any network with at least two consecutive linear layers in the end, sufficiently small loss and approximate balancedness of the linear layers suffice to guarantee NC1. Our network defined in (3) has two final linear layers, but due to the entropic regularization, all stationary points of the free energy are not balanced, which means that the techniques in (Jacot et al., 2024) can- not be applied to our setup. To demonstrate this, we prove in Appendix B.4 the following result.
2
which is equivalent to (pt: W.) being an es--statooary point. Then, for any es 0, there exists T(es) > 0 s.t. for all > T(cs) except a finite Lebesgue measure set,
é' S Se.
The proof of Lemma 4.6 is provided in Appendix C.I. This result directly implies that lim inf c3 0.
Next, Theorem 4.8 shows that, by picking large enough y and training long enough, we achieve 0(3-1) empirical loss. This requires the following mild assumptions that imply the positive definiteness of the kernel K, in (8) at initialization, as showed in Lemma 4.7.
Lemma 4.5. Let (p, W) be a stationary point of the free energy. i...
Vo Voggenlo En W) 0, V. W) = -0, 5p
Assumption 2. Assume (2k) f 0 for all k >> 0 and that there exist s E [d] s.t. (i) 23[s] # # for all + E [m]. and (ii) x, [s] o for all i E [n].
with [a < co. Then, any stationary point satisfies
AaE 1 - Xw WW = =3-11pp
(19)
In words, the activation function is required to be smooth and have non-zero even derivatives, which is satisfied by e.g. tanh or the sigmoid (also fulfilling Assumption 1). As for the training data, we assume that it is non-degenerate and not parallel, which holds for most practical data sets.
The result in (19) implies that the network cannot be bal- anced, i.c., Ep(ace |aa"" cannot be proportional to WWT In
fact, assume that Ao and dw are of same order as 8 i.e.,
5

---

Neural Collapse beyond the Unconstrainted Features Model
The next technical lemma, which comes from (Nguyen & Mondelli, 2020, Lemma 3.4)', shows the required positive definiteness of the kernel.
We remark that also provide a convergence rate for the loss in (24). The problem is challenging due to its non- convex nature, and to our best knowledge, explicit rate of convergence is known in the mean-field regime beyond the two-layer case (which has a convex free energy).
Lemma 4.7. Under Assumption 2, let K(X,X) Eueny (o T u))((u X)], where "'a- = N0, Ia). Then,
A = Amin (X(, X)) > O.
Comparison with related work. While the strategy de- scribed above is motivated by and similar to that used in (Chen et al., 2020, Theorem 4.4), its technical implemen- tation differs, due to differences in the problem setting. In fact, Chen et al. (2020) consider two-layer neural networks whose free-energy landscape is strongly convex in p. In contrast, the presence of the last linear layer W implies that the free-energy landscape is non-convex in (p, W), which makes it less obvious that gradient flow converges to low free-energy manifold. As consequence, Chen et al. (2020) show that the gradient flow dynamics stays in ball around initialization with a certain radius for infinitely long time, In contrast, in our case, the gradient flow dynamics stays in the ball only for finite time, but this finite time suffices to ensure small enough free energy.
We are now ready to state our result showing the conver- gence of the empirical loss to a low loss manifold by running gradient flow for long enough time.
Theorem 4.8. Let Assumptions 1, hold, set Xo = dw 3-1 and
7> C3. to 3Css.
(21)
Then, for any B and any! 1. 2 L0, we have
(p,n Wi) S 8 C4,
(22)
where Cs are constants that depends on n, d,p, A, but not on B.
The expression of C3, C2, C3 is provided in Theorem C.1, whose statement and proof are in Appendix C.2. We note that the choice >p = Aw = 8 is only for technical convenience, what matters here is that Ap, dw = 0(8
Finally, the combination of Theorem 4.8, Lemma 4.6 and Corollary 4.4 gives that NCI provably holds under gradient flow training.
Proof sketch. We start by defining a first-hitting time t, = mini inf(t: WW, W, - Wo Wollop > Rw h, (23) inf(t: : (pello))) > R,Y).
Corollary 4.9. Consider the setting of Theorem 4.8 and, for
any 0 < 60 < 1, let
8 > max (2C7nC1) (44) (#))
where (-I(- denotes the KL divergence. Intuitively, (23) means that, for t < t., the gradient flow stays in ball around the initialization. The crux of the argument is to show that, with a suitable choice of Rw, Rp, the loss becomes small before the dynamics has exited the ball.
1 64(gCa)², (640C3 C#
)
3)
50
where 7: as chosen as in (21). Then, there exists T(3) O s.t. for all t >> max(T to} except a finite Lebesgue measure set,
To do so, we first prove in Lemma C.2 that, for < ta,
(pe, W.) I exp(->2A1)) + 23- A2,
(24)
NC1(H,)) S 80.
for some A1, Ag that do not depend on 7, 3 (but only on A,,P, d, nt). This implies that the empirical loss converges exponentially fast (in t) to an error of order 8-2, as long as the gradient flow is inside the ball. We then show that, by picking a proper y, f. is large enough so that the first term in (24) is of order 3-2 for some 40 < l..
The proof of Corollary 4.9 is deferred to Appendix C.3, and the result implies that lim inf NC1 S 80- Corollary 4.9 implies that, the three-layer model (almost) always achieve NC1 solution, for long enough training, which explains the prevalence of neural collapse in practice.
Finally, by combining the upper bound on the empirical loss with the fact that D (Pt, Po). (W.,. I|? are bounded by a constant independent of B, we obtain that (pt,, W.,) ) = 0(8-1). Thus, since the free energy decreases along the gradient flow, the upper bound in (22) follows from the relationship between empirical loss and free energy proved in Lemma A.3.
Imbalancedness after gradient-based training. The nu- merical results of Figure 1 show that, even if the solution obtained via gradient descent is not balanced, its training loss and gradient norm are still small and, therefore, as pre- dicted by our analysis, it satisfies NCI. As a normalized balancedness measure, we use
Note that the lemma is contained in the v1 of the paper, avail- able on rxiv.
NB W) = ||Epaa J
c. WWT lop
(25)
|Ep /aa" 11lop
6

---

Neural Collapse beyond the Unconstrainted Features Model
NCS
Aanage NE1
10
Volues
Valuea
Spochs
Spochs
(a) MNIST
(b) CIFAR-100
Figure 1: Average training loss (blue), NC1 (orange) and gradient norm (green) during SGD training. We report the average for 4 independent experiments, as well as the confidence interval at standard deviation.
with
is d 2048, the number of neurons in the first layer is N 16384, and the number of neurons in the second one is p = 64. We train the model with noisy SGD of batch size 64, pick weight decay Aw X = 3 = 10-4 and learning rate E {0.001, 0.0001). using the smaller (larger) learning rate for the first (second) half of the epochs. As in the previous case, the network does not achieve balancedness throughout training: at convergence, N B(p, W) (0.5386, 0.2779, 0.4141, 0.5257} in the 4 independent experiments; see also Figure 2b in Appendix E for a plot of the metric in (25) as a function of the number of training epochs. Nevertheless, the NC1 metric decreases with the loss and the gradient norm, reaching a small value at the end of training, see Figure 1b.
C. arg min ||Eplaa 11 cWWT 1.2.
(26)
This captures the extent to which E, J and WWT are proportional. We then train the three-layer neural network (z) = W hy (r), where the output h (z) of the first two layers is given by (2). Specifically, we consider the following two settings.
Setting (a): MNIST. We relabel the dataset into q = classes taking the original label modulo 3, and randomly pick 10000 samples in each new class for training. The input dimension is d = 784, the number of neurons in the first layer N = 6272, and the number of neurons in the second one is = 16. We train the model with SGD of batch size 64 and learning rate 7) E {0.001, 0.01), using the smaller (larger) learning rate for the first (second) half of the epochs. We pick weight decay Xw Ap 10 44 but add no noise (8-1=0) =0-100 0) and fix the last linear layer at initialization, which produces an imbalanced network. In fact, at convergence, B(p W) = {0.8093, 0.7149, 0.5895, 0.8531) } in our independent experiments. We also plot the evolution of the normalized balancedness metric in (25) as function of the number of training epochs in Figure 2a of Appendix E, which shows that the network does not achieve balancedness throughout training. However, even if the network is not bal- anced, Figure la still shows that the NC1 metric decreases and flattens to rather low value, following the same pattern as the training loss and the gradient norm.
Similar results reported for the training of ResNet- 18 and VGG-11 11G-11 in Appendix E. The implementation of the experi- ments is publicly available at the GitHub repository https: //github. wwweco/o///////e. Dwt.//Dyyaaaaaaeaaeea ht///cmlll 25 ___expr.
5. Within-class Variability Collapse and Generalization
While neural collapse is widely known as a phenomenon occurring at training time, does not necessarily imply that the test error is small (Hui al., 2022, Section 4). We now show that, for well-separated datasets, training via gradient flow implies both approximate NC1 and small test error.
Problem setting. We make the following additional as- sumptions (consistent with Assumption 1).
Setting (b): CIFAR-100. We perform classification on super- classes using pretrained ResNet50 features. Specifically, we consider the super-classes aaquatiic mammals " "large harnivores" super-class containing 5 original classes and 500 samples in total. We then take a ResNet50 pretrained on ImageNet- IK extract the penultimate-laye features of the training set, and use such features as training data. The input dimension
We set 1, and assume the activation function to be the sigmoid function, i.e., (=)) Ite-r- We further assume there are q classes and n data points (rj, Uj) sampled i.i.d. from D, with m points for each class and 27 ~ +D(x)lu))- Each class is balanced, in the sense that DD. ex) = 1/g for all k.
Assumption 3.
"people"], with each

---

Neural Collapse beyond the Unconstrainted Features Model
Algorithm 1 Two-stage gradient flow
Proposition 5.1. Under Assumption 1, for any fixed non- zero W, E W) is strongly convex in p. and there exist unique global minimizer p with the following Gibbs form:
Initialization: Let Wn = [e1.-...,G!! ".."
Stage 1: Fix Wo and obtain P1 = arg min, Cn(p,, Wo).
Stage 2: Initialize at (Wo.P1) and run the following Wasserstein gradient flow:
p ( xexp (- 6 (-aTata) W(WT (Y
3 u) 2 101
dW, Vi w Lm, (Pt, W.)d;;
(27)
5 d0:- = Vo 8p (r (1) + V2B-IdB,
Proof. The result follows from the strong convexity of (p, W). In fact, 1 H YI is convex in H.P H is linear in p. the L2-regularization is convex and the entropic regularization is strongly convex. Then, the claim is consequence of (Hu et al., 2021, Proposition 2.5). O
For technical reasons, we consider the approximated model obtained by truncating the second layer, i.c., (x)
Test error analysis. We start by introducing (T. linearly separable data, which intuitively corresponds to each class being linearly separable w.r.t. the others.
EplTn(a)or(w' ()),, where TR R + R is a smooth function applied component-wise such that
for |=| 2 R,
Definition 5.2. We say that the data distribution D of a g- class classification problem is bounded and (T. - separable if, for each k: there exist dg s.t. | 3 < M² and
TR())= = R + Co:
for |=| 2 2R,
smooth interpolation,
for R < |2| < 2R.
2 if I E supp(D(-lex)),
Notably, for large enough R, the derivative of 7 satisfies
0g
if r E supp(D((lex:)), for all k' + k.
1, for |2| S R,
Given a predictor f Rdd - 124, we aim to bound the mismatch error:
[2(2)= =
0, for |=| 2 2R,
S Co. for R |z| < 2R.
Pr (One-Hot([(z)) # ex]: erricst (J:D) = D((ee)
erTisa(J;D)
We also denote H = [h (x1).......
A (rx) E RRx"" and define the loss and free energy w.r.t. the approximated sec- ond layer as
k=1
where we define the function One-Hot : R - R as
1. if = arg max,Ifl.
[One-Hot(f)ls 0, else.
(P, W) 2n
[ H YIE,
While the (r, -ilinnar separability of the data seem to be restrictive, we remark that in general it is not true that NCI and vanishing test error co-occur under MSE loss without any assumptions on the data distribution. In fact, the occurrence of NC1 implies that the model overfits the data, and overfitting is not always benign without additional assumptions (Bartlett et al., 2021). In this sense, showing the ro-occurrence of NC1 and vanishing test error may be regarded as a harder problem than benign overfitting.
(o, 3) = L (p, W) + 3-1
2 W|
8
11 : + 3 1og pl. 2
We remark that the technical reason for having an approxi- mated second layer is to ensure the uniqueness of the global optimum for the Gibbs minimizer as discussed in Propo- sition 5. Nevertheless, formly for large enough R, which means that we also expect the same conclusion for the original model which corre- sponds to R = +00.
results in Section 5 hold uni-
We now show that training on (T, M) M separable dataset leads to both neural collapse and test error vanishing in the number of training samples n.
Theorem 5.3. Under Assumptions I and 3, let the data separable as
distribution be bounded and (TMM) per Definition 5.2. Pick R >> 1 large enough, large enough, and B (640CGnC3 t.) Then, for any (pts W.) obtained by Stage 2 of Algorithm 1, we have
Two stage training algorithm. Our result holds for the two-stage training described in Algorithm 1. Specifically, in Stage 1, we aim to find the global optimum of & (Wo, p) having fixed Wo, and in Proposition 5.1 below show that, for all fixed non-zero Wo. E, p) has unique global minimizer in (1BPt)) Furthermore, such global minimizer is achieved by noisy gradient flow as studied by (Suzuki et al., 2024a). In Stage 2, we run a gradient flow on the free energy, as we did in Section 4.2.
I
eTr W;); S Cio Coo
2n
log (og(2/6)
+ 60g
8

---

Neural Collapse beyond the Unconstrainted Features Model
with probability at least 1 - 6. Farthermore, there exists ((3) s. for all t > T(3) except a finite Lebesgue measure set,
Acknowledgements
This research was funded whole or in part by the Austrian Science Fund (FWF) 10.55776/C0E12. For the purpose of open access, the authors have applied a CC BY public copyright license to any Author Accepted Manuscript sion arising from this submission. The authors would like to thank Peter Súkeník for general helpful discussions and for pointing out that all the stationary points are approximately proportional in the case without entropic regularization.
p.)(H S0-
The constants Cg, Cio. Cn depend on d, Co, C1. M, but not on n. Their expression is provided in Theorem D.6, whose statement and proof are Appendix D.I. The ar- gument uses Rademacher complexity bounds for neural networks in the mean-field regime as in (Chen et al., 2020; Suzuki et al., 2024b; Takakura & Suzuki, 2024), and the key component is to control the dependence of the constant C10 on n. This is achieved by noting that, for a (T, MM )separated data distribution, a two-layer network with constant number of neurons approximately interpolates the data.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
In a nutshell, Theorem 5.3 provides a sufficient condition on the data distribution to achieve both NC1 and vanishing test error. Although for simplicity in the statement we pick a specific value for B, we note that a similar result would hold for (640C 50
References
< 3 < (p)oy(nn)).
Andriopoulos, G., Dong, Z.,, Guo, L.,, Zhao, Z., and Ross, K. W. The prevalence of neural collapse in neural multi- variate regression. In The Thirty-eighth Annual Confer- ence on Neural Information Processing Systems, 2024.
6. Conclusions and Future Directions
In this work, we consider a three-layer neural network in the mean-field regime and give rather general sufficient condi- tions for within-class variability collapse (namely, NCI) to
Araújo, D., Oliveira, R. I., and Yukimura, D. A mean-fied -ield limit for certain deep neural networks. arXiv preprint urXiv:I906 001600 2019.
occur. We then show that (i) training the three-layer neural network with gradient flow satisfies these conditions, and (ii) a vanishing test error is compatible with neural collapse at training time. Taken together, our results connect represen- tation geometry to loss landscape, gradient flow dynamies and generalization, offering new insights into gradient-basee optimization in deep learning.
Bartlett, P. L., Montanari, A., and Rakhlin, A. Deep learning: a statistical viewpoint. Acta Numerica, 30:87-201, 2021. doi: 10. 10.1017/099229922100000
Beaglehole, D., Súkeník, P., Mondelli, M., and Belkin, M. Average gradient outer product as a mechanism for deep neural collapse. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Three interesting future directions include: (i) establishing more general conditions (either necessary or sufficient) that guarantee both neural collapse during training and vanishing test error; (ii) tackling the challenging case in which there non-linearyty between the last two layers - setting where the properties of neural collapse have been proved in the UFM framework for binary classification (Súkeník et al., 2023); and (iii) extending the results to cross-entropy loss, which is more commonly used for classification. The main technical difficulty for the latter is to rule out the possibility that different data points in the same class could have different logits. which appears challenging even when the loss is small.
Chen, Z., Cao, Y., Gu, Q., and Zhang, T. A generalized neu- ral tangent kernel analysis for two-layer neural networks. Advances in Neural Information Processing Systems, 33: 13363-13373, 2020.
Fang, C., He, H., Long, Q., and Su, W. Exploring deep neural networks via layer-peeled model: Minority col- lapse in imbalanced training. Proceedings of the National Academy of Sciences, (18433)e2103391
2021a.
Fang, C., Lee, J., Yang, P., and Zhang, T. Modeling from features: a mean-field framework for over wver-parameterized deep neural networks. In Conference on learning theory, pp. 1887-1936. PMLR, 2021b.
Han, X., Papyan, V., and Donoho, D. L. Neural collapse un- der MSE loss: Proximity to and dynamics on the central path. In International Conference on Learning Represen- tations, 2022.

---

Neural Collapse beyond the Unconstrainted Features Model
Hong, W. and Ling, S. Beyond unconstrained features: Neural collapse for shallow neural networks with general data. arXiv preprint 3rXiv:2409.08832 2024a.
bounds and kernel limit. In Conference on learning the- ory, pp.
PMLR, 2019.
Mixon, D. G., Parshall, H., and Pi, J. Neural collapse with unconstrained features. Sampling Theory, Signal Processing, and Data Analysis, 20(2):11, 2022.
Hong, W. and Ling, S. Neural collapse for unconstrained feature model under cross-entropy loss with imbalanced data. Journal of Machine Learning Research, 25(192): 1 1-48, 2024b.
Nguyen, P.-M. and Pham, H. T. A rigorous framework for the mean field limit of multilayer neural networks. Math- ematical Statistics and Learning, 6(3):201-35/) 2023.
Hu, K., Ren, Z., Šiska, D., and Szpruch, t.. Mean-field langevin dynamics and energy landscape of neural net- works. Annules de I'Institut Henri Poincare (B) Proba- bilites statistiques, volume 57, pp. 2043-2065. Institut Henri Poincaré, 2021.
Nguyen, Q. N. and Mondelli, M. Global convergence of deep networks with wide layer followed by pyramidal topology. Advances in Neural Information Processing Systems, 33:1 1961--1999,, 2020.
Hui, L.,, Belkin, M., and Nakkiran, P. Limitations of neural collapse for understanding generalization in deep learn- ing. arXiv preprint arXiv:2202.0888 2022.
Pan, L. and Cao, X. Towards understanding neural col- lapse: The effects of batch normalization and weight decay. arXiv preprint arXiv:2309.04644 2023.
Jacot, A., Súkeník, P., Wang, Z., and Mondelli, M. Wide neural networks trained with weight decay provably ex- hibit neural collapse. arXiv preprint arXiv:2410.04887, 2024.
Papyan, V., Han, X., and Donoho, D. L. Prevalence of neural collapse during the terminal phase of deep learn- ing training. Proceedings of the National Academy of Sciences, 117(40)
2020.
Javanmard, A., Mondelli, M., and Montanari, A. Analysis of a two-layer neural network via displacement convexity. The Annals of Statistics, 48(6))3619-338)3 -3442224,,,, 2020.
Pham, H. T. and Nguyen, P.-M. Global convergence of three-layer neural networks in the mean field regime. In International Conference on Learning Representations, 2021.
Kim, and Suzuki, T. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the atten- tion landscape. In Forty-first International Conference Machine Learning, 2024.
Seleznova, M., Weitzner, D., Giryes, R., Kutyniok, G., and Chou, H.-H. Neural (tangent kemel) collapse. Advances in Neural Information Processing Systems, 36, 2024.
Kothapalli, V.
Neural collapse: A review on modelling principles and generalization. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
Shevchenko, A. and Mondelli, M. Landscape connec- tivity and dropout stability of sgd solutions for over- parameterized neural networks. In International Con- ference on Machine Learning, pp. 8773-8784. PMLR, 2020.
Kothapalli, V. and Tirer, T. Kernel vs. kernel: Exploring how the data structure affects neural collapse. arXiv preprint arXiv:2406.02110 105, 2024.
Shevchenko, A., Kungurtsev, V., and Mondelli, M. Mean- field analysis of piecewise linear solutions for wide relu networks. Journal of Machine Learning Research, 23 (130):1-55, 2022.
Lu, Y., Ma, C., Lu, Y., Lu, J., and Ying, L. A mean field analysis of deep resnet and beyond: Towards provably optimization via overparameterization from depth. In International Conference on Machine Learning, pp. 6426- 6436, 2020.
Sákeník, P., Mondelli, M., and Lampert, C. H. Deep neural collapse is provably optimal for the deep unconstrained features model. Advances in Neural Information Process- ing Systems, 36, 2023.
Ma, C., Wu, L.,, et al. The barron space and the flow- fow-induced function spaces for neural network models. Constructive Approximation, 55(1):369-406 2022.
Súkeník, P., Lampert, C. H., and Mondelli, M. Neural collapse vs. low-rank bias: Is deep neural collapse really optimal? In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Mei, S., Montanari, A., and Nguyen, P.-M. A mean field view of the landscape of two-layer neural networks. Pro- ceedings of the National Academy of Sciences, 115(33):
E7665-E7671, 2018
Suzuki, T., Wu, D., and Nitanda, A. Mean-field langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction. Advances Neural Information Processing Systems, 36, 2024a.
Mei, S., Misiakiewicz, T., and Montanari, A. Mean-field theory of two-layers neural networks: dimension-free
10

---

Neural Collapse beyond the Unconstrainted Features Model
Suzuki, T., Wu, D., Oko, K., and Nitanda, A. Feature learn- ing via mean-field langevin dynamics: classifying sparse parities and beyond. Advances in Neural Information Processing Systems, 36, 2024b.
Takakura, S. and Suzuki, T. Mean-field analysis on layer neural networks from kernel perspective. In Forty-first International Conference on Machine Learn- ing. 2024.
Thrampoulidis, C., Kini, G. R., Vakilian, V., and Behnia, T. Imbalance trouble: Revisiting neural-collapse geometry. In Advances in Neural Information Processing Systems, 2022.
Tirer, T. and Bruna, Extended unconstrained features model for exploring deep neural collapse. In International Conference Machine Learning, pp. 11478-21500 PMLR, 2022.
Wainwright, M. High--iimessioal Statistics: A Non- Asymptotic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematies. Cambridge University Press, 2019.
Zhou, J., Li, X., Ding, T., You, C.,, Qu, Q., and Zhu, Z. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In International Conference Machine Learning, pp.
PMLR, 2022.

---

Neural Collapse beyond the Unconstrainted Features Model
A. Technical Lemmas
Lemma A.1 (Properties of Kronecker product and vectorization), The following properties hold:
1. vec( = (CT & A)vec(B) B).
2. ()C @ D) (AC) & (BD) if one can form matrix product AC and BD.
3. Singular spuce of Kronecker product: given two matrices E RRxx and B E IRPx9, let their SVD be
A= UASAVA
B= B-UbSEVn.
where UA E &mmxm SA E RRXX VA E R and Un E RRxP, SB E RRxx, VB E Rox9, Then, the SVD of A 1 B reads
A (UA @ USA (V)
Proof.
The first two claims be easily verified. For the third, we have:
(UA (U)(A @ (V)(Va @ VB)= (U'a & UB)(SA & (R) 8
-(AAUU)(((AAT)) @ (S
=(UUSAVA) & (UnSBVE))
=AQB B,
and
(U 8 (UA@U)) (UTUN) (U]UU) Im Lp Imps
which gives the desired result.
Lemma A.2. Given two matrix A, B E Rnxn, assume that A is invertible and A + B is invertible then we have:
()+ A (A+ B) BA-1
(A Proof. Let (A + B) -1 A-1 +1 C where we aim to compute C. Then, we have (A +)+ B)C 1. This
implies that BA- + B)C = 0, and we have C -(A + B) BA- which gives the desired result.
funcciioon Lemma A.3. Let p(0) E P2P be absolutely continuous measure, C(p) P2 - R be non-negative functional, and
E(p) = C(p) + 1 + 3 -.'EE,,logr))
Then, for any p E P2(RP), we have that
2T
((p) & E(p) + 3 log
2
X3
D
2T
DRI(ellPO) SB (E()) B
log
2
A
E,110 -41-1((6) 1-1
1 +
log
\
D
2x
C(p) 2 10 < 38(p) + 3
log
28-1
1 D log
2
X
X3
where Po x exp(-3X1l013/2)
12

---

Neural Collapse beyond the Unconstrainted Features Model
Proof. Note that
B 3A
BB log p]
P -Bp log (xx(
D
log
2T
(2
31
D
2m
2 B B-
log
2
A
where the last passage follows from the non-negativity of the KL divergence. This implies that
E(p) = C(p) + 8 DKL(plleo) 3-122 D 2x
log
C(p)
-3-1
log 2
2m
X
X
which gives the first two inequalities. The third inequality on the second moment follows from (Mei et al., 2018, Equation 10.12), and the final equality comes from combining the first and third inequality.
B. Proofs for Section 4.1
B.1. Proof of Theorem 4.2
First, given p. W, we define
(0: V. 6 V =
"e Hp - Ww (X u) + Xaa 8- p(),
(28)
and we will use the shorthand A in the rest the proof. By the definition of cs-stationary point, we have
Bp 12 S Es
By rearranging terms in (28), we have
a TAp'
log W(NW -)))(())))(( - 1-18-1V
V. +170)+1-100 2T P a.s.
n
which implies that
Hp = (u X)]
= x- Hp - Y) Kp(X.X)
---33-1E, log T X)] + X,'E,AA.0 x)). (x)).
n
We first show that the term V.a log p(0)o(u T z)] = 0, for any To see this, it is sufficient to show that
0n, log ((0) (T) p(do) 0. /
Indeed, we have
f
0a log ((0)(((((u)) p(d0) = (o,p(0)o((T))))) x))0= - On (x) = 0,
which implics that
Hp =x==>;''(GW)) H, V) Kx(X,X) + (
(29)
n
By multiplying both sides of (29) with T and subtracting Y, we
T H. Y = Ww H, -Y) K K,(XXX) >m' "..'WWBBloA BolAso(u x)]
13

---

Neural Collapse beyond the Unconstrainted Features Model
An application of the first property stated in Lemma A. gives that
vec(>W] Hp Y) --= -
(12 (n2 (x
2 (W (WTW) W) Aptin (vec(Y) -x >xp'vec(>W (eec>)""plAaotte EplAao (x))) (30)
n
Plugging the expression for vec(nW HH Hp
Y) back to (29), we have
vec(H,) -xp' (7 X)(X,
(W) bec(TWTTH, V) + (/'eee(B,pAAA x)])
Xx(X,X) n
K,(X,X)
@ (WTW)) W) + Aptna) -1
8 (WW
(vec(Y) Xp 5='ee BOoWTTE,le ou x)))
Xp vepp (p (x])
X)])
x 2 W)
(>2 K2(X,X)
8 (WTW) + Aptra vec(Y) + Ej (Es, Api W),
n
where we define the error vector as
Biles:Ap:11 W) = -xp² 7 K
(12 (,(X,X) & (WTW) W) + Aptng)
@ W/
vec(yWTB, Bpl x))) xx))
+ >=='vee(E,AA, X))).
Now we aim to upper bound the error. To do so, we write
(es, Api7, W)ll
< 22x,²
(1 Xx) @ (w) (
X)
(WTW) + W)+ Aptra)
vec(WTTE,lee (X)]
n
n
+ (1,,
T x)1)11
Xx
2
max (1x
6 (WT W) + 2ptna) )-) 2 1o: (u x)]1
S 222,20 20 max
00 w)
n
n
+ 2x, 3A [x)]1
< (2x, (2xp">² (W)* + 2x,²) ||E|Aan Tx)1I
(31)
Then, upper bound |EE|A.o [x]] as follows:
|E T x)11? < o (x)
= E. 10 x
(32)
I CinE
S onndz.
Combining (31) and (32), (6)-(7 readily follow.
To obtain (9), we first show that, when W W is full rank, the following equality holds
K XX,X) (ww
K (X,
2 W WTW W + ApIra) )
n
(x,(XX)
= = 2 (WWW - -,7 > (Ia @ (W(WTW)-1)
8 W W + xptng
n
(33)
14

---

Neural Collapse beyond the Unconstrainted Features Model
To do so, define the eigen-decomposition of (,(X,x) and SVD of W as follows:
X) -KKEKUK yW = -U-WS Vw:
where UK E R Σx E RRxE, Uw E Roxp, Sw E ERXO, Vw Rox0 By Lemma A. we have that
K X) &W = (Us0Uw12
70 x W = W
Thus, we have the following equalities:
XX, n
② @w) (52 (x2 & WTWW + Aptin)
n
(UK & (Ex & Sw)(Uk @ Vw) TUUx @ (w @ (WTSSw + Aplma) @ VWN)
-UK @ (EK @ (Ex @ Sw ) + )pIng) ( (Ux @ Vw))
For simplicity, we write the matrix Sw ( 9n) and define Sw (
Here, given
0p-0.9
0p-4.4
integers n, we define O,,m as the x matrix containing zeros. Clearly, we have that
= [1 I/
0q.p-9 Op-9.7-9
SWSW
1p-qn
Next, we observe that
(In @ (w'SST) WTx & Sw) = Ex @ STT = Σx @ Sw.
Thus, write
(Ex @ Swx @ SW) \
= (In @
(Ex @ Sw) 8 SW) + Aplng) (In Sw (L, (Ex @ (w @ SW) + XpIng)
= (Tn s Sw (Ex 2 Sw)) @ (ST ) + ApTa))
(In @ sw' - Ap @ Sw) + \phng)
which implies that
-1 n X) so w (>> (x(, WTW + Ap
n
a
= (UK 2 (In @ Sw)) @ VwwT - >,,UUx @ (l) @ SW @ SWW + ) (UK @ Vw)
Finally, we can verify that:
7 I & (W(W ((W-) (UxUR) 8 (UWS
(Uk Uw) @ (
and
(Ia @ (W(WTW) (, (22 K,(X,)) @ W W + XpTin)
=(Ux & Uw) Sw) Viv)T(xx Vw (Ex
Xpla)) @ Tw)
=(U & UT) @ (Fx SW) + ) [UK @0 Vw)T.
which gives (33).
15

---

Neural Collapse beyond the Unconstrainted Features Model
From the above decomposition, we know that
X)
8 @w) (10 Kx(X, 8 W W + Xp vee(N)
n
-In & (W(W ([W"))-l))odff Ap7 (In @ (W(TTW) g2 Xx(X,X) X)
-1
WWT W + +xpIng)
vec(Y)
K)(X,X)
vec (WW w) 7 (In 1
(W (W (W(WWW)-1) w)
W W ApIna
vec((),
n
where we use the first item of Lemma A.I in the last passage. Let us now define
Apies.Api7hpt W) == -Xp7 & (W(W (W(WTW)-3) (x2 (2 2 W + Aptma
vec(Y)
"T, @ (W(WTW)) [vec(TW [ H - -Y) (12 Xx(x,
8 W W + Aptna
vec vec E (x'x))))
n
where the second passage follows from (30). Using (6) (that we proved above), we have
vea Hp) = vecly- Wy- + E2(cs, Api W) + Apss,,
remains upper bound |E((s Ap: 7, P. W To this aim, we write
EA W) - 2y 2 0 max Ln @ (W (W w) ")) TI, H1 -Y " (
2n7 2 Cn W) 2 | E (x)1I
min (2)
2n W) + 2x= o max (w) x)Il
mmin (W)2
Plugging in the bound in (32) gives
2 (p, W) + 2x,²0 max (W) (W)2Cfne3
(es, Api b,
(34)
Mmin (W)2
By combining (7) and (34) with an application of the triangle inequality, the proof is complete.
B.2. Proof Lemma 4.3
Proof of Lemma 4.3. To upper bound Fmax (W), we directly use the definition in (4):
m 2 < W < 2A (p W) < (AW (1o
To lower bound mmia (W), we start by showing that
H 1H < g
(35)
To see this, assume by contradiction that |HI > 8 Then, there exists such that [1,p(x;)] > bounds hold:
20 and the following
(x:)(=)) = [[E xx)]]13
= la a'a(( T( 2;)1 (0' is an independent copy of 0)
I Ba [aaaT|allla( xx [l)((((
I Cf (r [al]
s Cf 11 [|]]|]]]
cf |||a||]]
s o E. |||0|2),
16

---

Neural Collapse beyond the Unconstrainted Features Model
which implies that E Epllaa
gro c
By combining (15) with (4), have
Epllaa[Í] S 1 - 11B(2)-1(100 33
(36)
where we recall that 2 = 203-1. 19.3-.Nooe Note that, for all co E (0, 1) and for all B 2 e e 50
log 2
geo (log B)a
2 Beo/2
(37)
Thus, by taking
B max (e log (2C3nB(2))
as in (14), we have that
8"0 Of
is strictly larger than the RHS of (36), which is contradiction.
Now, we are ready to argue that, for large enough 8, miin (W) 2 60 Assume by contradiction that m miin(W) < 80 and w.l.o.g assume [Ew]a < 60 Then, the following lower bound holds
La W) 2 Ln(p,W)
| - 2
2n
VV ETy U Hp 1
2n
Hp - VI
2n
(38)
2
[[Ewla. Hpla: - [V 1x:
2n
2
I[V 1gg
||[E 1g. Hw
4n
2n
1
1
[EW
(W H pl
4q
2n
1
r
4g
2n
(44) 1/2
Note that by taking B 2
we have
1
1 3
2 1
Furthermore, for all B 2 edo log((n), we have
4g
2n
B1/a
log B > B1/(2a).
Thus, by taking
la log((2)))
Aq n
64()B)2
max
as in (14), we have that the RHS of (38) is strictly larger than 138-1(log) 3) which is a contradiction and concludes the proof.
B.3. Proof of Corollary 4.4
Proof of Corollary 4.4. From Theorem 4.2, we have
H, = -,=
Wy -W + Ez.
17

---

Neural Collapse beyond the Unconstrainted Features Model
Note that the NC1 metric is scale invariant, thus it is equivalent to compute NCI(HH,). Later on, will write Hp :- 7111p with slight abuse of notation. Some manipulations give
H, - W(W W(WT ( (r- - q 19tn + trB2
In In117
Mc = 1 WWW
Y Y 1 rT + 7 E
In - 1 1n Y.Y
m
9
Hp Mc -E2
In - 1 Inth
(Im
In 1 yTv)
Thus, we have
TH ((H--M)* M.) -Ma)
=HH, - Mcll
= -1 In 1
IT
(I-
1
m
- E
Tr Hp Hp}
H
1 W(WTW)- w)-1
1111
2
11
In
--11 113
1
9
n
1
(9- 1)m
4y² 1
20
(W)
q
which concludes the proof.
B.4. Proof of Lemma 4.5
Proof of Lemma 4.5. From the stationary condition, we obtain
V. 5 En W) = W HG 1-)o('oX' u) +>pa+ --pa+ V.lgvvo =0, P a.s.
n (p) - H, -Y + +w =0
80
Rearranging the terms gives
\ W(W H2- YYo(X --'2,"V p a.s.
W >wh1 7 H,,,WWT,,-Y'' -Y.
Then, compute:
X
E 1 - Ep
(V. log (( W(W HH-HH -YHH, 3-1 EE, ((V log p(0))a²
W(T - D)((XTuja -
1
3-'2''E,
n
Xw
WWT
W(W 10-Y)H,
18

---

Neural Collapse beyond the Unconstrainted Features Model
Thus, Aw WWT Bp ] B-1 [V
log (())a J. and next we compute Ep [(V (0))a))aa
]]
To do this,
note that:
Ep [V. loglog// ((0))(())aa))) 11.s Ep (0., log p(0)a)]
/ (0))aa log (())a d0
a(0,,
/ p(0)0,,a, do
-1, i=j
0, itj
Thus, Ep [(V log p(0))a Ipp. which gives the desired result.
C. Proofs in Section 4.2
C.1. Proof of Lemma 4.6
Proof of Lemma 4.6. We note that (p. W) is lower bounded and we have:
(pT,, WT), Eno, Wo) - S'' OEEn (pEn(pe, W.)
A standard computation gives
8 Ve Enno En wa)(e)) 5p
(p,, Wt) = -Ep.
CV Wo)l
which implies that En (pr, Wr) is lower-bounded monotone decreasing sequence. Thus, limm & (pT, WT) =C C < xo.
The existence and boundedness of lim En (pr, Wr) implies that, for any es 0, there exists ((ss) > 0 s.t. for all t > T(cs) except finite Lebesgue measure set,
27 + V ( < ES.
Ep.
V
(p., (i)( 8p
which finishes the proof.
C.2. Proof of Theorem 4.8.
Theorem C.1 (Full statement of Theorem 4.8). Let Assumptions 1, 2 hold, set Xo Aw g-1 and
> Cs, to = 3Cs.
Then, for any B and any 2 to we have
Cx, We) B BCa
19

---

Neural Collapse beyond the Unconstrainted Features Model
where
4B1
2/2BB Bpp
C3 max
Rw
B3 B1
g +1
p d 2
C4 = 6y 2 ? max
By 2
+3
+ 3R +
log 2m + 22 + (p d) log 8T)),
B? B
2
C5 = y min {
B1 B3
B2 B1
1
X?
Rw = 2'
Rp = min
p p+ d,
C4n²
X.
/2n Rw + R) (2222dddp++++ +d+1)
A1
A2 32
2n'
X
B1 / /Rw Rw + L/pp d(p d) 20 22 R,AT"
n
B2 2(W + +1) + VRW Rw + 111p
d) 2
BN = 2p Rw + +14411
d(p + d) +) 1 Vp=d+2c |,01 B4 = np Rw + (14 Vp
To prove the above Theorem C.I, we first define the following first hitting time for any fixed Rw 1, Rs < <p+d d:
t. min : [|W, W, W Wollop > Bww,inn(t DKz(Pilleo) > Rp}}.
From the above definition of f,, we have that, for & S t..
IW, W. W.J Wollon S Rw, (RI(Pellpo) s Rp.
The next two lemmas (proved in Appendices C.2.1 and C.2.2) control the behavior of the dynamics before 4,,
Lemma C.2. Let Rp S min + d, c Rw S 1. For S tx: we have
X²
(en W.) < exp(->2AA0)) (>²1t)) + -
where
X
A2 A2= 2pn [Rw
A
4n'
Lemma C.3. Let R2 I min + d, c4n
X²
), Rw I 2/- For S to, we have
IW, w 16c-WW Wo Wollop 5> B1 + 3- Bet.
DDKL(P,,p) -- B3 3- But
where
B = - Rw + 1 1 +) RR,
B2 2(Rw +1) -1) + n
/ Rw 1/pp(44C V (p d) + 201) / /R,A2.
B3 22pV Rw + 1(801 p
B₁ = V 2pv Rw + p +
Now we are ready to prove the main theorem.
20

---

Neural Collapse beyond the Unconstrainted Features Model
Proof of Theorem C.I. We pick to 7 min an application of Lemma C.3 gives
Ba ), and consider two cases. If Lx lo: then for any s 4, < lo:
B3 B4
I[, W, - Wo Wollop - + B B2l S + B2lo S 22 B1,
B3 8 3- Bato) v 47 BB3. (p,( & (7 1 B3 + 8- Bal)² - 1B
By picking
4B1 2/2B3 Rw Rp
7 2 max
we get that:
Rw
R
IW, W W Wollop
DRz(Pi.lPo)) -
2
2
for all S to with to > t., which contradicts the definition of t.. This implies that t, 2 to and, by Lemma C.2, we have
B1
B3
Ln(pto: Wio) S 2 exx
Ay min
3
+ 2 28-
B2
B4
B1
B3
-2
2 yA1 min
3
+ +2y 23-
B2
B1
B2
B
- 2y 23-22 max
+ 2n 28-
B
B3
We also have the following upper bound [W 1/3:
W II: < gl(W T Wiollop
S < to Wi W. Wollop + +IW, Wollop)
S a( Rw + 1).
Thus, we can upper bound the free energy for all 2 to as
8-1
En(pt, En (pto: W (P Wt,) +
W 11? + +
(D'iii((tollpp)
2
B2
Bi
8-1
2y A, max
2y 2- ?A2
1 [(Bw + +1 3-1 RRp
B
B
2
Applying Lemma A.3 gives that for lo:
Cxmule Wi) S (pt, We) + --1 log 2x 2 (1 + dl S g-
with
B2 B
B B
g Rw +1) (Rw
3Rp + p + d
-2 C4 -- max
6y
log 2m + 2 (p + d) log 87),
2
2
where we use 3 1. This completes the proof.
C.2.1. PROOF OF LEMMA C.2
We first compute the evolution of Ln (pt, Wt)
2n
W., Hp. - Y under gradient flow:
1
d
dt Cn,, W.)
=
-r,
W. H
12
dt
( d
+ (1 rr, W. TTT (a
d
We
Hp:
H
dt
dt
F
21

---

Neural Collapse beyond the Unconstrainted Features Model
where we define ry = W Hp: Y R The evolution of W is computed as:
1 Hp - W W.
dt
and the evolution of Hp, can be computed as:
d
/ ao X) ()
- Hp pz
di
ao X)Vo (p.())VoV((I)) de
Va Valo(o)( a (VV. XDiag(D (u X)) pc(d0).
where we define the potential Vi(.): Ro+d + to be the first variation of the free energy
8 vi(0)= (pt, (.)(0) r aa x)) + + 8 log P((),
6p
F
and Diag(o'(u X)) E R=x to be the diagonal matrix with o'ww' x;)) on the i-th diagonal entry. The gradient of the potential given by
VVV((0) 2 (X w) + 3 B px(),
(39)
V V(0) 2 2'Diag(o'(" x))r WT a 8 --' + g Vv.. u log (0).
Thus, we can express the evolution of H,, as follows:
d
- 7 / WXno(XTu)ot uuo(u x) pd(0) ni
Hp.
dt
- --1 / (a+ Va Po (0))o x) Pc(d))
7
an T War::::iag(o (m X))Xx XDiaglo' (u x))(((d)
n
/ [auu+ V, log p.) XDiag(o (u X)) pc(d0).
g-1
Now, we can write the evolution of the empirical loss function as
d
72
HpD 7/-1 (re, (re,W W, HaDr
Cnt,, WT) = Tre
n² H..
(re,
dt
72
o u)o(u x) (i((o)
n²
78-1
Wn. Va log x) x))r(())
22
/ W, W,JaaTwae WarrDDiiagl( (x))x XDiag(o'(vl x))e((())) (T)))(((())))((((
n?
Tti
7/8-1
(r.,W. / aauu + Vn log (0)) XDiag(o (u XX) pe(d0)
We first control the potential positive terms via following lemma.
22

---

Neural Collapse beyond the Unconstrainted Features Model
Lemma C.4. Let R2 sd+ P. Then, for S lies we have
(Te, W/ Had a + (r..1// / [a Valog x) p)()))))
+ WT / aauu Vn (0)) XDDagloo (u X)) pc(
$ /8/2mm2 + 1) 1) Ln / Wy).
Proof. We first note that
log 2z) pc(d0) U V(VV.llog "Daggl (u T X)) pe()))
/
/v. V.
/ Volo(u (e pc((0) 0, = / pe T z (u' xz)
ij
(40)
/ pc(O)V. (ajo' (u' 2z)zz)) do
f a,o"
x2))( (xx))\) (Br(
For simplicity, we define the function
(0) = 20, 2z) + ar "ajTTzo'uu 22) a,"" -=x""u[T)))I 21111 [G(O)) (0) E Roxn.
Then, have
Wwo.. W. Hpddo re, W. f(a (o)gp (x)())))((xx)(((( p)
+ re. (no.T WT / V(uu+V. XDDig(o' X)) p((d))
K / (c0)) F
Sllllll/|Walloo || / G(0) (((0))
Σ (ʃ (e.(0) Peco0) 2.f
Σ ( (e: (o:-0(o0))))))))) 2 2.j
SV2nCul0r, (.)
where in the last step we use that
E [2a:o(u (2,)) + au Tu, 22) - (
= =Epo luilee [uo 2z) + u (u o" [ = O,
since Po ~V(0, Ipta).
Following the computations in (Chen et al., 2020, Lemma A.1, Equation C.4 and C.5) and using Assumption 1, we have
|| (0)l1 S (C + 1).
23

---

Neural Collapse beyond the Unconstrainted Features Model
Thus, by (Chen al., 2020, Lemma B.2), obtain
( (pr po)(do)) s (SCiVD+d+ (p))
Hence, we conclude that
1) 1)VLn)pp, Ln(p;,W), Kn, Kn.WT W.T / ((0))((00)) s 8 (Rw + 1)R, (2
where we use that, for < l,s
W2(pt: Po) S (PD((p))) & Rp,, (by Talagrand's inequality, see (Chen et al., 2020, Lemma 5.4) W = IIW, Wellop & I Wollop + [|W, W2 Wo Wollop S Rw + 1.
This concludes the argument.
Next, we lower bound the negative terms. We first recall the definition the kernel:
Kx(X,X) = f f / u)o(uu (ouuu X) x(d)).
Furthermore, by Lemma 4.7, Ami ( (oo (X, 2 do > 0. As Ami (Wo Wo) 1, this implies that Amin (K (X, X X) 6
WO Wo) 2 A.. We then have the following lower bound time tx.
xa and Rw S 2/2 Then, for t S t., we have
Lemma C.5. Let R2 I min(p + d, can
Amin (K (X, (W, W,)) 2 4
Proof. First, by Weyl's inequality we have
Amiu (W, Wi) 2 Amin Amin (W. Wo) - II Wo W, Willop 2 1 Rw.
remains to lower bound Amin' (Kp, X)). To do so, note that
(x;, Kn (x1,
[x,T)) [o u:)o(u x;y] Epplo(" ( xx)o(" azy)]
< (CC.
< / po) S 4C7 Rpp
where in the first inequality use Kantorovich-Rubinstrii duality. Thus, have
|| K (X, x) Ko, (X, < |K (X, - K S
which implies that By picking R2 min(p + d, o4nec1 } and Rw < 2: the claim follows:
Amin (K,. (X,X)) 2 Amin (Koo (X, - |I( (X,
(X, - Kp, (X, 2 h 4nC2 CCfVVo Rep
x²
X.
Amin Amin(K,, (X, @ (W, W.)) Amin (Kp, X)) (W, W.)
4
24

---

Neural Collapse beyond the Unconstrainted Features Model
By combining the results of Lemmas C.4 and C.5, we have that, for Rp S <min(p + d, 64ntey Biicg
}. Rw S 2 and t & 4,
d
x2
Amin (Kp, (X, (X, wolllell + 73-1 Wr Hp.)
(pr, Wi)
dt
n²
W, / ( (abgp()))))
re, WT, /afu+ + V. log ((()) XDiaa(d x)) pc((0) p))(()))))))))))))))))))))))))
n² 2nL(ps Wi) 3 8 2pn ( W.).
2 Xa
n
By dividing both sides by 2/ Ln (pe, Wx) and defining
A_ 2
2 (Rw ()Rg,(CCiD++++]
Z2
4n
get
1
d (ec Wi) S - - (pr, W,) + 7 722
2/ (pr Wi) dt
Note that
Cn (pt, W.) - de v (pee, W,). Hence, an application of Gronwall's Lemma gives
W. W.)
Cp Wi) S exp(->²Z1) -n² (VEn(po. Wo) -7-11-111111 Z2 4
2o Z1
Z1
< exp( (-2tt))))) +--1 2z
Zi
which gives the desired result.
C.2.2. PROOF OF LEMMA C.3
We first control the term |[W W. - W.J Wollop- Note that, as W, W. - W.J Wo is symmetric, we have
IWW T W. WT Wollop = max(Amax (W, W. W.J Wo), Amax A Wo w. W.)
Then, for any fixed v E 89-1, we have
d d T (W, W. Wo Wo)c di
W. W.o
d
TT ((.w.) w, WT ( (2w:))
23 2 TTT, Wev WWo-2o7e (oT(rr W, W. HorTe) Jo n
100T (r. W, W. Hprr))o
27 1ITT (TTTelll Weell2
27 ||r ||e||H, Welle
27 (2nL. W. W.
n
To upper bound | Walle, by using the same techniques in Lemma C.2, we have:
|| Wallr & [[ llpll Hp, H Holllr
Pe
25

---

Neural Collapse beyond the Unconstrainted Features Model
which, as |Wellop VRw I and VVa,o( (axno((TT,))aa 22)1l2 s 201 + 1), E [p), [n]. gives
|1 Well S RRw + -t
(41)
S <Rw Rw + +1/m1CC///44++22222/
Similarly, we have
d 1 (w,, Wo - W.TW)o = d
dt - -u ..w,Twaa Wee
dt
= ((#.:)
(
(d.:)) v
+ + W. dt
28--0 'o TW, Weo IoT T (r Ws + W. Hprr )
27 V2nCu(Pr. Wi) Wille
28 '1,, Wellop +
n
27
s 388-(((1 + +1)
22L W W Walle.
a
Thus, for any & < t, and any v E S4-1, we have
47 v Rw +/pn(4i///==)++000 /R, 1 W.)l n
lo (W, W. WoTWo) $ 23 (Rw 1)t
which implies
47 V Rw + + Rp / VCulps, W.) ds.
(Rw + W,, W. - Wo Wollop < 2
1)t+
For simplicity, let us define
4 v Rw + T/pn(ACi (p++ 201) Rp
Za 2(R + 1), Z4
n
By plugging in the result of Lemma C.2, have:
IW. W. - To Wollon Zz + 724 / (exp(->" + -3-1A22
A2) ds
= = + ZIA2)tt - expp exp(=>>2 (>
<8- + -
Next, we control the term (pe(pellpo). First, the time derivative of (Pellp))o)) is computed follows:
_ DRi(eelle))
d
= d
di 2 Ep, 1101
+ E,, [log pt]
= / ( 11o1 + log pr (0)
Ti. (r))
/ (> 11011 + pr)
Vo - (()
f (0 V log pr (oV(())))))))
26

---

Neural Collapse beyond the Unconstrainted Features Model
By recalling (39), we have
d D
/ |00 + Volog
2
vo (or p:(d0)
11
7
(u + V. log (aiag/) XDiag(o'(u (u X))r, W. pc(d))
n1
rr. Wi / (a (a+ (alop(()))) X) p(((0)
W. /alu V., V., X X)) pc(d0)
n
$ Kro (c10)a(00)
where by recalling (40) and defining
(0)(0) a: xz) + aqu" aaau au (r [G(O))& (0) E 12xx
a,o (ITT(x)z)lee 21)111
we have
(Kc(pe)lpo) < 2 WT. G(O)) (d0) dt
d
s 2 |Willoo///nCnll Wi),
Σ f
(f po))a)))
Following the computations (Chen et al., 2020, Lemma A.1, Equation C.4 and C.5) and using Assumption 1, we have
(o)), 4Cr(1 1).
Thus, by (Chen et al., 2020, Lemma B.2), we obtain
I (e (p:
(d) (o)) < (GGG)))+++)))M(O.)))) (000/tt----/))//////oo
Hence, we conclude that
d df
(tte:llaa))ssss///oxoo.lllll Cn, W,),
with
2. 22p Rw 1(8C
Thus, we have
d dt
VDKL(Pllpu) S 772s (..,
which implics
'' VLulps, W.
VDKe(Pillpo) 5 5yZs
= + -
thus concluding the proof.
27

---

Neural Collapse beyond the Unconstrainted Features Model
C.3. Proof of Corollary 4.9 Proof of Corollary 4.9. By Theorem 4.8, we have that, for t to.
(pt, Wi) S C4.
Then, by Lemma 4.3, we have that, for any co < 1/2, by picking
8 2 max (207nCs) ,. { 1
64((C1)3
we have
miin (W) 23,,
max (W)) 1 200,
Plugging this in (10) gives
Ea(ss, P: W.) S 2 + (83*
233 (a
By Lemma 4.6, we know that, if we pick és small enough, there exists T(cs) s.t. for all > T(es) except finite Lebesgue measure set,
[2(ss, 8:7, Pt. w S
Consequently, taking
82 (6400
1
50
ensures that (16) is satisfied and, hence, we can apply (17) which gives that, for all t>> T(es) except finite Lebesgue measure set,
(CI(HA) S
S 60.
(g-1)m 4gCa
- 16nCz - 1+20
Finally, by taking co = 3- we finish the proof.
D. Proofs in Section 5
Throughout this appendix, given E RP, we define the partial order z u if Us S u, for all i E [p]. Additionally, if R1, we write v < Rt as a shorthand. Similarly, we write v Á u (resp. v B R)) if there exists some i such that e; >> u4 (resp. U; > R). The symbols < and A are defined analogously. Furthermore, for a given vector v, we define the Jacobian as (R) = diaggiigggg (v) (v1)....... T/ (H(p)) E R2x..
D.1. Proof of Theorem 5.3
We start with a result (proved in Appendix D.2) controlling the generalization error for data distributions that satisfy Assumptions and 3. recall that given a function class F, the Rademacher complexity is defined as
Ry (F) E,
sup fEF n
Gu Ba (TTTx:)))]
1=1
where 6 are i.i.d. Rademacher random variables.
Lemma D.1. For i (1,....9), , let F, be class of functions from R R. Let D be data distribution satisfies Assumptions I and 3, and let I,,..., D. Then, for any f 12d - R29 s.t. [f], Fr, we have
2/² V Cn( + Σ RR - /
log(2/8)
err err (f: <
1=1
with probability > 1 5. Here, we define L.
In Σ (1; ) - 11112.
28

---

Neural Collapse beyond the Unconstrainted Features Model
We then define the following functional class corresponding the output function of our model:
F(M Mp) = {f IR 7 R f(z) = w hp()), || S Mw, Epl1001 - Mp).
(42)
The next lemma (proved in Appendix D.3) upper bounds the Rademacher complexity of the functional class in (42).
Lemma D.2. Given Mw. Mp > O, we have:
Ra (Mr Mp))
2n
The next theorem (proved in Appendix D.4) shows that, if the stationary point we achieve has small regularized loss, then the test error is small.
Theorem D.3. Let (p, W) satisfy
(x, S ")
Then, for any 0 < Eo < 1/2 and
B Z max
log va
20 (2CrnB)
[g
64(gB)2
n
the following upper bound on the test ervor holds
CtT
log(2/6)
(log, 3)a/2 + Bg 3)8V erTrese W); -B (/9
+ 6bp
2n
with probability 2 1-5.
Theorem D.3 implies that it is necessary to control how B scales with in order to control the generalization error. We now show that one can do so for (r, 1f -linearrll separable data.
Lemma D.4. Let 'D be bounded and (r, --lliearlyy separable data distribution as per Definition 5.2. Pick R > q and let (22) T+c Then, for any O, there exists with
M (log( (log(v//c))2 2
p + d
c-
Ep 11 I 9²
+ +2,
Ep, [log P1
<
log
2
2Te
such that [R(=; Wo) = Wo R (x) approximates well the true data distribution: p1
Ri..)>lllf((s Wo) - vl 5 26 2 + 4(gd CipIOE..
Lemma D.5. Consider bounded and (r, -illirarly separable data distribution as per Definition 5.2, Let (pt, Wt) be obtained by Stage of Algorithm Then, for any 3 0, we can pick large enough such that
(pt W₂) S C.B81111000 log 3,
with
Cs 3 3 =1 1+2(g2d+66p1 +
27 +
M
P P (dd
+ p+d d
log(2
4
2
The proofs of Lemma D.4 and Lemma D.5 are provided in Appendix D.5 and D.6 respectively. Lemma D.5 shows that the constant in the upper bound of the free energy does not blow up with n.
We are now ready state and prove the full version Theorem 5.3.
29

---

Neural Collapse beyond the Unconstrainted Features Model
Theorem D.6 (Full statement of Theorem 5.3). Under Assumptions I and let the data distribution be bounded and
(T, - separable as per Definition 5.2. Pick >> 1 large enough, n large enough, and = (640C3nC3 ) Then, for any (pt. Wt) obtained by Stage of Algorithm I, we have
I
1og(2/5)
erriest efTiest(f(:Pt Wi)D) I
+ 6bg
(43)
2n
with probability at least 1 - 6. Furthermore, there exist T(3) s.t. for all t > T(3) except finite Lebesgue measure set,
NC1 p.) S 60-
(44)
The constants Cg, Cio, Cn1 are given by
M
p+ d
3(p + d) log(2x) p d
Co 9 2 4(d C(p)Ci+ +
+
+ 2(1 + (p d) log 8x),
272
4
2
Cou = 50gCg v C7T,
Cn 64
Proof of Theorem 5.3. By Lemma A.3 and Lemma D.5, have
(pr, Wi) S 38, (Pt. W.) + -----+
log + 2 + (v + log S Co,9-1 log B,
2
with
Cg = 3Cs + p d
+ 2(1 + (p + log
2
Thus by Theorem D.3,
Ftth 2n
+ 6qg 12g(2/0)
erTiest (f(e: (f: Pe. D) < 16(Cg,8~ Bog + SqCy log 3
(45)
6 and that n is large enough, the desired bound (43) readily follows. Finally, by proceeding as in the argument of Corollary 4.9, we also obtain (44) and the proof is complete.
with probability 2 - 8. By using that (640CFnC) mo
D.2. Proof of Lemma D.1
Proof of Lemma D.I. It is easy see that (F: 1 E upper bound:
We have the following
1ne Ho))f())))
= 10re (1)
S [110)) (Note that 2c - E (+1))
1 11(2
where by 1g here denote the indicator function an event E.
We now follow the approach (Chen al., 2020, Lemma 5.6) and define the surrogate loss
1, yn < 0;
Cramp (V.
2yn/ 1, 0 5 nv < 1/2;
0, yy 2 1/2.
Then, Cramp is 2-iischhitz in of and, for y E {+1),
1 1y S Cramply' v) S by - -/1-
30

---

Neural Collapse beyond the Unconstrainted Features Model
Thus, have that, for any æ,
(120))))-)))000011111000 S
Cramp(2[f(2)) - 1, 2e 1) - Σ -f(z))--xI[i)
1)
- 22 -x113
which leads to the following generalization bound
1
(f:ies((fD)
Eoe Home- (a(()))vl
9
ED(-lew)
9
2
leww
Tamp (2[((=)) 2ex 1)/2
9
2
n
(ramp(2( (2[[(x;)] - 1, - 1 ]
4
log(2/8)
<
R. (F.) + 60
with probability > 1
i=1
(2/()/5)
[[ 15l12 + 4 Σ
ni (FF)) + +
with probability > 1 -1-6
1=1
1=1
2/2
1
1og(2/5)
[xxa)) 1,,12
(F,) 6gg
with probability 1
V
2n 1=1
n
2/2
9
log(2/5)
V
(f) 4
Σ
Ri (Ji) /0/
with probability 1 8.
/4
D.3. Proof of Lemma D.2
Proof of Lemma D.2. The proof is modification from (Takakura & Suzuki, 2024, Lemma 4.3). We first use the fact that the Rademacher complexity can be upper bounded by Gaussian complexity, which is defined as
G, (F(Ma, Mp)) = E,,
Σ c,w Eplao x.)] (TT)TTa)) iml
sup
n
with e NOO, 1). For any function class F. and any we have (Wainwright, 2019, section 5.2):
T RR)
AT (.))
V
31

---

Neural Collapse beyond the Unconstrainted Features Model
Thus it is sufficient to upper bound the Gaussian complexity of the function class:
1
Bn (F( Mp)) = Ee,
Eplaa("" ((TTT;)))) x:)]
sup (MM, 7
= Res
sup
Σ
G: T TE Bppla (u s;)]
(TTa:))
(1.0):||1113 M. M, n2 i=1
< Er;
sup
||wl12
eg Eepppaa mi)l
(p,0):]|ar01355 Mw. R.
Ma
sup (110
C EpEp (u as)]
M
Ee
sup P:B, 1
|| Bp (a (u))ll2
where we define Z(u)
1
Σ
colu x4)
vn 1=1
Mu
Eei
sup p:BB,110 M
EplaTa'Z(uu) a' a'Z(u))(w()
M Ee,
(Æpl(a (Ep(laTdd)1,, ((918)22((32222))))
sup 1110 11
M Bes
sup piB 1110011)MM, M
Mu Ees
(E (E2_1a(313,
sup pat M
M E
[p](((u)))
1101111MM
Note that Z(u) ~ NNO0 (())2) with p(u) = E. o(" 2;)²
(uTz))² < C...
Thus,
Mw-M.c3
(F Mp)) S
and the desired claim readily follows.
D.4. Proof of Theorem D.3
Proof of Theorem D.3. By Lemma 4.3, we have that
m mmin(W) 2 g-co
|WI < 2Bog 8)
Recall that (e, W) < B3 (log Thus,
1,110011 S 23Lx, () < (Bog 3)
Consequently, for each i, (f(2; W)), e (M,, with
Ma. 2Boggg 33),
M, Blog "
Hence, combining Lemma D.I and Lemma D.2, we get
Ct
log(2/0)
erTtese p, p.; < 22 1/2 (log /) + Blog 3)a 3)8V
+ 69g
2n
with probability > 1 - 5, which completes the proof.
32

---

Neural Collapse beyond the Unconstrainted Features Model
D.5. Proof of Lemma D.4
Proof of Lemma D.4. Using the definition of linearly separable data, we first show that, for each k, we could use one neuron to approximate the perfect classifier.
particular, having fixed let og = qeR, un log( ///e) ug. Then, we have
1
(ug x) 2
for supp(D(Hlew),
1 1 exp( log vale))
1+c/V9
1
c/v9
au x)
for supp(D(-lex)) with M'' * k.
1 + exp (log (vq/c))
It c/v
Define P. = ( (4 (ax mx)) and Po = Po * Tc, where y ~
By picking > q. we know that,
under the distribution Pa. TR(a) a and thus:
h (x) = (x) = Pu, +
Σ
PERREW
for any x E supp(P((lex) with PE. 2 Tte/ Pkr, S </v9
This implies that:
T40/v0
EolIfR(E; Wo) -1121 IlpKARE + Σ
PH - -xxl2 ex1 se
k'tk
We then have the following upper bounds:
|
R
Np. (x)1
(z)
= EalTR(a (G) Tx)] - E,, (u T x)]1 where (Ga, Gw) ~ N(0, Ipt041)
= E - (Tnl))(((((++(())))))) (G.) T x)] - Epo lao =)]1[²
< Ep. EclTe()o(((o(u((C@@ x)] Ep. luo T =)]]|" [BL.Bol(JJA@GGGCC((CCCCC =)]l[²
For the first term, we have
[|E__Eolrnaao(((ut (G.) - Ep (u 2)]1(² = || (llu
Ep. [ao
(u x)]
Boo((u + (G.) x)) - oup x;)²
S cot (Eal(G.Ia)) s gcdd
For the second term, we have
Ep Ee (Slaaasss <>11 C2GOCp..
Thus, we have the following bounds on the test error:
Ep []f Wo)
u13] 2ED [[n( por Wo) - y| + 2ED [[ Wo) fn(x;; Wo
2 + 4(gd - C3p)
Finally, we need to upper bound the second moment and the negative entropy of (p.. Wo). For the second moment term, direct computation gives:
(log
(V
Ep. || < q²+
+ 2
72
For the entropy term, have:
p + d
log S
Ep, log(po) s Eyg log(?c)
2
2
Taking Ç = and P1 = Åo concludes the proof.
33

---

Neural Collapse beyond the Unconstrainted Features Model
D.6. Proof of Lemma D.5
We first upper bound the approximated free energy (01, Wo) obtained by Stage of Algorithm using the optimality of P1, Wo-
(p1, Wo) - f Wo) + 1 9-1 E2 |||
g-1
|
+ 9 - log p1-
2
Next, by the construction in Lemma D.4 with e² 3 and by letting D be the empirical distribution of training samples, have
(7,, Wo) - (1 + 2( (8p)
M log(g3)²
Ep, [1101]]] S q²²
g-1
47²
p+d d 2
Ep, [log
log ( 1 /Tte)
which implies that
(p1, Wo) S Co.8 log 3,
(p), Wo) + 9-1
(46)
2 |Woll? S 17.-''og1 log
M² with Co (1 2( CiP)C 3 _ and C7 Co ptd + log 2n.
2
To conclude, we need upper bound (n1 Wo). which is obtained using the following intermediate result (proved in Appendix D.7) that controls the approximation error caused by the approximated second-layer.
Lemma D.7. Let (p1, Wo) be the unique global minimizer achieved by Phase 1 of Algorithm 1. Then, it holds that
HPp H < (4Co log B + + d 8T)) Cip² V2
VTR
R2
exp (R CR (Wo)
In particular, for any fixed B, we have
lim R-+
|H
HH 0.
Finally, to bound Ex Wo), write
(p1, Wo) (p,, Wo) + (n Wo) - (R, Wo)
1
< (p,, Wo) + (IIW)) Hp, - Y1 - 1| H Y
2n
1
1
s (pi, Wo) + - Hp: - W. H 122 +
W T H 112
n
2n
I (p, Wo) + T max (w
11p, H 11 + C (p1, Wo)
2E (p1, Wo) + 3 P
log(2x)) + H
R2 p1 1|2
2
where the last line follows from Lemma A.3. By Lemma D.7 we have that lim 11 1o,p-H,110 0, which means that
H
one can pick R large enough such that H|H HH S E Wo). Noting that the free energy is non-increasing when running Stage of Algorithm concludes the proof.
D.7. Proof of Lemma D.7
Proof of Lemma D.7. In the proof, we will write (p. W) instead of (P1, Wo) given that there is confusion. We start by showing that
J pi pi(a, u) dadu <
PV2
R²²
exp
(2/2pC78log (C Co) -
(47)
/
34

---

Neural Collapse beyond the Unconstrainted Features Model
Recall from Proposition 5.1 that p has the Gibbs form in (27),
p(a, u) -n(p)=1eep exp (-2rn(a)
1 (0 (arna))T(T"Tr) W(W H ()o("T
u) -2" [1aaa 2- + -llulB)
where (p) denotes the normalization constant. Thus, we have
1 1a
p(a, u) dadu
(-LTR(O)TW(N)) W(W -Y)o(xX -5(1a
( + +llulB) dadu
-Z((P)-1
exp
n
{ /supexpl
3 -T(()) up - ( 101
exp ( + |lull3) dadu.
- ZR(P) -1
We upper bound the various terms separately. First, we have
sup a,
Tn(a) W(W H Y)o(X u)
- n
max (W)] H -X "))l
3
(R + Co)C1ynnpomas
H 1|
n
B 2pc W)C, Co)o (W)
3B /pCr. 3C1 (R + Co) m (W),
where in the last passage we use (46).
Next, we lower bound the normalization constant as
3 Ta(a) W(W H - r)o(X u) - - (1lal12 + [lull²) dadu
Zn(p)= = exp exp
8 (n(a) T W(W H yo((TT1) wl) / exp
2 exp
sup
(1lal12 + [|w]18) dadu
a,
exp
-BV log 3C1 + Co) max (W))
/ exp
(lall + || (3) dadudbo
(p+d)/2
1 (W)
exp
SBV/ypC7-1Lol 3C m))(
2m
Finally, we bound
exp (-2
( (1|al13 + =|ul(B)
dadu
p
$ 12
exp
///2²) daj
exp
las1" 12 + [|u113))
daz ...... da
10.12
d2
1
< 2p
exp
2 |a11?
daj
2T
R
d2/2
1
trx
< 2p
a1 exp
-
da.
R
d22
exp
R2
2p
2T
R
Combining all the bounds. the desired result (47) follows.
35

---

Neural Collapse beyond the Unconstrainted Features Model
The following chain of inequalities holds
H - H < B [uo (u X) lal 2
Eo,o o ol |a (u x) u')l
lal (a))zl
(0' an independent copy of
0 la A |a'l *
Eo ( ( (T))) aal 10'1
Rollla | 1
u)
12 la trd
Eallla|
1
pla, u) dadu
|a/
C2p²v2
R²²
s E |0|
VTR
exp
23 2 10 (R + Co)o Co)omax max (W)
2
where the last passage follows from (47). Finally, we upper bound Epllla| |u|))
Epllaa [21 la gal - Bp Epllaal [2]
Ipllaaa] 122 |u|62
- 43E (p. W) +) + 4 + (p +p+ d) log 8T)
43 (Cop-1 log 8) + 4 (1 + (p d) log 8T),
where the third line follows from Lemma A.3 and the fourth line from (46). Combining these two bounds gives the desired result.
E. Additional Numerical Results
E.1. Extra Experiments Three-layer Networks
Figure plots the normalized balancedness metric N B(p, W) defined in (25) during training. The results clearly show that the network does not become balanced at convergence.
1.0
0.5
0.4
255
500
152
1000
1250
1550
1750
2000
902
1000
1992
2000
2003 1056
Epochs
Epochs
(a) MNIST
(b) CIFAR-100
Figure 2: Normalized balancedness (see (25)) as a function of the number training epochs, with each color representing an independent experiment.
Figure plots the approximately linear relation between the log of the NC1 metric and either the log the gradient norm (Figure 3a) or of the training loss (Figure 3b). This clearly shows the polynomial relation between the gradient norm/training loss and NC1 as in Corollary 4.4.
36

---

Neural Collapse beyond the Unconstrainted Features Model
1754
1750
750
12:10
130
13.14
13.49
250
390
33400
3210
13
18
10
33
33
Training
(a) Grad norm v.s. NC1
(b) Training loss v.s. NC1
Figure 3: Scatterplot of average gradient norm/training loss and average NC1 for training a three-layer network on MNIST.
The color bar indicates the training time step.
E.2. Experiments ResNet-18 and VGG-11
Figure plots the evolution of training loss, gradient norm and NC1 metric during training for ResNet-18 (Figure 4a) and VGG-11 (Figure 4b). The curve indicates that the decrease of the NC1 metric with training loss and gradient norm (as predicted by Corollary 4.4) also holds for real networks.
Anerag Trening
Aeerag Treinig
10
10a
30-3
10-1
30a1
1.0
(a) R8
VGG-18
Figure 4: Average training loss (blue), NCI (orange) and gradient norm (green) during SGD training for ResNet-18 and VGG-11 on CIFAR-10. We report the average for independent experiments, well as the confidence interval at 1 standard deviation.
Figures and 6 plot the relation between the log of gradient norm (Figures 5a and 6a), training loss (Figures 5b and 6b) and log of NC1. In the terminal phase, the log of gradient norm and training loss have an approximately linear relation with the log of NC1, which implies a polynomial relation between the gradient norm/training loss and NC1 as in three-layer networks. However, the early phase training exhibits a different non-lined lineir relation compared to three-layer networks.
37

---

Neural Collapse beyond the Unconstrainted Features Model
1211.
2211.
140
NC1
12:17
22117
18410
1310
1E:14
12:10
Gradient Norm
Training
(a) Grad norm v.s. NCI
Training loss v.s. NC
Figure 5: Scatterplot of gradient norm/training loss and NC1 for training ResNet-18 CIFAR-10. The plot starts from
epoch 51, right after the learning rate is reduced by factor 10. The color bar indicates the training time step.
12117
12
128
1A0
NC1
124.7
121
Gradient Norm
Training Loss
(a) Grad NC1
Training loss NC
Figure 6: Scatterplot of average gradient norm/training loss and average NCI for training VGG-11 CIFAR-10 10. The plot starts from epoch 51, right after the learning rate is reduced by factor 10. The color bar indicates the training time step.
38


Reviews:
[{"review_id": "PS5DqlFCm4", "reviewer": "Reviewer_GqGb", "summary": "This paper studies the emergence of variability collapse in the penultimate layer representations of a three layer mean-field NN. The authors show that points for which the gradient norm is small (approximate stationary points) also show variability collapse, and that the level of variability collapse is controlled by the gradient norm. Next, they show that gradient flow on the parameters of a three layer mean-field NN wrt MSE loss converges to approximate stationary points. Finally they consider the case of linearly separable data and show that variability collapse is also related to small test error in this case.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for finding that our paper advances the study of neural collapse and recommending a strong acceptance! We are happy to have a further discussion in case additional questions or comments come up.", "claims_and_evidence": "All of the above claims come with theoretical proofs and are convincing. The authors also consider differences between their entropy regularized free energy landscape and the landscape of L2 regularized networks, leading to solutions with imbalanced layers.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "I followed the proof sketches in the main paper but did not work through the appendix.", "experimental_designs_or_analyses": "Yes, the experiments measuring balancedness of the solution seem to be sound.", "relation_to_broader_scientific_literature": "This paper advances the study of neural collapse in two ways - first by moving beyond the unconstrained features model they allow for considering the role of the data/kernel. Next, they consider approximate stationary points and show that variability collapse still holds. This is in contrast with prior works that typically consider stationary points where the gradient is exactly zero.", "essential_references_not_discussed": "The discussion is comprehensive", "other_strengths_and_weaknesses": "None", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "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": "gmVSMAmpEP", "reviewer": "Reviewer_g3N9", "summary": "The paper theoretically studies the phenomenon of neural collapse in classification, focusing on its most basic property: the vanishing of within-class variability. Unlike data-agnostic prior work analysing the unconstrained features model (UFO), this work adopts a data-specific perspective by considering a three-layer neural network with a feature-learning component in the mean-field regime. The authors establish a connection between NC1 and the loss landscape, proving that points with small empirical loss and gradient norm are roughly NC1 solutions. They also demonstrate that gradient flow dynamics on the three-layer model converges to NC1 solutions. Finally, they show that NC1 and vanishing test error can coexist for well-separated data distribution, connecting neural collapse to generalisation.\n\n**Update after rebuttal**\n\nI thank the authors for the detailed response. I consider the paper a solid and relevant contribution, so I recommend its acceptance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review and positive feedback. We address questions and concerns below.\n\n**1. Difficulty in extending results to CE loss, and expected qualitative results.**\n\nThank you for the question. It is indeed interesting to extend our results to CE loss and we summarize two major difficulties below:\n\n* The intuitive difference between MSE loss and CE loss is that small MSE loss implies that the output of the neural network is roughly the same for all data points in the same class, as the label of data points in the same class is the same. In our model, we stack a linear layer after the feature learning part, which implies that data points in the same class having approximately the same output is a **necessary but not sufficient condition** for NC1 to happen. To see this, suppose $h(x_1), h(x_2)$ are the features of two data points in the same class. Then, $||h(x_1)-h(x_2)||\\_2^2 \\geq  \\frac{1}{\\lambda_\\max(W^\\top W)}||W^\\top h(x_1)-W^\\top h(x_2)||_2^2.$ Now, in order to extend our results to CE loss, one would have to rule out the possibility that different data points in the same class could have different logits, and it is unclear how to do so in our setting even when the loss is small. However, we believe that this is just a technical issue coming from the fact that we consider a linear classifier in the last layer. \n\n* Another technical issue is that our results require the characterization of features in Equation (6). This only holds for MSE loss, and a different characterization would be needed for CE loss.\n\nThat being said, we do expect a similar characterization to hold for CE loss. As a first step towards handling the CE loss, we note that there are indeed papers in the mean-field regime training with CE loss [1,2] or training with general losses that include CE [3]. However, all those papers consider two-layer networks, [1] requires no regularization terms, and [2] focuses on the specific problem of learning parities.\n\n**2. The technical reason for considering the approximated model truncating the second layer in Section 5, and whether the following conclusions would be affected if the full model was considered.**\n\nHaving an approximated second layer implies that there exists a unique global optimum for the Gibbs minimizer as in Equation (27) and such a global optimum is achievable by mean-field langevin dynamics [3]. Specifically, one needs that $\\nabla_\\theta \\frac{\\delta}{\\delta \\rho} \\mathcal{L}_n(\\rho)(\\theta) $ is Lipschitz in $\\theta, \\rho$ (cf. Proposition 2.5 in [4]), and without truncation this cannot be guaranteed. However, we note that our results in Section 5 hold uniformly for large enough $R$, which means that we also expect the same conclusion for the original model which corresponds to $R = +\\infty.$ Rigorously speaking, one cannot interchange the order of the limits (in the time $t$ and the truncation parameter $R$), but we believe this to be a purely technical issue.\n\n**3. Typo in the title.**\n\nThanks for spotting this, we will correct it in the revision.\n\n[1] Chizat, Bach, Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss.\n\n[2] Suzuki, Wu, Oko, Nitanda, Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond.\n\n[3] Suzuki, Wu, Nitanda, A. Mean-field langevin dynamics: Time-space discretization, stochastic gradient,\nand variance reduction.\n\n[4] Hu, Ren, Siska, Szpruch, Mean-field langevin dynamics and energy landscape of neural networks.", "claims_and_evidence": "Although I have not thoroughly checked the proofs, the claims seem reasonable and supported by rigorous mathematical analysis. The authors make use of common assumptions in the literature, present proof sketches and a fair discussion of related work.", "methods_and_evaluation_criteria": "The methods and evaluation criteria make sense for the theoretical analysis conducted in this work.", "theoretical_claims": "I did not check the proofs in detail.", "experimental_designs_or_analyses": "The paper presents comprehensive numerical experiments using MNIST and CIFAR-100 to validate the theoretical results. The experiments are well-designed, and the results are consistent with the theoretical claims.", "relation_to_broader_scientific_literature": "The paper seems to be well-positioned within the broader scientific literature. The authors provided a comprehensive review of the related work and are very clear about the difference between their work and previous work. They also highlight the fact that their work can offer new insights into gradient-based optimization in deep learning.", "essential_references_not_discussed": "I am unaware of any essential references that were not discussed in the paper.", "other_strengths_and_weaknesses": "Strengths:\n\n- The paper bridges neural collapse to loss landscape, gradient flow dynamics and generalisation. It presents a rigorous and fairly general analysis. \n\n- The idea of taking a two-layer neural network in the mean field regime and stacking a linear layer is quite interesting.\n\n- The paper is well-written, and the results are well-stated. It is also well-motivated (going beyond UFO) and clearly tries to review, connect and extend previous literature. \n\nWeaknesses:\n\n- It is probably not feasible analytically, but a small remark on the difficulties of extending the results to cross-entropy loss would be interesting. I believe the mean-field regime has been explored only for the squared loss, which would be a complication for such an analysis (see questions for authors).", "other_comments_or_suggestions": "I believe there is a typo in the paper's title: \"Unconstrainted\" should be \"Unconstrained\".", "questions_for_authors": "1)  If a similar analysis for cross-entropy loss were possible, would the authors expect drastic changes in the qualitative picture?\n\n2) Could you clarify the technical reason for considering the approximated model truncating the second layer in Section 5? Would the following conclusions be affected if the full model was considered?", "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": "CdRHJ1QdFx", "reviewer": "Reviewer_AaqF", "summary": "The authors study the how a certain aspect of neural collapse (NC) - namely, within-class variability tending to zero, can be provably associated with convergence to nearly stationary points of the loss function in noisy gradient flow dynamics for a 3 layer neural network, where the last two layers have linear activations.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review. We address concerns below.\n\n**1. Lack of soundness in experimental design.**\n\nWe followed the suggestions and did additional experiments, please download pdf in https://github.com/conferenceanonymous152/icml25\n\n(a) For three-layer networks, Figure 3 shows the linear relation between log of NC1 and log of gradient norm/training loss throughout training. Different points correspond to different epochs of training (purple ones in the top-right correspond to early stages, and yellow ones in the bottom-left correspond to late stages), and we average the results of 4 independent experiments. The plots indicate a polynomial relation between the quantities as predicted by Equation (17).\n\n(b) For practical models, we show that ResNet-18 and VGG-11 trained on CIFAR-10 exhibit a similar behavior as three-layer networks. Figure 4 is similar to Figure 1 in the paper. The scatterplots in Figures 5, 6 plot the $\\log-\\log$ relation between NC1 and gradient norm/training loss, and they indicate a polynomial relation in the terminal phase of training (while the early phase of training may be a little different). \n\n**2. NC2-NC3.**\n\nTo prove NC2-NC3, [1] uses that the number of linear layers is large. If we stack many linear layers on top of the feature-learning layer and run gradient flow, we expect that NC2-NC3, as well as NC1, can be shown. As for NC1, our approach could be adapted as follows. Let the network be $f(x; \\rho, W_2, \\dots, W_L) = W_L \\dots W_2 h_\\rho(x),$ and define $ W^\\top =W_L \\dots W_2$. Then, the characterization in Theorem 4.2 still holds, since the condition of Definition 4.1 only involves the gradient computed over $\\rho$. The lower bound in Lemma 4.3 holds and the upper bound becomes $\\sigma_{max}(W) \\leq O( (\\log \\beta)^L ) = o(\\beta)$, which implies Corollary 4.4. Theorem 4.8 can be proved by letting  $t_*$ be the minimal first hitting time for $W_2, \\dots, W_L$ and $\\rho$ and following a similar strategy. Finally, since approximate NC1 holds for $H_\\rho(x),$ it also holds for $W^\\top H_\\rho(x)$, for any $W$ satisfying Lemma 4.3.\n\nAs for NC2-NC3, we note that $W_2, \\dots, W_L$ are balanced for all stationary points $(\\rho, W_2, \\dots, W_L)$. Thus, if the condition number of $W$ is bounded by a constant independent of $L$ (as in Theorem 3.1 of [1]), one can show NC2-NC3 following the approach in [1]. We note that **$W_1 := \\mathbb{E}_\\rho[a a^\\top]$ is not balanced with $W_2, \\dots, W_L$** due to Lemma 4.5 of our work. Thus, to show NC1 for a network with many linear layers, the results of [1] cannot be applied. \n\n**3. Major differences compared to previous work, additional entropic regularization term and question about line 112.**\n\nThe entropic regularization term is not implicitly induced by SGD, but it is explicitly added in the form of a Brownian noise. In fact, we consider noisy gradient flow for training, i.e., mean-field Langevin dynamics. This is standard in the mean-field regime [2, 3], and noisy gradient flow is the limiting dynamics of noisy SGD with step size $\\eta$ and any batch size (including full-batch) on networks of width-$N$ as $\\eta \\to 0$ and $N \\to \\infty$ [3].\n\nWhat distinguishes our work is that we show NC1 for a network that (i) is trained end-to-end and (ii) exhibits feature learning throughout training. In contrast, previous work focuses either on UFM (which does not capture the effect of data) or on a dynamics reaching near-interpolation in the NTK regime [1] (which does not exhibit feature learning [4]). Furthermore, the occurrence of NC1 in [1] relies on the balancedness of all linear layers, which may not hold, see Lemma 4.5. Instead, our result only needs small training loss and gradient norm, which is commonly satisfied as long as the training dynamics converges. Finally, we show for the first time the co-occurrence of NC1 and vanishing test error for a class of data distributions. For additional comparison of our work to [1], please see point 6 of the response to Reviewer sgNN.\n\n**4. Connection between theoretical results and real-world networks.**\n\nSee point 1.\n\n**5. Line 116.**\n\n$Law(a,u)$ means the joint measure of random variables $a,u.$\n\n**6. Assumptions (A3).**\n\nReLU does not satisfy (A3), but we still expect our results to hold by taking the limit of a sequence of approximations to ReLU. We note that (A3) is purely technical (it ensures the well-posedness of the Wasserstein gradient flow) and standard in works [2, 3] considering the mean-field regime. \n\n[1] Jacot, Sukenik, Wang, Mondelli, Wide neural networks trained with weight decay provably exhibit neural collapse.\n\n[2] Suzuki, Wu, Nitanda, Mean-field langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction.\n\n[3] Mei, Montanari, Nguyen, A mean field view of the landscape of two-layer neural networks.\n\n[4] Yang, Hu, Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks.", "claims_and_evidence": "The claims made in the submission seem to be well supported by formal proofs.", "methods_and_evaluation_criteria": "Since the paper is mainly theoretical in nature, the claims relate to bounds and convergence properties. The evaluation criteria are therefore mainly the voracity of the proofs. The empirical evidence are only meant to show that even imbalanced networks tend towards small loss/gradient norm and NC1 together, and for that they roughly make sense.", "theoretical_claims": "I checked the main theorems given in the paper, but did not re-derive them.", "experimental_designs_or_analyses": "I've checked the soundness of the experiments given, and found them to be lacking, in the sense that the experiments do not verify the claims made in the paper. \nWhile Fig.1 shows a qualitative behavior that aligns with the statement that NC1, gradient norm, and the empirical loss all converge to small values together, they do not constitute a sufficient support for these claims in practice. Two things can be done when providing empirical results for a theory paper:\n1) All statements are exact under the assumptions made, and no further verification is needed, aside from perhaps a numerical simulation that shows, for instance, that in a completely controlled setting, SGD on the 3 layer network with linearly separable data (say very clustered gaussian mixture) converges to NC1 with predicted dependence on the training/test loss that is given by gradient flow on the free energy.\n2) An attempt to make a broad statement connected to real networks beyond the setting studied, not knowing all the constant factors, but showing that the relation between NC1/train loss/gradient norm is maintained by looking at quantities that scale with (for example the $\\beta$ parameter) by performing a scan over many values.\n\nThis paper chose neither, and only shows a qualitative result, that I feel does not really contribute much, given the theoretical results.", "relation_to_broader_scientific_literature": "The results of these authors relate to the broad phenomenon of neural collapse which is observed in trained classifiers  across different networks and datasets, and attempts to connect a specific NC property in a tractable setup with the loss landscape both at training and test, thus relating NC to generalization performance, which is of broad interest.", "essential_references_not_discussed": "None that I am aware of.", "other_strengths_and_weaknesses": "**Strengths**:\n\n- The paper is well written and fairly easy to follow, apart from minor clarifications that might be useful to see.\n- The theoretical results seem to be correct.\n\nThe main contributions are extending previous results. In particular, under mild assumptions of small gradient norm and small loss, the authors show that a 3-layer network with the last two layers being linear converge under gradient flow to NC1.\n\nThis relaxes assumptions from previous works, specifically:\n\n1. Remove the balancedness assumption on all the linear layers and show NC1 occurs even for unbalanced layers\n\n2. Mean field in the regime of only the first layer width tending to infinity, while the width of the second layer remains of constant order.\n\nThey further show that for linearly separable (and bounded) data the test error decreases in correspondence with NC1 convergence.\n\nThese results indeed shed further light on the phenomenon of neural collapse, and attempts to connect it further with generalization performance.\n\n**Weaknesses:**\n\n1. As opposed to previous works (Jacot et al, Sukenık et al.) there is no discussion regarding NC2 and NC3.\n\n2. It seems that the major difference between this and previous works is the extension of the two linear layer setting to the SGD case, where noise induces an effective entropy regularization term in the dynamics.\n\n3. The connection made between the theoretical results and real-world networks is not studied seriously, but seems a bit of an afterthought.\n\nThe combination of these weaknesses makes the contribution, in my opinion, overly incremental.", "other_comments_or_suggestions": "I did not find any glaring typos.", "questions_for_authors": "116 - What does Law(a,u) mean? it is never specified clearly.\n\n112 - I think the phrasing implies you introduce an entropy regularization term but it comes simply from performing SGD on the predefined loss, so the dynamics are done on a free energy which includes the entropy penalty term, right?\n\nDoes condition A3 mean the results don’t hold for ReLU? how would they change?", "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": "Op9UqSGfOx", "reviewer": "Reviewer_sgNN", "summary": "This paper provides a theoretically rigorous proof on the neural collapse (NC) phenomenon with a three-layer neural network in the mean-field regime under mean-square-error (MSE) loss. It shows that under gradient flow (GF), the within-class variability (namely NC1) vanishes as the time $t$ tends to infinity. Also, it establishes an upper bound on the vanishing test error, ensuring both NC and good generalization performance can be achieved simultaneously.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the valuable and detailed comments, we address concerns below.\n\n**1. Supplementary materials.**\n\nWe will add this in the revision. \n\n**2. NC not surprising under MSE loss.**\n\nWe politely disagree with the claim that “it is quite intuitive that the features concentrate at their class mean when the train loss is small under MSE loss.” It is true that small MSE loss implies that the output of the neural network is roughly the same for all data points in the same class. However, this is only a **necessary but not sufficient condition** for NC1 to occur. The technical subtlety is that, since the feature dimension is larger than the number of classes, there might be components of the features in the space orthogonal to the last linear layer. **Without the balancedness condition**, it is not obvious to show that these orthogonal components vanish at stationary points. Intuitively, these components make the weights more isotropic, thus they could potentially decrease the entropic penalty, but increase the L2 penalty, meaning that the tradeoff needs to be carefully considered. Thus, our results in Section 4.2 require a suitable characterization of the learned feature and careful manipulations to show that the orthogonal components are indeed vanishing with training loss and gradient norm.\n\n**3. Extension to CE loss.**\n\nPlease see point 1 of our response to Reviewer g3N9. \n\n**4. Measuring rate of convergence.**\n\nWe provide a convergence rate for the loss in Equation (24). This is already non-trivial given the non-convex nature of the problem and to our best knowledge, no explicit rate of convergence is known in the mean-field regime beyond the 2-layer case (which has a convex free energy). Characterizing the convergence rate of the gradient norm is harder, as it requires controlling the curvature along the whole trajectory.\n\n**5. Linear separability assumption in Theorem 5.3.**\n\nIn general, it is not true that NC1 and vanishing test error co-occur under MSE loss without any assumptions on the data distribution. In fact, the occurrence of NC1 implies that the model overfits the data, and overfitting is not always benign without additional assumptions: classic benign overfitting results for simple models such as linear regression require a decay in the data spectrum and proper alignment between data spectrum and signal. In this sense, showing the co-occurrence of NC1 and vanishing test error may be regarded as a harder problem than benign overfitting, and the latter phenomenon has not been investigated in the mean-field regime to our best knowledge.\n\n**6. Comparison of our work to [1].**\n\nFirst, [1] operates in the NTK regime where the behavior of the weights is fundamentally different from the mean-field regime that we consider. More precisely, the proof of Theorem 4.4 in [1] distinguishes two phases in training: in the first phase, the loss becomes small and no feature learning takes place; in the second phase, the weight matrices in the linear part of the network become balanced and NC1 occurs. In contrast, our work operates in the mean-field regime where feature learning takes place from the beginning and, therefore, there is no need for a second phase of the dynamics: NC1 occurs as soon as the loss and its gradient are small. Our experiments on ResNet-18 and VGG-11 also indicate this relation (see point 1 of our response to Reviewer AaqF for more details).\n\n\nSecond, the balancedness condition in [1] is not just a technical requirement, but the key reason for the occurrence of NC1. In fact, the balancedness condition rules out the possibility of having any components of the features in the orthogonal space to the last linear layer. However, our results both theoretical (see Lemma 4.5) and experimental (see the end of Section 4.2) conclusively demonstrate that NC1 holds even if the network is not balanced.\n\n**7. Impact of weight decay.**\n\nFirstly, non-zero weight decay is necessary to ensure non-vanishing $W$ and features at all stationary points. Secondly, the weight decay affects the order of the empirical loss we converge to: when $\\lambda_\\rho = \\beta^{-1}$, the empirical loss at convergence is of order $\\beta^{-1}$. This is reflected in Corollary 4.9, as better NC1 guarantees (smaller $\\delta_0$) require smaller weight decay $\\beta^{-1} \\leq poly(\\delta_0).$\n\n**8. Extension to deeper/more complex networks.**\n\nWe expect that our results can be extended to a deep network with many linear layers, see point 4 of our response to Reviewer AaqF. In contrast, analyzing deep networks with many nonlinear layers would require fundamentally different techniques and we regard it as an exciting future direction. Towards this goal, two possible strategies are (i) to characterize features at an $\\epsilon_S$-stationary point, or (ii) to directly consider the evolution of features during training.\n\n[1] Jacot, Sukenik, Wang, Mondelli, Wide neural networks trained with weight decay provably exhibit neural collapse.", "claims_and_evidence": "The theoretical results are proven in the main text as well as in the appendix. Experiment results also show a collapse in NC1.", "methods_and_evaluation_criteria": "Yes, the experiment is done on a 3-layer MLP, matching the setting in the theory. Standard datasets MNIST and CIFAR-10 are used.", "theoretical_claims": "I only checked the proof sketch in the main text. They seem reasonable and straightforward. Due to time limit and my lack of expertise in mean-field dynamic theory, I am not able to check all proofs from line to line in the appendix. \n\nHowever, if other reviewers could point out a potential mistake in the proof, I would check it closer and participate in the discussion.", "experimental_designs_or_analyses": "The experiment results match and validate the theoretical claim.", "relation_to_broader_scientific_literature": "I think the contribution of this paper is limited, when compared to broader scientific literature. \n\nIt is known that NC is very significant when the network is trained under MSE loss. In particular, it is quite intuitive that the features concentrate at their class mean when the train loss is small under MSE loss. (It would be a different story when other loss, say cross entropy loss, is considered.) Hence, the qualitative statement in the main result (corollary 4.9) is not very surprising. The statement could have been more significant when it could also quantitative, like measuring the (non-)asymptotic rate of convergence etc. \n\nAnother result (Theorem 5.3) on the vanishing test error is nicer in this sense, however, it requires the data to be almost linearly separable.  \n\nMore importantly, [1] already extended the MSE NC result from unconstraint feature model (UFM) to a deeper one with several linear last layers, which is of the same goal of this paper. This paper has also discussed deeply on the comparison between [1] and their work, which I really appreciate, but I feel like the main improvement of this paper over [1] is just removing the weight balancedness condition in [1], which is only a technical requirement. Also, [1] considered networks with any depth while this paper only considered a 3-layer network.\n\nThat is why I think this paper has limited contributions related to the broader scientific literature. \n\n\n---\n[1] Jacot, Arthur, et al. \"Wide neural networks trained with weight decay provably exhibit neural collapse.\" arXiv preprint arXiv:2410.04887 (2024).", "essential_references_not_discussed": "I think most related works have been discussed thoroughly in the paper.", "other_strengths_and_weaknesses": "This paper does not require the weight balancedness assumptions used in [1].\n\nThe theorems and lemmata are presented and elaborated in a clear manner, which I really appreciate. \n\n---\n[1] Jacot, Arthur, et al. \"Wide neural networks trained with weight decay provably exhibit neural collapse.\" arXiv preprint arXiv:2410.04887 (2024).", "other_comments_or_suggestions": "I have no other comments.", "questions_for_authors": "I would like to ask:\n\n1. Could you explain more about the difference between the result from this paper and [1]? \n\n2. Related to [1] where they emphasized weight decay (WD) in the title, how could one interpret the effect of WD on the main result, say Corollary 4.9?\n\n3. Are there any technical difficulties to extend the setting to deeper and more complex networks?\n\n4. What is the hardest part when one tries to extend the argument in the mean-field setting to other losses, say CE?\n\n\n---\n[1] [1] Jacot, Arthur, et al. \"Wide neural networks trained with weight decay provably exhibit neural collapse.\" arXiv preprint arXiv:2410.04887 (2024).", "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:

Meta-review: This paper studies the emergence of neural collapse in penultimate layer representations. Using a three-layer neural network in the mean-field regime, the authors demonstrate that (1) approximate stationary points exhibit neural collapse (i.e., within-class variability collapse), (2) gradient flow converges to these approximate stationary points, and (3) neural collapse is associated with small test error in this setting.

Overall, the paper offers a good theoretical analysis that advances our understanding of neural collapse. While some concerns have been raised regarding the limited scope—specifically, the focus on a three-layer MLP architecture in the mean-field setting with MSE loss—these restrictions may reflect the broader challenge our community faces in developing rigorous theoretical descriptions of deep learning models.
Decision: Accept (spotlight poster)