Neural (Tangent Kernel) Collapse
Authors: Mariia Seleznova
Dana Weitzner
Raja Giryes
Gitta Kutyniok
Hung-Hsu Chou
Venue: NeurIPS 2023 poster (2023)
Keywords: Neural Collapse
Neural Tangent Kernel
NTK alignment
Local Elasticity
Gradient Flow
TL;DR: The Neural Tangent Kernel alignment drives Neural Collapse.

Abstract: This work bridges two important concepts: the Neural Tangent Kernel (NTK), which captures the evolution of deep neural networks (DNNs) during training, and the Neural Collapse (NC) phenomenon, which refers to the emergence of symmetry and structure in the last-layer features of well-trained classification DNNs. We adopt the natural assumption that the empirical NTK develops a block structure aligned with the class labels, i.e., samples within the same class have stronger correlations than samples from different classes. Under this assumption, we derive the dynamics of DNNs trained with mean squared (MSE) loss and break them into interpretable phases. Moreover, we identify an invariant that captures the essence of the dynamics, and use it to prove the emergence of NC in DNNs with block-structured NTK. We provide large-scale numerical experiments on three common DNN architectures and three benchmark datasets to support our theory.

Neural (Tangent Kernel) Collapse
Mariia Seleznova * Dana Weitzner
Raja Giryes² Gitta Kutyniok' Hung-Hsu Chou'
"Ludwig- Mudwiig-Mximmllinns.. Universität München
TTel Aviv University
Abstract
This work bridges two important concepts: the Neural Tangent Kernel (NTK), which captures the evolution of deep neural networks (DNNs) during training, and the Neural Collapse (NC) phenomenon, which refers to the emergence of symmetry and structure in the last-layer features of well-trained classification DNNs. We adopt the natural assumption that the empirical NTK develops block structure aligned with the class labels, i.c.. samples within the same class have stronger correlations than samples from different classes. Under this assumption, we derive the dynamics of DNNs trained with mean squared (MSE) loss and break them into interpretable phases. Moreover, we identify an invariant that captures the essence of the dynamics, and use it to prove the emergence of NC in DNNs with block-structured NTK. We provide large-scale numerical experiments on three common DNN architectures and three benchmark datasets to support our theory.
Introduction
Deep Neural Networks (DNNs) are advancing the state of the art in many real-life applications, ranging from image classification to machine translation. Yet, there is no comprehensive theory that can explain multitude of empirical phenomena observed in DNNs. In this work, we provide theoretical connection between two such empirical phenomena, prominent in modern DNNs: Neural Collapse (NC) and Neural Tangent Kernel (NTK) alignment.
Neural Collapse. NC 39 emerges while training modern classification DNNs past zero error to further minimize the loss. During NC, the class means of the DNN's last-layer features form a symmetric structure with maximal separation angle, while the features of each individual sample collapse to their class means. This simple structure of the feature vectors appears favourable for generalization and robustness in the literature 12 31 40 47). Though NC is common in modern DNNs, explaining the mechanisms behind its emergence is challenging, since the complex non-linear training dynamics of DNNs evade analytical treatment.
Neural Tangent Kernel. The NTK 30 describes the gradient descent dynamics of DNNs in the function space, which provides dual perspective to DNNs' evolution in the parameters space. This perspective allows to study the dynamics of DNNs analytically in the infinite-widtt limit, where the NTK is constant during training [30 Hence, theoretical works often rely on the infinite- width NTK
to analyze generalization of DNNs 1 20 28 49). However, multiple authors have argued that the
infinite-width limit does not fully reflect the behaviour of realistic DNNs (2 10 22
27 36
since
constant NTK implies that no feature learning occurs during DNNs training.
NTK Alignment. While the infinite-width NTK label-agnostic and does not change during training, the empirical NTK rapidly aligns with the target function in the early stages of training In the context of classification, this manifests itself as the emergence of a block structure
45
*Correspondence to: Mariaa Seleznova (selezemath. 1mu de).
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

in the kernel matrix, where the correlations between samples from the same class are stronger than between samples from different classes. The NTK aligmment implies the so-called local elasticity of DNNs' training dynamics, i.e., samples from one class have little impact on samples from other classes Stochastic Gradient Descent (SGD) updates [33. Several recent works have also linked the local elasticity of training dynamics to the emergence of NC (3 53). This brings us to the main question this paper: Is there a connection between NTK alignment and neural collapse?
Contribution. In this work, we consider model of NTK alignment, where the kernel has a block structure, i.., it takes only three distinct values: an inter-class value, an intra-class value and a diagonal value. We describe this model in Section 3 Within the model, we establish the connection between NTK alignment and NC, and identify the conditions under which NC occurs. Our main contributions are as follows:
We derive and analyze the training dynamics of DNNs with MSE loss and block-structured NTK in Section 4 We identify three distinct convergence rates in the dynamics. which correspond three components of the training error: eπror of the global mean, of the class means, and of each individual sample. These components play a key role in the dynamics.
We show that NC emerges DNNs with stockestructured NTK under additional assumptions in Section 5.3 To the best of our knowledge, this is the first work to connect NTK alignment and NC. While previous contributions rely on the unconstrained features models (21 38 48) other imitations of DNNs' training dynamies [53] to derive NC (see Appendix A for detailed discussion of related works), we consider standard gradient flow dynamics of DNNs simplified by our assumption the NTK structure.
We analyze when NC does or does not occur in DNNs with NTK alignment, both theoreti- cally and empirically. In particular, we identify an invariant of the training dynamics that provides a necessary condition for the emergence of NC in Section 5.2 Since DNNs with block-structured NTK do not always converge to NC, we conclude that NTK alignment is more widespread phenomenon than NC.
We support our theory with large-scale numerical experiments Section 6 Source code to reproduce the results is available in the project's '' GitHub repository
2
Preliminaries
We consider the classification problem with C' E N classes, where the goal is to build a classifier that returns a class label for any input x E X. In this work, the classifier is DNN trained on a dataset {(ri, yı ) where T; X' are the inputs and y; RC are the one-hot encodings of the class labels. We view the output function of the DNN f X - RC features h : X' - Ih"" and a linear classification layer parametrized by weights W E CxCx and biases be E _ Then the logits of the training data X (a.) be expressed as follows:
composition of parametrized last-layer
f(x) = WH+ b1
(I)
where H E | are the features of the entire dataset stacked columns and 1N E R-W a vector of ones. Though we omit the notion of the data dependence in the text to follow, i.c. we write H without the explicit dependence on X, we emphasize that the features H are function of the data and the DNN's parameters, unlike in the previously studied unconstrained feature models [2 38 48
We assume that the dataset is balanced. i.e. there m N/C training samples for each class. Without loss of generality, we further assume that the inputs are reordered so that Tic- 1ot-1mm++....-... Tem belong to class for all c E [C]. This will make the notation much easier later on. Since the dimension of features is typically much larger than the number of classes, we also assume C' in this work.
2.1 Neural Collapse
Neural Collapse (NC) is an empirical behaviour of classifier DNNs trained past zero error
39
Let (h) == N- Σi h (xx;) denote the global features mean and (N), := m E Eclass hxx)))
e E [C] be the class means. Furthermore, define the matrix of normalized centered class means as M [l) (T)c//l [ E RaxC where (T), = (h), - (h), e E [C]. We say that DNN exhibits NC if the following four behaviours emerge the training time increases:

---

(NC1) Variability collapse: for all samples z; from class E [C], where i [m], the penultimate
layer features converge to their class means, i.c [|h(z;) (h)ell2
(NC2) Convergence to Simplex Equiangular Tight Frame (ETF)) for all [C], the class means converge to the following configuration:
C (Ic (Ic- c (old). C C-1
1
(I()) - (h) (( (h)ll2 - 00, M
(NC3) Convergence to self-duality: the class means M and the final weights WT converge to each other:
M/ - w / Tell 0
(NC4) Simplification to Nearest Class Center (NCC): the classifier converges the NCC decision rule behaviour:
argmax argmax(Wh(4) (Wh(x) + b). argmin|(((2) - (h)ell2-
Though NC observed in practice, there is currently conclusive theory on the mechanisms of its emergence during DNN training. Most theoretical works on NC adopt the unconstrained features model, where features H are free variables that can be directly optimized [2] 38 48 Training dynamics of such models do not accurately reflect the dynamics of real DNNs, since they ignore the dependence of the features on the input data and the DNN's trainable parameters. In this work, we make a step towards realistic DNN dynamics by means of the Neural Tangent Kernel (NTK).
2.2 Neural Tangent Kernel
The NTK @ of DNN with the output function X - R.C and trainable parameters w R.P (stretched into a single vector) is given by
(2, mz) =
I4, X, k, E [C].
We also define the last-layer features kernel @h, which is a component of the NTK corresponding to the parameters up to the penultimate layer, as follows:
h (Z,, = (V (x,)) Vwho(z))), Iy E Ex K E [n].
(3)
Intuitively, the NTK captures the correlations between the training samples in the DNN dynamics. While most theoretical works consider the infinite-width limit of DNNs [0 52), where the NTK can be computed theoretically, empirical studies have also extensively explored the NTK of finite-widh
networks 199 36 45 49 Unlike the labol-agnostic ffinite-width NTK, the empirical NTK aligns with the labels during training. We use this observation in main assumption (Section
2.3 Classification with MSE Loss
We study NC for DNNs with mean squared error (MSE) loss given by
(W,H.b) 1 (1(()) Y
(4)
where Y RCNN is matrix of stacked labels Vi. While NC was originally introduced for the cross-entropy (CE) loss (39), which is more common in classification problems, the MSE loss is much easier analyze theoretically. Moreover, empirical observations suggest that DNNs with MSE loss achieve comparable performance to using CE [14 29 41 which motivates the recent line of
research MSE-NC 21 38
Block Structure of the NTK
Numerous empirical studies have demonstrated that the NTK becomes aligned with the labels Y during the training process 45 This alignment constitutes feature learning and associated with better performance of DNNs 9 13] For classification problems, this means that the empirical

---

(x,,(X)
b) E.9(())
|00 1
Aligament
95
0L5
400
A1 Ap
100 301
300
12403
Figure I: The NTK block structure of ResNet20 trained on MNIST. a) Traced kernel Σ computed on random data subset with 12 samples from each class. The samples are ordered as described in Section 2 so that the diagonal blocks correspond to pairs of inputs from the same class. b) Traced kernel Σ k=1 O (X) computed on the same subset. e) Norms of the kernels (X)) for all k, E [C]. d) Norms of the kernels (X) for all k s E [n]] The color bars show the values in each heatmap as fraction of the maximal value in the heatmap. e) The alignment of the traced
kernels from panes and with the class labels.
NTK develops an approximate block structure with larger kernel values corresponding to pairs of samples (a :r)) from the same class 44 Figure D shows example of such a structure emergent MNIST@Motivan Motivated by these observations, we assume that the NTK and the last-layer features kernel exhibit block structure, defined follows:
the empirical NTK of ResNer20 trained on MNIST
Definition 3.1 (Block structure of a kernel). We say kernel
x X + RK*K has block
structure associated with (A1, A2, Aa), i X2 > d 2 0 and
(x,, x) AL
(x),, = >IIK, (5)) A3IK:
(5)
where xf and v are two distinct inputs from the same class, and uf is an input from class e' +
Assumption 3.2. The NTK θ X X RCXC has block structure associated with (Na: Yer (n). and the penultimate kernel Oh X X + Raxn has a block structure associated with (Ra, Ries Kn).
This assumption means that every kernel (x, = (0,,,(a;, xy kK)]ee[[] corresponding to an output neuron fi.. le E [C] and every kernel (X) corresponding to a last-layer neuron hp. P E [n] is aligned with Y Y (see Figure 1 panes a-b). Additionally, the "non "non-diagone -iag-naa kernels "X..((") and (X)). 3...A. & f are equal to zero (see Figure 1 panes c-d) Moreover, if > and Ne > Kn-
Assumption 3.2 can be interpreted as local elasticity of DNNS, defined below.
Definition 3.3 (Local elasticity (23) A classifier said be locally elastic (LE) if its prediction or feature representation on point of from class e. E [C] is not significantly affected by performing SGD updates on data points from classes c! cc
To see the relation between Assumption 3.2 and this definition, consider Gradient Descent (GD) step of the output neuron fe, k e [C] with step size performed on a single input =5 from class * c. By the chain rule, block-structured e implies locally-elastic predictions since
(((z)) 0f(25)
(x)) fz(x)) (10(x,.00 a5)
+ O(7"),
(6)
ie., the magnitude of the GD step of f(xf) is determined by the value of Oxx 25 Similarly, block-structured kernel 0 implies locally-elastic penultimate layer features because
Ol(x)) + 0 0f(rg)
++ (xf) - h (2()) = -100
This observation provides connection between our work and recent contributions suggesting connection between NC and local elasticity [3 53).
We provide figures illustrating the NTK block structure on other architectures and datasets in Appendix 'We discuss possible relaxations to our main assumption, where the "non -oa-diaaanal' components of the last-layer kernel O are allowed be non-zero, in Appendix

---

Eigenvalue
Eigenvector
Multiplicity
Asingle = 7d 7
vg m-I I m 1,
-
C
i>0, class CCO
Actass Asingle m(7e 7nd
Ve
((C 1)1
N-
C-1
class cs
others <0
Aglobal = Actass + Nrn
Vo 1,
Table I: Figendecomposition of the block-structured NTK.
Dynamics of DNNs with NTK Alignment
4.1 Convergence
As a warm up for our main results, we analyze the effects of the NTK block structure on the convergence of DNNs. Consider GD update of output neuron SK, k [C] with the step size ry:
f f((X) (X)(((() Yx) +07"", ==1....,..
(8)
Note that have taken into account that θ is zero for &t * by assumption. Denote the residuals corresponding to fa as r f (x) Y E R. Then we have the following dynamics for the residuals vector:
++ -1- (1-
+(1)²)
The eigendecomposition of the block clock-structured kernel *x, (X) provides important insights into this dynamics and is summarized Table We notice that the NTK has three distinct eigenvalues Aglohal 2 Actnss 2 Asinge.. which imply different convergence rates for certain components of the error. Moreover, the eigenvectors associated with each of these eigenvalues reveal the meaning the error components corresponding to each convergence rate. Indeed, consider the projected dynamics with respect to cigenvector Vo and eigenvalue Aglobal from Table [
(rk vo) = (1 (1-11)globann 7/Aaloaal ( (Tk. vo).
(10)
where we omitted (17²) for clarity. Now notice that the projection of r onto the vector Vo is in fact proportional to the average residual over the training set:
(rl vo) (rl. IN) N(rl))
(11)
where (.) denotes the average over all the training samples x; E X. By similar calculation, for all E [C] and i [m] we get interpretations of the remaining projections of the residual:
(rl.ve) = C - ((rk)e - (rl)), (rk.
m - (ri(2)) (rk)))
m
(12)
We where (-)c denotes the average over samples xf from class c, and r (xf) is the kth component of f (zf) - vf. Combining (I0). II and (2). we have the following convergence rates:
(ri+)) ) = (1 - 1) 1) (T(e)
(13)
(r(t)) (ri+2) = (1 ryAes ) ((rk) -rit))
(14)
(xf) (rk (1 )) (1 - m>sinnn (IY (2 (r())).
(15)
Overall, this means that the global mean (r) of the residual converges first, then the class means, and finally the residual of each sample r r(23). To simplify the notation, we define the following quantities:
R = f(X) Y = [r(æ1). r(NN)),
(16)
1
R
RY Ty (r),.- (r) I
(17)
m
1 (r) & 1.
R Relobal
R1r1N
(18)
N
5

---

where R E RCXN is the matrix of residuals. Relass E RCXN are the residuals averaged over each class and stacked m times, and Rglsbal E ROXN are the residuals averaged over the whole training set stacked N times. According to the previous discussion, Ryiahal converges to zero at the fastest rate, while R. converges the slowest rate. The last phase, which we call the end of training, when Rlass and R giobal have nearly vanished and can be treated as zero for the remaining training time. We will use this notion in several remarks, as well as in the proof of Theorem 5.2
4.2 Gradient Flow Dynamics with Block-Structured NTK
We derive the dynamics of H, W, b under Assumption 3.2 in Theorem 4.1 One can see that the
kernel greatly simplifies the complicated dynamics of DNNs and highlights the role
of each of the residual components identified in Section 4.1 We consider gradient flow, which is
close to gradient descent for sufficiently small step size [6]. reduce the complications caused by
higher order terms. The proof is given in Appendix B.I
Theorem 4.1. Suppose Assumption
3.2 holds. Then the gradient flow dynamics of DNN can be
written as
H- -WT - Ke) Re)R + (1e
Kyg MeRaasss + Kn N.g.oball
W -RH
(19)
b= -Rglobal
We note that the end of training, where Retass and Rggloaal are zero, the system
19 reduces to
H (Ma Re VHĆ.
W = W--Vw
((W. == 2 WWH + b1 - Y (20)
and b = O. This system differs from the unconstrained features dynamics only by a factor of KJ Re before H. Moreover, such form of the loss function also appears in the literature of implicit regularization 4 where the authors show that WH converges to low rank matrix.
5 NTK Alignment Drives Neural Collapse
The main goal of this work is to demonstrate how NC results from the NTK block structure. To this end, in Section 5.1 we further analyze the dynamics presented in Theorem 4. in Section 5.2 we derive the invariant of this training dynamics, and in Section 5.3 we finally derive NC.
5.1 Features Decomposition
We first decompose the features dynamics presented in Theorem 4.1 into two parts: H1, which lies in the subspace of the labels Y, and H2, which is orthogonal the labels and eventually vanishes. To achieve this, note that the SVD of Y has the following form:
YQ YYQ [Vmlc, C], P
(21)
where O RCx C) is matrix of zeros, and RCxC and Q R N are orthogonal matrices. Moreover, we can choose P and Q such that P = Ic and
Q QQ1, Q2),
Q
Lc @ 1m Vm
E E NXC
Q2 = Ic & Q₂ E NX(N-C) o
(22)
where @ is the Kronecker product. Note that by orthogonality, Q2 E mm has full rank and 1 m Q2 = O. We can now decompose HQ into two components as follows:
HQ V H Hz),
Hi
HQ
H₂
HQ
(23)
mm
mm
The following equations reveal the meaning of these two components:
H [(h) (h)...... (h)c). H2
H ( (QQ
(24)
where (h), R is the mean of h over inputs of from class [C], and H()) Raxm is the
submatrix of H corresponding to samples of class c, i... H (1) ..... H HO)) We see that H,

---

is simply the matrix of the last-layer features' class means, which is prominent in the NC literature. We also see that the columns of H Q2 are m different linear combinations of m vectors h( i E [m]. Moreover, the coefficients of each of these linear combinations sum zero by the choice of Q2. Therefore, H2 must reduce to zero in case of variability collapse (NC1), when all the feature vectors within the same class become equal. We prove that H2 indeed vanishes in DNNs with block-structured NTK as part of our main result (Theorem 5.2h.
5.2 Invariant
We now the former decomposition of the last-layer features to further simplify the dynamics and deduce a training invariant Theorem 5.1 The proof is given in Appendix [.2]
Theorem 5.1. Suppose Assumption 3.2 holds. Define H, and H2 as (23). Then the class-means of the residuals (defined in (17) are given by R1 WH1 + bl llc, and the training dynamics of the DNN can be written as
H1 --W R1 B1 [Petas + k.mlcl
H₂ MMingle W WHz
(25)
W =-m(R,HL HJ + WH,H,
b = --mR, Ic
where Msingl == Rid and Melass Msingle + m(Ke- Kn) are the two smallest eigenvalues of the kernel O K (X) (X).or for any k [n). Moreover, the quantity
E
WWW-
H₂H
(26)
m
Helass
Msingle
invariant in time. Here a
Nn
+Ckam
We note that the invariant E derived here resembles the conservation laws of hyperbolic dynamics that take the form Enyp == a² 122 consl for time-dependent quantities and b. Such dynamics arise when gradient flow is applied Since the solutions such minimization problems, given by ab = exhibit symmetry under scaling ya, b -> b/7,, the value of the invariant Enyp uniquely specifies the hyperbola followed by the solution. In machine learning theory, hyperbolic dynamics arise as the gradient flow dynamics of matrix factorization problems 3 15). Moreover, the end of training dynamics
loss function of the form L(a, b) (ab q)² for some q-
linear DNNs [42),
defined in (20 has a hyperbolic invariant given by
E
W
WWW-
HH
(27)
Asingle
Therefore, the final phase of training exhibits a typical behavior for the hyperbolic dynamics, which also characteristic for the unconstrained features models [21 Namely, "scaling" W and H by an invertible matrix does not affect the loss value but changes the dynamic's invariant. On the other hand, minimizing the invariant Egor has the same effect joint regularization of W and H (48).
However, we also note that our invariant provides new, more comprehensive look at the DNNs' dynamics. While unconstrained features models effectively make assumptions on the end-of-training invariant Eea to derive NC [21 38 48), our dynamics control the value of Ega through the more general invariant E. This way we connect the properties of und-of-training hyperbolic dynamics with the previous stages of training.
5.3 Neural Collapse
We are finally ready to state and prove our main result in Theorem 5.2 about the emergence of NC in
DNNs with NTK alignment. We include the proof in Appendix B.3
B.3
Theorem 5.2. Assume that the NTK has a block structure as defined Assumption 3.2 Then the
DNN's training dynamics are given by the system of equations in 25 Assume further that the
last-lau layer features are centralized, ie (h) = O, and the dynamics invariant (26) zero, i.e., E O.
Then the DNN's dynamics exhibit neural collapse defined in (NC4).
7

---

Below we provide several important remarks and discuss the implications of this result:
(1) Zero invariant assumption: We assume that the invariant (26) is zero in Theorem 5.2 for sim- plicity and consistency with the literature. Indeed, similar assumptions arise in matrix decomposition papers, where invariant guarantees "balance" of the problem [3]. However, our proofs in fact only require weaker assumption that the invariant terms containing features H are aligned with the weights W T i.i.
W
X H
H2H
(28)
Melass
Msingle
where we have taken into account our assumption on the zero global mean (h) = 0.
(2) Necessity of the invariant assumption: The relaxed assumption the invariant (28) is necessary for the emergence of NC in DNNs with block-structured NTK. Indeed, NCI implics H = O, and NC3 implies H x W W. Therefore, DNNs that do not satisfy this assumption do not display NC. Our numerical experiments deseribed in Section 6 strongly support this insight (see Figure 2 panes a-e). Thus, we believe that the invariant derived in this work characterizes the difference between models that do and do not exhibit NC.
(3) Zero global mean assumption: We note that the zero global mean assumption (h) 0 in
Theorem 5.2 ensures that the biases are equal to b = - lc at the end of training. This assumption is common in the NC literature [2 (88 and is well-supported by our numerical experiments (see figures in Appendix [ pane i). Indeed, modern DNNs typically include certain normalization (c.g. through batch normalization layers) to improve numerical stability, and closeness of the global mean to zero is a by- ey-rruduct of such normalization.
(4) General biases case: Discarding the zero global mean assumption allows the biases b to take an arbitrary form. In this general case, the following holds for the matrix of weights:
(WWT)²2
(Ic-ales alclo (Ic-alcto-
(C)) big lcb"))
(29)
Helass
For optimal biases = C lc. this reduces to the ETF structure that emerges NC. Moreover, biases are all equal, i.e. b = B1c for some 3 E R, the centralized class means still form an ETF (i.e., NC2 holds), and the weights exhibit certain symmetric structure given by
WWT (1o-r1013)
M (Io (lc-d1013)
(30)
where y y (1-1111 -1-1111- aC) The proof and a discussion of this result are given in Appendix B.4 In general, the angles of these two frames are different, and thus NC3 does not hold. This insight leads us to believe that normalization is an important factor in the emergence of NC.
Our proofs and the discussion suggest that all the four phenomena that form NC do not have to always coincide. In particular, our proof of NC1 only requires the block-structured NTK and the invariant be P.S.D, which is much weaker than the total set of assumptions in Theorem 5.2 Therefore, variability collapse can occur in models that do not exhibit the ETF structure of the or the duality of the weights and the class means. Moreover, shown above, NC2 can occur when NC3 does not, i.e., the ETF structure of the class means does not imply duality.
(5) Partial NC:
6 Experiments
We conducted large-scale numerical experiments support our theory. While we only showease our results on single dataset- architeccr pair in the main text (see Figure 2 and refer the rest to the appendix, the following discussion covers all our experiments.
Following the seminal NC paper 39 we use three canonical DNN architec- tures: VGG [46.. ResNet (24) and DenseNet (26). Our datasets are MNIST [35), FaSHionMNIST (51) and CIFARIO 10 34.. We choose VGG11 for MNIST and FashionMNIST, and VGG16 for CIFARIO. We add batch normalization after every layer in the VGG architecture, set dropout to zero and choose the dimensions of the two fully-connected layers on the top of the network as 512 and 256. We use ResNet20 architecture described in the original ResNet paper [24 and DenseNet40 with bottleneck layers, growth k 12, and zero dropout for all the datasets.
Datasets and models.
8

---

a) Invoriontt nocm
b) Inv. cligranent to WTW
cxmomemaa lrt.
That amntrary
1000
LeChun
Re 89% (97.3 13%)
3.8
LeCan numaal let, 12300337
Avc: 09 69% (e (9)) 0.93
LR
10
Anc: 99 n0% (o5,33 0.51
0.35
041
zxt (9h-31%))
10-3
Ant. 09.5155 (e) y1"", 0.88
H: Hemand LFw0000
1.00
Uniooe inis.
1.5
10.0
Ant. 99.04%, (ef,13")
He nend 18--002)
0.23
1.0
c)
00:00
1.50
1,50
0.88
1.95
1.25
41.15
1.00
1.00
1,00
bJ) [[c bocetag Inir
1.5
41.00
1.75
0.75
1.0
0.39
0.50
0.50
32
00.05
0.5
0.25
0.98
100
2001
2000
000
0.00
114)
2010
2200
410
1501
2000
2010
4044
107 10"
opoch
spoch
opcch
Leaming rate
Figure 2: ResNet20 trained on MNIST with three initialization settings and varying learning rates (see Section 6 for details). We chose model that exhibits NC (red lines, filled markers) and a model that does not exhibit NC (blue lines, empty markers) for each initialization. The vertical lines indicate the epoch when the training accuracy reaches 99.9% (over the last 10 batches). a) Frobenious norm of the invariant |E][r_ b) Alignment of the invariant terms as defined in (28). c) NC1: standard deviation of h(()) averaged over classes. d) NC2: ||M M/ ||MM| M|] & F. where & is an ETF, e) NC3: |[W || - [| The legend displays the test accuracy achieved by each model and the last-layer features kernel alignment given by (@ /||@|||.. Y Y Y F)) The curves in panes a-e smoothed by Savitzky- filter with polynomial degree I over window of
size 10. Panes f, g and h show the NC metrics and the test accuracy as functions of the learning rate.
Optimization and initialization. We use SGD with Nesterov momentum 0.9 and weight decay x 10-4. Every model is trained for 400 epochs with batches of size 120. To be consistent with the theory, we balance the batches exactly. We train every model with set of initial learning rates spaced logarithmically
the range 77 E [10-4, 100.25). The learning rate divided by 10 every 120 epochs. On top of the varying learning rates, try three different initialization settings for every model: (a) LeCun normal initialization (default in Flax), (b) uniform initialization on [-vk, vk), where & = 1/n for a linear layer, and k = 1/(Kne-1) for a convolutional layer, where K the convolutional kernel size (default in PyTorch), (c) He normal initialization in fan_ mode.
Results. Our experiments confirm the validity of our assumptions and the emergence of NC as their result. Specifically, we make the following observations:
While most of the DNNs that achieve high test performance exhibit NC, able to identify DNNs with comparable performance that do not exhibit NC (see Figure [ panes f-h). We note that such models still achieve near-zero error on the training set in our setup.
Comparing DNNs that do and do not exhibit NC, we find that our assumption the invariant (see Theorem 5.2 and (28) holds only for the models with NC (see Figure panes a-e). This confirms our reasoning about the necessity of the invariant assumption for NC emergence.
The kernels θ and O are strongly aligned with the labels Y Y Y in the models with the best performance, which is in agreement with the NTK alignment literature and justifies our assumption on the NTK block structure.
We include the full range of experiments along with the implementation details and the discussion of required computational resources in Appendix C Specifically, we present figure analogous to Figure 2 for every considered dataset-architeccure pair. Additionally, we report the norms of matrices H H.H and (h) (h) well as the alignment of both the NTK θ and the last-layer features kernel h in the end of training, to further justify our assumptions.

---

7 Conclusions and Broad Impact
This work establishes the connection between NTK alignment and NC, and thus provides mechanis- tic explanation for the emergence of NC within realistic DNNs training dynamics. It also contributes the underexplored line of research connecting NC and local elasticity of DNNs' training dynamics.
The primary implication of this research is that it exposes the potential to study NC through the lens of NTK alignment. Indeed, previous works on NC focus the top-down approach (layer-peeled models) [88 21
38 48) and fundamentally cannot explain how NC develops through earlier layers of DNN and what are the effects of depth. On the other hand, NTK alignment literature focuses on the alignment of individual layers [7]. and recent theoretical results even quantify the role of each hidden layer in the final alignment [37]. Therefore, we believe that the connection between NTK alignment and NC established in this work provides a conceptually new method to study NC.
Moreover, this work introduces novel approach facilitate theoretical analysis of DNNs training dynamics. While most theoretical works consider the NTK in the infinite- limit to simplify the dynamics [1] 20 28, 49), our analysis shows that making reasonable assumptions the empirical NTK can also lead to tractable dynamics equations and new theoretical results. Thus, we believe that the analysis of DNNs' training dynamics based the properties of the empirical NTK promising approach also beyond NC research.
8 Limitations and Future Work
The main limitation of this work is the simplifying Assumption 3.2 on the kernel structure. While the NTK of well-trained DNNs indeed has approximate block structure (as we discuss in detail in Section B). the NTK values also tend to display high variance in real DNNs 22 44 Thus, we believe that adding stochasticity to the dynamics considered in this paper is a promising direction for the future work. Moreover, the empirical NTK exhibits so-called specialization, i.e., the kernel matrix corresponding to certain output neurons aligns more with the labels of the corresponding class (5). btock-structured kernels, specialization implies different values blocks corresponding to different classes. Thus, generalizing our theory to block-structured kernels with specialization another promising short-term research goal. In addition, our theory relies on the assumption that the dataset (or the training batch) is balanced, i.c., all the classes have the same number of samples. Accounting for the effects of non-balanced datasets within the dynamics of DNNs with block-structured NTK is another possible future work direction.
More generally, we believe that empirical observations are essential to demistify the DNNs' training dynamics, and there are still many unknown and interesting connections between seemingly unrelated empirical phenomena. Establishing new theoretical connections between such phenomena is an important objective, since it provides more coherent picture of the deep learning theory as whole.
10

---

Acknowledgments and Disclosure of Funding
R. Giryes and G. Kutyniok acknowledge support from the LMU-TAU International Key Cooperation Tel Aviv University 2023. R. Giryes is also grateful for partial support by ERC-StG SPADE grant no. 757497. G. Kutyniok is grateful for partial support by the Konrad Zuse School of Excellence in Reliable AI (DAAD), the Munich Center for Machine Learning (BMBF) as well as the German Research Foundation under Grants DFG-SPP-2299 KU 1446/31-1 and KU 1446/32-1 and under Grant DFG-SFB/TR 109, Project C09 and the Federal Ministry of Education and Research under Grant MaGriDo.
References
[1] Ben Adlam and Jeffrey Pennington. The neural tangent kernel in high dimensions: Triple descent and multi-scale theory of generalization. In Proceedings of the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 74-84. PMLR, 2020.
[2] Laurence Aitchison. Why bigger is not always better: on finite and infinite neural networks, In Proceedings of the 37th International Conference on Machine Learning, ICML. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 156-164 PMLR, 2020.
[3] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässaan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages
PMLR, 2018.
[4] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Advances Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 7411 1-7422, 2019.
[5] Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks kernel learners: The silent alignment effect. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[6] Bubacarr Bah, Holger Rauhut, Ulrich Terstiege, and Michael Westdickenberg. Learning deep linear neural networks: Riemannian gradient flows and convergence to global minimizers. Information and Inference: A Journal of the IMA, 11(1):307-353, 02 2021.
[7] Aristide Baratin, Thomas George, César Laurent, R. Devon Hjelm, Guillaume Lajoie, Pascal Vincent, and Simon Lacoste-Julien. Implicit regularization via neural feature alignment. The 24th International Conference on Artificial Intelligence and Statisties, AISTATS 2021, April 13-15, 2021, Virtual Event, volume 130 of Proceedings of Machine Learning Research, pages 2269-2277. PMLR, 2021.
[8] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Monderrman-Minng and Qiao Zhang. JAX: composable transformations of Python+NummPy programs, 2018.
[9] Shuxiao Chen, Hangfeng He, and Weijie J. Su. Label-aware neural tangent kernel: Toward better generalization and local elasticity. In Advances Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[10] Lénaïc Chizat, Edouard Oyallon, and Francis R. Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NewrIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 2933-2943, 2019.
[11] Hung-Hat -Hsu Chou, Carsten Gieshoff, Johannes Maly, and Holger Rauhut. Gradient descent for deep matrix factorization: Dynamics and implicit bias towards low rank. arXiv preprint: 201 1.172,, 2020.
[12) Moustapha Cissé, Piotr Bojanowski, Edouard Grave, Yann N. Dauphin, and Nicolas Usunier. Parseval networks: Improving robustness to adversarial examples. In Proceedings of the 34th
11

---

International Conference on Machine Learning, ICML 2017. Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 854-863. PMLR, 2017.
[13] Nello Cristianini, John Shawe-Taylor, André Elisseeff, and Jaz S. Kandola. On kernel-target alignment. In Advances in Neural Information Processing Systems 14: Natural and Synthetic, NIPS 2001, December 3-8, 2001, Vancouver, British Columbia, Canada, pages 367-373. MIT Press, 2001.
[14] Ahmet Demirkaya, Jiasi Chen, and Samet Oymak. Exploring the role of loss functions in multiclass classification. In 54th Anmual Conference on Information Sciences and Systems, CISS 2020, Princeton, NJ, USA, March 18-20, 2020, pages 1-5. IEEE. 2020.
[15] Simon S. Du, Wei Hu, and Jason D. Lee. Algorithmic regularization in learning deep ho- mogencous models: Layers are automatically balanced. In Advances in Neural Information Processing Systems 31: Annual Conference Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 382-393, 2018.
[16] Omer Elkabetz and Nadav Cohen. Continuous vs. discrete optiizzation of deep neural networks. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 4947-4960, 2021.
[17] Tolga Ergen and Mert Pilanci. Revealing the structure of deep neural networks via convex duality. In Proceedings of the 38th International Conference on Machine Leurning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 3004-3014, PMLR, 2021.
[18) C Fang, H He, Long, and WJ Su. Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences of the United States of America, 118(43), 2021.
[19] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[20] Mario Geiger, Arthur Jacot, Stefano Spigler, Franck Gabriel, Levent Sagun, Stéphane d' Ascoli, Giulio Biroli, Clément Hongler, and Matthieu Wyart. Scaling description of generalization with number of parameters in deep learning. Journal of Statistical Mechanics: Theory and Experiment, 2020(2):034001, 2020.
[21] X. Y. Han, Vardan Papyan, and David L. Donoho. Neural collapse under MSE loss: Proximity to and dynamics on the central path. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[22] Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[23] Hangfeng He and Weijie J. Su. The local elasticity of neural networks. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia. April 26-30, 2020. OpenReview.net, 2020.
[24] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recog- nition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 770-778. IEEE Computer Society, 2016
[25] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2020.
[26] Gao Huang, Zhuang Liu, Laurens van der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017, pages 2261 2261.226. IEEE Computer Society, 2017.
12

---

[27] Jiaoyang Huang and Horng-Tzer Tzer Yau. Dynamics of deep neural networks and neural tangent hierarchy. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 1542-45517 PMIR, 2020.
[28] Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks generalize better than deep feedforward networks? A neural tangent kernel perspective. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[29] Like Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross- entropy in classification tasks. In 9th International Conference on Learning Representations. ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[30] Arthur Jacot, Clément Hongler, and Franck Gabriel. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pages 8580-8589, 2018.
[31] Yiding Jiang, Dilip Krishnan, Hossein Mobahi, and Samy Bengio. Predicting the generalization gap in deep networks with margin distributions. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
[32] Dmitry Kopitkov and Vadim Indelman. Neural spectrum alignment: Empirical study. In Artificial Neural Networks and Machine Learning ICANN 2020- 29th International Conference on Artificial Neural Networks, Bratislava, Slovakia, September 15-18, 2020, Proceedings, Part II, volume 12397 of Lecture Notes in Computer Science, pages 168-179. Springer, 2020.
[33] Vignesh Kothapalli, Ebrahim Rasromani, and Vasudev Awatramani. Neural collapse: A review on modelling principles and generalization. arXiv preprint arXiv: 04044,, 2022.
[34] Alex Krizhevsky et al. Learning multiple layers of features from tiny images, 2009.
[35] Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann. www.:allww.n....ettttieeeee com/exdb/nmisst 2, 2010.
[36] Jachoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Advances in Neural Information Processing Systems, 1111556115722 2020.
[37] Yizhang Lou, Chris E Mingard, and Soufiane Hayou. Feature learning and signal propagation in deep neural networks. In International Conference on Machine Learning, pages 14248-14282. PMLR, 2022.
[38] Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. arXiv preprint arXiv:2011.11619, 2020.
[39] Vardan Papyan, XY Han, and David L. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40):24652-244 246633, 2020.
[40] Federico Pernici, Matteo Bruni, Claudio Baccchi, and Alberto Del Bimbo. Fix your features: Stationary and maximally discriminative embeddings using regular polytope (fixed classifier) networks. arXiv preprint arXiv: 1044, 2019.
[41] Tomaso A. Poggio and Qianli Liao. Explicit regularization and implicit bias in deep network classifiers trained with the square loss. arXiv preprint arXiv:2 10100077, 2021.
[42] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
[43] Mariaa Seleznova and Gitta Kutyniok. Analyzing finite neural networks: Can we trust neural tangent kernel theory? In Mathematical and Scientific Machine Learning, 16-19 August 2021, Virtual Conference / Lausanne, Switzellnnd, volume 145 of Proceedings of Machine Learning Research, pages 868-89: 895. PMLR, 2021.
[44] Mariaa Seleznova and Gitta Kutyniok. Neural tangent kernel beyond the infinite-width limit: Effects of depth and initialization. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 19522-19560 PMLR, 2022.
13

---

[45] Haozhe Shan and Blake Bordelon. A theory of neural tangent kernel alignment and its influence on training. arXiv preprint arXiv:2105.14301, 2021.
[46] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[47] Yifan Sun, Changmao Cheng. Yuhan Zhang, Chi Zhang, Liang Zheng, Zhongdao Wang, and Yichen Wei. Circle loss: A unified perspective of pair similarity optimization. Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 6398-649889 0478-6407470 2020.
[48] Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural collapse. Proceedings of the 39th International Conference on Machine Learning, volume 162, pages 21478 21550, 2022.
[49] Tom Tirer, Joan Bruna, and Raja Giryes. Kernel-based smoothness analysis of residual networks.
In Mathematical and Scientific Machine Learning, volume 145, pages 921-954, 2021.
[50] Tom Tirer, Haoxiang Huang, and Jonathan Niles-Weed. Perturbation analysis of neural collapse.
arXiv preprint arXiv:22 10.::2280.6668, 2022.
[51] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint
2017.
[52] Greg Yang. Tensor programs II: neural tangent kernel for any architecture. arXiv preprint arXiv:2006.. 14548, 2020.
[53] Jiayao Zhang, Hua Wang, and Weijie J. Su. Imitating deep learning dynamics via locally elastic stochastic differential equations. In Advances Neural Information Processing Systems 34: Anmual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 1021, virtual, pages 6392-6403, 2021.
14

---

A Related works
NC with MSE loss. NC was first introduced for DNNs with cross-entropy (CE) loss, which is commonly used in classification problems [39]. Since then, numerous papers discussed NC with MSE loss, which provides more opportunities for theoretical analysis. especially after the MSE loss was shown perform on par with CE loss for classification tasks [4 29
Most previous works on MSE-NC adopt the so-called unconstrained features model 22 38 48). In this model, the last-layer features H are free variables that are directly optimized during training, i.., the features do not depend on the input data or the DNN's trainable parameters. Fang et al. [8]] also introduced generalization of this approach called N N-layer-peeled model, where features of the N-th-to-last layer are free variables, and studied the 1 1-layer-peeled model (equivalent to the unconstrained features model) with CE loss as special case.
One line of research on MSE-NC in anconstrained/layr--eeelll mizers of optimization problems associated with DNNs [17 18 48). In particular, Tirer el al. [48 showed that global minimizers of the MSE loss with regularization of both H and W exhibit NC. Moreover, Ergen & Pilanci [17] showed that NC emerges in global minimizers of optimization model. In comparison to our work, these contributions do not consider the training dynamies of DNNs, i..., they do not discuss whether and how the model converges to the optimal solution.
models aims to derive global mini-
problems with general convex loss in the context of the
Another line of research on MSE-NC explicitly considers the dynamics of the unconstrained features models 21 38). In particular. Han et al. [21] considered the gradient flow of the unconstrained renormalized features along the "central path", where the classifier is assumed to take the form of the optimal least squares (OLS) solution for given features H. Under this assumption, they derive a closed- form dynamics that implies NC. While they empirically show that DNNs are close to the central path in certain scenarios, they do not provide a theoretical justification for this assumption. The dynamics considered in their work is also distinct from the standard gradient flow dynamics of DNNs considered in our work. On the other hand, an earlier work by Mixon et al. [38] considered the gradient flow dynamics of the unconstrained features model, which is equivalent (up to rescaling) to the end-of- training dynamics (0) that we discuss in Sections 4.2 and 5.2 Their work relies on the linearization of these dynamics to derive certain subspace, which appears to be an invariant subspace of the non non-linearized unconstrained features model dynamics. Then they show that minimizers of the loss from this subspace exhibit NC. We note that, in terms of our paper, assuming that the unconstrained features model dynamics follow a certain invariant subspace means making assumptions on the end-of-training invariant (27). In comparison to these works, make step towards realistic DNNs dynamics by considering the standard gradient flow of DNNs simplified by Assumption 3.2 on the NTK structure, which is supported by the extensive research on NTK alignment 77 9 44 45]. In our setting, the NTK captures the dependence of the features on the training data, which is missing in the unconstrained features model. Moreover, while other works focus only on the dynamics that converge to NC, we show that DNNs with MSE loss may not exhibit NC in certain settings, and the invariant of the dynamics 26) characterizes the difference between models that do and do not converge to NC.
Notably, works by Poggio& & Liao [4] adopt a model different from the unconstrained features model to analyze gradient flow of DNNs. They consider the dynamics of homogeneous DNNs, in particular ReI.U networks without biases, with normalization of the weights matrices and weights regularization. The goal of weights normalization in their model is to imitate the effects of batch normalization in DNNs training. In this model, certain fixed points of the gradient flow exhibit NC. While the approach taken in their work captures the dependence of the features on the data and the DNN's parameters, it fundamentally relies on the homogeneity of the DNN's output function. However, most DNNs that exhibit NC in practice are not homogeneous due biases and skip-connections.
NC and local elasticity. A recent extensive survey of NC literature 333 discussed local elasticity as a possible mechanism behind the emergence of NC, which has not been sufficiently explored up until now. One of the few works in this research direction is by Zhang et al. [53), who analyzed the so-called locally-elastic stochastic differential equations (SDEs) and showed the emergence of NC in their solutions. They model local elasticity of the dynamics through an effect matrix, which has only two distinct values: a larger intra-class value and a smaller inter-class value. These values characterize how much influence samples from one class have on samples from other classes in the SDEs. While the aim of their work is to imitate DNNs training dynamics through SDEs, the authors
15

---

do not provide any explicit connection between their dynamics and real gradient flow dynamics of DNNs. On the other hand, we derive our dynamics directly from the gradient flow equations and connect local elasticity to the NTK, which is well-studied object in the deep learning theory.
Another work by Tirer al. 150 provided perturbation analysis of NC study "inexact collapse". They considered a minimization problem with MSE loss, regularization of H and W, and additional regularization of the distance between H and given matrix of initial features. In the "near-collapse" setting, i.e., when the initial features are already close to collapse, they showed that the optimal features can be obtained from the initial features by certain linear transformation with block structure, where the intra-class effects are stronger than the inter-class ones. While this transformation effect matrices in locally-e elasstic training dynamics, it does not originate from the gradient flow dynamics of DNNs and is not related to the NTK.
matrix resembles the
B Proofs
B.1 Proof of Theorem 4.1
Proof of Theorem 4.7 We will first derive the dynamics of h1 (x which is the s-th component of the last-layer features vector sample xf X from class E [C]. Let w R be the trainable parameters of the network stretched into a single vector. Then its gradient flow dynamics is given by
C ME (f(X)
-w[(f
=
- Ykr)) V (X))x)
(31)
where V f(X)(x) ke R. P is the component of the DNN's Jacobian corresponding to output neuron and the input sample x.. Since entries of f(X) can written as
f(X)ss Σ
Wks
= Σ
Wkoth (x²) + bk.
(32)
we obtain
C n
w Σ
Σ
Σ (f(X)
V
WKerhs (x b&).
(33)
By chain rule, we have (xf)) (V (zf), w.) Then, taking into account that
(x Vw ( hW (x.)) bx)) Wks (V (wh,(x))) (xf). V (
(34)
and that (Wwh. Wwh (x$)) = 0 (x/,,ş) x by definition of Oh, we have
n
ha(x;) = E
(F(X) YaU W WE Oh (c),
(35)
=1==11
Now by Assumption 3.2 we have 0h 3,8'
= 0 if + s'. Therefore, the above expression simplifies
N
C (F(()) - Yau)
hx(z)) =
4 (zf
1=1
N
Σ
W (WH b1 Y)s (ag,
To express H [+.(a:)] s,t R in matrix form, it remains to express O (xf, xx') as the iith entry of some matrix. We will separate the sum into three cases: 1) = i', 2) ب i' and ==d, d, and 3) c # c'. According to Assumption 3.2 the first case corresponds to the multiple of identity FallN. The second corresponds to the block matrix of size with zeros on the diagonal,
16

---

which can be written as Ki o Y [ln.. The third matrix equals to Kn (,,11 - Y). Therefore we can express the dynamics of H as follows:
H [WT (WH b Y) - kY - (IN1N Y YY)
(Nd Ni. W (WH b1 -Y)
(Ni Kn. Kn,) W (WHY") Y + +mbli -mY)
N Kn W (WHIx1N Nb1 lclx).
Now we notice that HY Y Y/m is the matrix of stacked class means repeated m times each and HIy N/N is matrix of the global mean repeated times. Therefore, we have
WHY Y mbl N Y Relssss
WHINIY + +Nb1 clccN = N NRRglaball
according to the definitions of global and class- mean residuals in (18) and (17).
The expressions for the gradient flow dynamics of W and follow directly from the derivatives of f(X) w.r.t. W and b. This completes the proof.
Li
B.2 Proof of Theorem 5.1
Proof of Theorem 5.1 Recall from
23 in Section 5.1 that have the following decomposition
HQ=
m H H.l.
H1
HQ
H2
HQ₂
V m
with orthogonal
Q1 Q2] E R NxN
We now artificially add QQ (= lx) to the dynamics
I9
in Theorem
and obtain
HQ
-Iad Ke W (WWHQQ + Q YQ)
-(Ke - Kn) m. WT WHQQQ Y
T b1 QQ
NK WHQQ 1 QQ + b1 -
(36)
W -(WHQ + b YQQQ
b -WWHQ + b1 YQ)Q
Let us simplify the expression. Since Q Voo c & 1m and Q2 = Ic 2 Q2, we have
1 1 = mm 0), YQ = /m
(37)
Plugging
37 into
36
the dynamics can be decomposed into
H = - (K. Ke) T (WH, ( + b1 Ic)
- -(Ke Kn, mm.. (WH, + bb] - Tc)
N WH, b1g (
(38)
H2 -(Ka Ki) W WH,
W -m(WH, + bl - H - - WH,
-m(WH, + bl Ic)le-
To further simplify (38). define the following quantities
Msingle = Ki
Figs
Helass
Msingle + m(K. k,),, R = WH, +b -Ic
(39)
Notice that Hsingss and Melass are the two largest eigenvalues of the block-structured kernel 0 (X) (see Table [ for the eigndccomposition of a strucsstret matrix), and R, is a matrix of the stacked class-mean residuals, which is also defined in (7). The the dynamics (38) simplifies to
H W (Heasss R2 + N WH,1c11 + b1T - (lc1))
Hg -single W WH,
(40)
W H, WH,H,
b -mR,1c.
17

---

It remains to simplify the expression for H1. By using the relation
WHilch + big c -
(41)
can deduce that the dynamics for H in (40) can be expressed (recalling that N mC)
H W R (Actass kmmllll)).
(42)
We notice that (Ic the invariant of the training dynamics by direct computation of the time-derivative E, where
1cl0) Lc alc I. where a fictect Ckam
Then can derive
Reius
E
W W
Hi (Ic - alc)))
H₂H,
(43)
Melass
Msingle
Since E O we get that the quantity remains constant in time. This completes the proof.
B.3 Proof of Theorem 5.2
We divide the proof into two main parts: the first shows the emergence of NC1, and the second one shows NC2-4.
(NCI). Following the analysis in Section 3 the dynamics eventually enters the end of training phase (see Section Then the dynamics in Theorem [5. simpliffie to the following form:
H1 = 0
H2 Msingle WWHH,
(44)
W -mWH,HH
0
As we note in Section 4 this dynamics is similar to the gradient flow of the unconstrained features models and an instance of the class of hyperbolic dynamics, which is discussed in Section [.2 During this phase the quantity
H mmMsingle E 1 (1
E = Msingle W W
(45)
Melass
does not change in time. Hence we can decouple the dynamic using the invariant as follows:
H2 Psingle (E + m mHgH, H2 w W (paingle W w -)E)
(46)
Since E is p.s.d (or zero, special case), E p.s.d well, and the eigendecomposition of the invariant is given by E E. C&UA. for some coefficients c 2 and set of orthonormal vectors en R Then also have H2H, Σx. OuLUENy where ar are symmetric (i.e. Or on)) and a&A 2 for all & = 1.... n (since H,H,J is symmetric and p.s.d.). Note that coefficients ck here are constant while coefficients au are time-dependent. Let then write the dynamics for ag using the dynamics of H,H,J
(HH = HH,H, H2H, É 2(H,H,))
(47)
Then for the elements of a have:
âR -aw(c + ci) -2 -2 _ axjaj
(48)
For the diagonal elements asks, this gives:
âgk -2c000k - Σ 2
(49)
18

---

Since CR 20,0x agt 2 0 and ag
2 0, we get that
Okk
4/
(50)
t-00
And, therefore, all the non-diagonal elements also tend to zero. Thus, we get that
HgH2
O
(51)
11000
and thus
H !
(52)
Now we notice that from the expression for H2 in (24) follows that H₂ = implies variability collapse, since means that all the feature vectors within the same class are equal. Indeed, H Q2 O E (m means that there is a set of 1 orthogonal vectors, which are all also orthogonal h2
to [=.))))..... ,,, (zfr,) only vector (up to constant) orthogonal to all the columns of Q2 in 1m and this vector is 1m- Therefore, [/,(x9),... (xm)) = 71m for some constant y for any = I....,n.
for any i = 1..... n, where zy are inputs from class c. However, there is
n. Thus, we indeed h(( which constitutes variability collapse within classes.
have hx(x)
(NC2-4). Set 3 = L. We first show that zero global feature mean implies b 33c. At the end of training. since R1 O, we have
C.
WH + b1 = Ic On the other hand, zero global mean implies H, (h) the right, we get the desired expression for the biases. Given the zero global mean, we have
(53)
O. Then multiplying (3 by 1
amC²
W TW
H
H E
E
(54)
m
MAlaass
Msingle
Pelass
By the proof of NC1, H2 + O. Together with the assumption that E is proportional to the limit of
W T W (or zero, as a special case), we obtain
Melass W T W MHHH - + W
(55)
for some 7 > O. Note that since H, H is p.s.d. this implies de = Melass 7 2 0. By multiplying the left and right with appropriate factors, we have
H WT w H H + 0 W W W mH H W - 0
(56)
Consequently (according to 53
AaC 21 HH H 0 3. (WW (Tc 311 0
(57)
Since both WWT and H. H1 are p.s.d., we have
H
V m m
(IG (1c
(58)
WWT
(Ic (1cl)
V T
To establish NC2, recall that H1 = _.......... [b)c] satisfies
and that M, as a normalized version of H,,
1
C (Tc -
1 lctT)
M
(Tc (1c1
1- 8
To establish NC3, note that from (55) and 588 together, it follows that the limits of M and WT only differ by a constant multiplier.
To establish NC4, note that using NC3
write argmax (Wh(z) + b). = argmax (Wh(z)).
= (b=31c)
argmax (M h(x)))
(NC3)
e
argmin [b(x) - (h)ell2.
This completes the proof.
19

---

B.4
General biases case
Proof. As in the proof of Theorem 5.2, at the end of training we have WH, + +b1 Ic. Moreover,
since E = O and H2 - O, we have
1 W m
(I + 0.
(59)
Melass
Multyplying the above expression to the left by W and to the right by WT, we obtain the general expression 29) for the matrix (WWT)² mentioned in the main text:
(WW T22
m
- [lc-alcld + + (l)-aCC(CCbb -b -1c-1ccb)) lobT)
(60)
Mellass
This expression implies that the rows of the weights matrix may have varying separation angles in the general biases case, i.e., there is symmetric structure is general. However, for constant biases b3,31c, B, the above expression simplifies to
=
(WW T2
m
10--0-0--000--003110
(61)
Helass
Since a 1/C and (1 (1-30)² 2 have that (1 (1-11-00(1-800)) 1/C. Therefore, the RHS of 6T) is always p.s.d. and has unique p.s.d square root proportional to Ic -1cc1 for some
constant < 1/C. Denote P (1- (1-(1=00)( - 3C)²) then we have -1-vI-CP)) Note that p < 1/C ensures that y is well defined. Then the configuration of the final weights given
by
WW
m
(Tc-71o13)
(62)
Melass
This means that the norms of all the weights rows are still equal, as in NC2. However, since < 1/C if / * 1/C, the angle between these rows is smaller than in the ETF structure.
We can derive the configuration of the class means similarly by multyplying (59) to the left by H and to the right by H1. In the general biases case, we get
H H. (Ic alcl0) H H1
Melass (Io- b1 lob + +1bb11117)
(63)
As with the weights, we see that this not symmetric structure in general. Thus, NC2 does not hold in the general biases case. However, for the constant biases b = Blc, the above expression simplifies to
H (To H H + Melass (Ic {1
-
(64)
Analogously to the previous derivations, we get that the unique p.s.d. square root of the RHS is given by To pl where p= - |1 BC < 1/C for 3 * 1/C. On the other hand, the unique p.s.d root of alold is given by llc plcld, where Φ = (1 vI aC)/C. Thus, we have the following
m
H H. bl.l) - Ic - plclo
(65)
Melass
Therefore, the structure of the last-layer features class means is given by
H H + Vi Melass m
$ [Ic [o-plc16) (TC (To 00 - 1016) 1 1 Molass (io-81ct))
(66)
where p + 6/(1 + oC) Cpp +C) 1/C for B # 1/C. Thus, similarly to the classifier weights W, the last-layer features class means form a symmetric structure with equal lengths and a separation angle smaller than in the ETF. However, the centralized class means given by M = H2 l/C) still form the ETF structure:
/ Melass (Lc-)1016)
M M
(67)
This holds since the component proportional to 1 on the RHS of equation (66) lies the kernel of the ETF matrix (Tc ICC). Thus, we conclude that NC2 holds in case equal biases, while NC3 does not.
D
20

---

Remark on - 0 case: Simplifying the expressions for constants y and 6, which define the angles in the configurations of the weights and the class means above, we get the following:
1 7 = (1- -1
1 (1 2 2- 2-v1-ac)
1 - SC1
0
(68)
Analyzing these expressions, we find that they are equal only if 1 - aC' i.c. a == 0. However, this can only hold if Kn = 0 by definition of a, i.e., when the kernel Q is zero on pairs of samples from different classes. While & 0 in general, there are certain settings where a approaches zero. Simplifying the expression for a, we can get the following
1 Nd Kip m7
(69)
+ (C-
One can see that a - 0 C + oc or when Ke/Ka Kn + cc. Since the kernel 6 is strongly aligned with the labels in our numerical experiments, the value of Ke/Kn, Ki, is large in practice. Thus, is not zero but indeed significantly smaller than 1/C. Thus, in our numerical experiments the angles 0 and Y are close to each other. However, we note that the equality of these two angles does not imply NC3, since the value of A characterizes the angles between the non-centralized class means.
Remark ar + 1/C case: If = 1/C, the equation (63) for the structure of the features class means with general (not equal) biases simplifies to
Melass (Ic (Tc- (10-)1018) MM M
(70)
m
iee, in this case the class means always exhibit the ETF structure, even without the assumption that all the biases are equal. Moreover, in this case 7 = 1/C as well. Thus, both NC2 and NC3 hold. While by definition o < 1/CC, we can analyze the cases when it approaches 1/C using the expression (69) again. One can see that when m - 0o and Ke/Kin - 1, we have a - 1/C. However, the requirement Kg/ Kin + implies that the kernel O does not distinguish between pairs of samples from the same class and from different classes. Such property of the kernel is associated with poor generalization performance and does not occur in our numerical experiments.
C Numerical experiments
Implementation details We use JAX (8 and Flax (neural network library for JAX) (25] to im- plement all the DNN architectures and the training routines. This choice of the software allows to compute the empirical NTK of any DNN architecture effortlestly and efficiently. We compute the val- ues of kernels θ and @ on the whole training batch (m = 12 samples per class, 120 samples in total) in case of ResNet20 and DenseNet40 approximate the values (id. Ter Yn) and (Ka, Ke: Kn). as well as the NTK aligmment metrics, and compute the invariant E using these values. Since VGG11 and much larger (over 10 million parameters) and computing their Jacobians is
VGG16 architectures very memory-intensive, we use m = samples per class (i.e., 40 samples in total) to approximate the kemels of these models. We compute all the other training metries displayed in panes a-e of Figures 8 9 [0 I1 on the whole last batch of every second training epoch for all the architectures. The test accuracy is computed on the whole test set. To produce panes f-h of the same figures, we only compute the NC metrics and the test accuracy one time after 400 epochs of training for every learning rate. We use 30 logarithmically spaced learning rates in the range 1) E [10-4 100:25 for ResNet20 trained on MNIST and VGG11 1 trained on MNIST. For all the other trchitecture-dataee pairs we only compute the last 20 of these learning rates to reduce the computational costs, since the smallest learning rates do not yield models with acceptable performance.
6
Compute We executed the numerical experiments mainly on NVIDIA GeForce RTX 3090 Ti GPUs, each model was trained on a single GPU. In this setup, single training run displayed in panes a-e of Figures 3 5 8 9 IO II took approximately 3 hours for ResNet20, 6 hours for DenseNet40. 7 hours for GGGI,, and IT hours for VGG16. This adds up to total of 312 hours to compute panes a-e of the figures. The computation time is mostly dedicated not to the training routine itself but to the large number of computationally-hent
metrics, which are computed every second epoch of a training
run. Indeed, to approximate the values of ② and O", one needs to compute C(C + 1) + +z((n + 1) kernels a sample of size mC from the dataset, and each of the kernels requires computing
21

---

gradient with respect to numerous parameters of DNN. Additionally, the graphs in panes f-h of the same figures take around 1.5 hours for each learning rate value for ResNet20, 3 hours for DenseNet40, and 4 hours for VGG11 and VGGI6, which adds up to approximately 1350 computational hours.
Results
We include experiments on the following architecture-dataset pairs:
Figure
3
VGGI1 | trained on MNIST
Figure
VGG11 trained on FashionMNIST
Figure 5
VGG16 trained on CIFAR10 10
Figure
ResNer20 trained on MNIST
Figure
ResNet20 trained on FashionMNIST
Figure 8
ResNet20 trained CIFARIO
Figure
DenseNet40 trained on MNIST
Figure
10
DenseNer40 trained on FashionMNIST
Figure
DenseNet40 trained on CIFAR1O
The experiments setup is described in Section 6 Panes a-h of 3
10, 11
analogous to the same panes of Figure 2 We include additional pane here which displays the norms of the invariant terms corresponding the feature matrix components H, and H2, and the global features mean (h) at the end of training. One can see that the global features mean is relatively small in comparison with the class-means in every setup, and the "variance" term H2 is small for models that exhibit NC. We also add pane j. which displays the alignment of kernels θ and @ for every model the end of training. One can see that the kernel alignments is typically stronger in models that exhibit NC.
C.1 Additional examples of the NTK block structure
We include the following additional illustrative figures (analogous Figure in the main text) that show the NTK block structure in dataset-architecture pairs covered in our experiments:
Figure 13
VGG11 trained on MNIST
Figure
VGG11 trained on FaSHionMNIST
Figure 15
VGG16 trained on CIFAR1O
Figure 16
ResNet20 trained on FashionMNIST
Figure 17
ResNet20 trained on CIFARIO
Figure 18
DenseNet40 trained on MNIST
Figure 19
DenseNer40 trained on FashionMNIST
Figure
DenseNer40 trained on CIFAR10
Overall, the block structure pattern visible in the traced kernels in all the figures. As expected, the block structure is more pronounced in the kemels where the final alignment values are higher. While the norms of the "non-diagonal" components of the kernels are generally smaller than the "diagonal" components in panes c) and d), we notice that there is a large variability in the norms of the "diagonal" components some settings. This means that different neurons of the penultimate layer and different classification heads may contribute to the kernel unequally in some settings. Moreover, certain "non-diagonal" components the last-layer kernel may have non-negligible effect some settings. We discuss how one could generalize our analysis to account for these properties of the NTK in Appendix(D D
C.2 Preliminary experiments with CE loss
While CE loss is common choice for training DNN classifiers, our theoretical analysis and the experimental results only cover DNNs trained with MSE loss. For completeness, we provide experimental results for ResNet20 trained on MNIST with CE loss in Figure [2] One can see that
22

---

smaller invariant norm and higher invariant alignment correlate with NC in the figure. However, DNNs trained with CE loss overall reach better NC metries but have much larger norm of the invariant in comparison with DNNs trained with MSE loss.
D Relaxation of the NTK Block-Structure Assumption
In this section, we first derive the dynamics equations of DNNs with general block structure assumption on the last-layer kernel @ (analogous to the equations presented in Theorem (.1 and Theorem 5.1] Then we discuss a possible relaxation of Assumption 3.2 regarding NC in Theorem 5.2 still holds.
and
under which our main result
D.1 Dynamics under General Block Structure Assumption
We first formulate the most general form of the block structure assumption O follows: Assumption D.1. Assume that O X x X Raxn has the following block structure
tn (x, A. An
θ (a A. +A
4 (x (1) AAn
(71)
where Ad.c.n E Rax= are arbitrary p.s.d. matrices. Here x; and æz are two distinct inputs from the same class, and ry an input from class d' +
This assumption means that every kernel matrix values, corresponding to the inter-class, intra-class, and the diagonal values of the kernel. However, these values arbitrary and may depend on the choice of k, [L.
(X), k, ,
[1, still has at most three distinct
Under the general block structure assumption, the gradient flow dynamics of DNNs with MSE loss takes the following form:
H W R m... W Relas + NA W Rglobal
W
-RH
(72)
b b= -Rgloba
This the generalized version of the dynamics presented Theorem 4.1 decomposed dynamics presented Theorem 5.1 takes the following form under the general block structure assumption:
Consequently. the
H1 = (Aa + mA.)WTTR mA, W R,1cl
H2 W WH,
(73)
W H + WH,H,)
b =-mR,lc.
The derivation of the above dynamics equations are identical to the proofs of Theorem 4.1 and
Theorem 5. presented in Appendix
Rotation invariance We notice that the dynamics of (W, H) in (22) has to be rotation invariant, ie., the equations should not be affected by change of variables W -Q. Q for any orthogonal matrix Q. This holds since the loss function only depends the product WH, which does not change under rotation. This requirement puts conditions on the behavior Ad, under rotation. Indeed, assume that the rotation W WQ, H - following change of the kernel:
Q H for some Q corresponds to the
Ad,cn dd.., (Q),
(74)
then the rotation invariance of the dynamics implies the following equality for any Q:
QÂ)(Q)Q T
T R mQA.(QQQ T + NQA,, (Q) T W Rglobal
(75)
= W R + mA, T Relass + NA, W Relobal-
(76)
These equations are satisfied trivially with our initial assumption, where Adce, = A Q x I... However, as we can see, any generalized assumption should specify the behavior of the kernel under rotation, and satisfy the above equation.
23

---

For general Ad.c, the following behavior under rotation trivially satisfies the above condition: Ad, (Q) (Q)= Q Ad.c.n Q. This behaviour of the kernel under rotation is intuitive, since it implies that the gradients of the last-layer features h are rotated in the same way as the features. However, we note that gradients of parametrized functions do not in general behave this way, since the rotation of the function has to be realized by a certain change of parameters. Consider, for instance, a one-hidden-layer linear network with weights V in the first layer. Then we have H VX, and rotation H + Q H corresponds to the change of parameters V + Q V. In this case, the kernel does not change under rotation, i.e., Ad (Q) = Ad,cn
Dynamics invariant We note that the dynamics in 73 does not in general have an invariant analogous to the one we identified in Theorem 5.1 Indeed, if we define quantity E == W
H H H2 for some constants 61 IR, and additionally assume centered global means H c = 0, we get the following expression for the derivative of E:
= (cI(Ad + mA.) - mlln W H
H W (a((A + mAa)) - mmlln)
(77)
+ (c2AA mila) w WH,H,, H.H, W W W (caAJ (2A]]mmlm)) -mlln).
(78)
which is not equal to zero with arbitrary matrices Ad,c-
D.2 Neural Collapse under Relaxed Block Structure Assumption
We now propose relaxation of our main assumption. under which our main result regarding NC D.1 in Theorem 5.2 still holds. In terms of Assumption D.1 on the general block structure of O*, our initial Assumption 3.2 in the main text is the special case with A, KnIns A. = (Ke Kn I., Aa = (Ka I.. The relaxed assumption can be formulated as follows in terms of matrices Ad,
Assumption D.2. Assume that A, is an arbitrary p.s.d. matrix and (Ac, Aa) satisfy the following conditions:
A Kolla + No, Aa= 10=-klh Kaln + Na: where N ker(R W)), i.e. Ne.d W R 0. Further, assume that the kernel changes under rotation with orthogonal matrix as follows:
(79)
Ad,., ( (Q) = Q Q.
(80)
Since A, is arbitrary, this relaxation allows arbitrary non-zero values of non-diagonal kernels O with & + s. The following observations justify the consistency of the above assumption:
Since W E RCx, R E RCXN and N > > C, T has a non-empty kernel (possibly time- demeddenndd.
The dynamics is rotation invariant under the assumption, i.e., the equation 75 holds.
The expression of the assumption is rotation invariant, in a sense that A.,a Q) Kd,ella Na,e(Q) where NT. E ker(R. WQ) for any orthogonal Q.
Under the above assumption, the derivative in becomes zero, so the dynamics has an invariant of the form E = W w ci H, e₂ H,, Moreover, the statement and the proof of our main
Theorem 5.2 remains unchanged. Thus, DNNs satisfying the conditions of Theorem 5.2 display NC
under Assumption D.2
D.3 Discussion
The analysis of the DNNs dynamics is simplified significantly by assuming that Oh has a block structure. However, formulating reasonable and consistent assumption the NTK and its compo- nents is non-trivial. The Assumption 3.2 that we used in the main text is justified by the empirical results but may not capture all the relevant properties of the NTK. We believe that studying DNNs' dynamics under a more general or a more reasonable assumption on the NTK a promising future work direction. The relaxed block structure assumption proposed in this section is the first step into this direction.
24

---

a] Invaciintt noem
10
b) Inv. aliggnmeent to
1.0
LoCan (9".99)))
The accumry
1,00
1.
LaCt meme me.
Cadsen ian.. LBw1100
100
Jan 90.88%, (9))))")) 0.70
0.06
Caffenn mn. Li-1110
99.88, (99,9)9)))
10-1
0.
90.50% (99.))))))))))
1.00
(5) Unilogu init.
1.5
L.B-41332
(99,333
0.50
c) NCI
d) NC2
1.50
e) NC3
1.50
0.5
1.25
4015
1.00
1.00
10 Hecorcl init
0.10
0.75
4.5
0.85
0.50
3.50
4005
1.25
3.25
0.06
100
200)
400
0.00
160
200
200
100
0.00
100
20.0 spooch
8.00
000
1.00
10
oporch
sparch
Leaninu til
i) Festures norms
i) Kernels aligumeet
111
|AA le
[Auug le
4.
(9/1)) //t///>7 2"//1""M.YY///p7""
30
Our
InC
Our
Figure 3: VGGI1 trained on MNIST, See Figure 2 for the description of panes a-h. i) Norms of matrices H, H, H2H and (h) (h)T the end training. Alignment of kernels O and @ the end of training. The color in panes i-j is the color of the same model in panes a-e.
a) Inveriant noem
b) Inv. alignment to W
1.0
0.94
Lecuun nomell init.
LG5
03 125 (P") 0.06
somal
0.96
LX
93.24% (e"r
0.50
Urafon mr 12-4133
115
Anc: 93.0% (e) FY") 0.54
Unlhen 18
Anc: 93 07% (9h.1")
100
10-4
0.94
1.8
0.2
(e). F1");
oer. inil
us 25% "3
0.50
1.0
c) NC1
XC2
1.50
NC3
10.
1.25
1.25
1.00
1.00
bb) 50 coocnal ijit.
40.0
0.77
0.75
0.500
0.50
0.06
0..55
0.25
0.50
2041
0,00
150
210
01.00
1.00
200
300
4.0
6.00
0.50
Loamiz 1000
3) Features norms
1) Kernels alignnnent
LD
46
10/
ICC
MCC
Eod
No
Ne
ONC
m0
Figure 4: VGG1 trained on PashionMNIST. See Figure 2 for the description of panes a-h. i) Norms of matrices H, H. H,H, and (h) at the end of training. j) Alignment of kernels O and 6 at the end of training. The color in panes i-j is the color of the same model in panes a-c.
25

---

Invariant corm
207
WTW
1.0
LaChes nemual ist. LRt-oon
Ana 18.72% (9 (8'-17");
acrumng:
Lef normalig
NEL
0..5
0.5
82.91%, (e):23");
12-000
0.7
Ann: 89.28% (89) y3""; u.s
Uniken net 1R-D
100
Io Woromeallnt [R0-5.005
4.2
Ancn 66 54% (0") y3""
Crnform int.
01.5
c) NC1
d) NC2
6) NC3
1.30
1.50
1.20
1225
0.15
1.50
1.00
Wr.g
+.10
0.75
0.25
morraa Hill.
0.50
0.50
0.8
0000
0.20
0:25
1000
300
3000
1.00
0.001
0
100
200 spoeth
300
0.00
0.00
15:00
200 opech
300
9.0
1000
0110
10
Leuninn 2ad0
5) Festures norms
1) Kernels aliggment
L5
|
12,,22 l
10000R0 I
(9))|H)) 2 /
Figure 5: VGG16 trained on CIFAR10. See Figure for the deseription of panes a-h. i) Norms of matrices H, H. H2H and (h) at the end of training. j) Alignment of kernels O and @ the end of training. The color in panes i-j is the color of the same model in panes a-e.
a) Invariant narm
h) Tay. alggmmet to W"W
1.0
Tee comnany
1.00
F) Lechun normol Iuit
1.5
1.0
211
89 nos (9h yV?)
0.3
0.99
0.6
Unilare ist LRt
Ano: 00 60% (9) yr")
0.51
105
Anc: 99 00% (99)) 0.86
111
11.2
(eh po3))
UP-0655
1.00
14) Unitorr init.
Anc. 09 64% (6h (33") 0.02
Hr senual Li-4000)
0:39
1.0
NCI
NC2
NC3
10.
10.0
bj He SOcTTAAL Init
0.00
0.5
0.95
0:58
0.0
spieds
1) Features porms
3) Keruels alignmmnt
10
[MAAlle
[ML
80
1049 (A)T|e
62
(8/19) "
GrC
Voor
10.3
200
Figure 6: ResNet20 trained on MNIST. See Figure for the description of panes a-h. i) Norms of matrices H. H.H, and (h) (h) at the end of training. j) Alignment of kernels O and @ at the end of training. The color in panes i-j is the color of the same model in panes a-e.
26

---

a) Javariant RoOm
b) Inv. alignment W
L3
91 83% (Sh,Y"";
18b00001
accurary
LeClun nommel ini..
128-4
93.64% (ef)3"")
10.5
Juc. 93 00% (eh, 0.68
Untfirm ty L8 w000
0.02
6.5
Anc: US 75%11% (9h,3%%% 0.70
Unafer Li-ars
400
10
-0.2
Uniform init.
10.3
Re scmall sit. UU--0ds
ns 645, (9P,) y3yn
0,01
1.0
<) NC1
d) XC2
NC3
0,02
605
400
1.00
[T] Her normal iiii.
0..75
00.70
1.0
0.50
0.50
0,25
0022
60
0.225
000
1.00
200 epoch
0.00
1241
2000
420
0.00
20-2
epech
1) Restures nocns
Kersels aligunenn
(8/1ela "Y
200
Coof
Figure 7: ResNet20 trained on FashionMNIST. See Figure for the description of panes a-h. Norms of matrices H,H, H2H and (h) at the end of training. j) Alignment of kernels Q and h at the end of training. The color in panes i-j is the color of the same model in panes a-e.
Invariant
b) aliguument W WT
LeGra somad asl, 18-0000
Tost aceurary
coll LeChm normel Irit
101
Anc, 69 19%, "Y") 0.06
3.8
Loh somaal 12--uas
Anc: 8t 04% (ef.FT))
0.2
40% (e".F1")
0.5
Ii nomeal net 18: pans Arc: 81 wus, (o) y) o.
3.2
Uniforu
He sonuel LEE-poos
10
0.0
Anc: 86 55%, jef FT""); 0.17
0.24
11.20
1.50
1.20
5.15
1.25
1.25
0.5
1.00
1.0
0.8
He
0.75
0.2
spoch
speril
<poch
i) Fratures norms
j] Kemels aliguenent
[|
00
AcD Se.
Lerba NC
Foot
Bo
Figure 8: ResNet20 trained on CIFARIO. See Figure for the description of panes a-h. i) Norms of matrices H, H2H,J and (h) at the end of training. j) Alignment of kernels O and O the end of training. The color in panes i-j is the color of the same model in panes a-e.
27

---

Invuriant
10
b) lov, alieranent WT
LeCh cornel lot, 1R-0 361
Avo 90 21%
(0t) M99) 0.53
1000
Let nartel init
0.8
LoCa (0 (373) 0.75
Arc: 99 68%
1R
0.20
1.0
200
405
0.58
Calfen tel.. LRweUNE
QU4
Aue 67.77 (6h, M3?") 0.70
100
Be nomual LE-46677
Aee 88 00S (0t p/90)) o
[3 Unilfyr
10-1
10.0
Re
98 0bs. Moy))
0.72
01501
c) NCI
0..00
d) NC2
1.30
c) NC3
1.50
10.0
0.50
1.25
1.25
0.15
1.00
1.00
1.10
(x) Heaaomaa
1.5
0.30
0.70
0.70
1.0
0.004
10.00
0.500
0.00
H0D
0.25
0.25
100
2010
350
100
0.00
100
203
930
100
2013 spoch
800
410)
50.0
spoch
rpoch
1) Restures nocns
Kersels aligunend
122,41 li
(8/1ele "Y
20 MCMCC
Figure 9: DenseNet40 trained on MNIST. See Figure 2 for the description of panes a-h. i) Norms of matrices H, H,H. and (h) (h) at the end of training. j) Alignment of kernels θ and @ at the end of training. The color in panes i-j is the color of the same model in panes a-e.
Invariant
algammen W-W
LeCon reenol LB
ThN uontray
[]
10
Nex: 98.003 (9,Y"" 0.03
(9%)
0MM
Ace: (9/))9/)
Cafsen inn. LR-0104
Jon: 94083 (99))99)) 0.75
100
Beternal me LB u807
Auc: 91.73 (69, 335T) 0.56
33 Chiform
14.0 (9,,
c) NC1
0.200
d) NC2
1.50
c) NC3
2.00
0.50
1.25
25
0.15
0.50
1.00
100
I)) Hir mormed
6.10
8.75
6.75
1.50
0.7A
0.06
0.25
300
2000
3000
400
100
2001
300
010
8.00
100
200
0.00
10
spods
opoch
spoch
Laurning mabn
i) norms
Kernle immment
08
|AAA lle
1100 (A37
04
10/1OL VY7JKY"
(e PYTA
KrCWi Aua
AC
Bist
Figure 10: DenseNet40 trained on FashionMNIST. See Figure 2 for the description of panes a-h. Norms of matrices H, H. H.H, and (h) (h) at the end of training. j) Alignment of kernels e and e at the end of training. The color in panes i-j is the color of the same model in panes a-è.
28

---

a] Invaciant noem
10
b) Inv. aliggmmen
1.0
Lehnn aomaal 12-0077
Anon 76 90% (9 1/ 0.57
The sonmmny
0.5
10.6
Ann. 78.715 (e" F1",,
Vniken att, 18-40077
077
Uniken 12-44s4
Anc 88 53%, (9 11" 0.08
100
10-1
28-20007
0.3
70 (12"""
Cauform Inic
88.94%, I
NC1
d) NC2
e) NC3
1.30
1.0
0.2
11.25
1.26
10.15
1.50
1.0
3.0
11 He mntrane inii.
40.10
0.75
0.75
0.30
0.50
0.25
0005
11.20
0.20
200
300
1000
160
200 epoch
300
000
11.00
1.00
200
300
0.0
9.00
10
opoch
Lowniig rote
1) Fearnees norms
1) Kernels aliemment
[[A REll
11000 (AP
10/10L VYTANY?
Secho uoe to
00
MC
Fayt
Figure 11: DenseNet40 trained on CIFAR10. See Figure 2 for the description of panes a-h. i) Norms of matrices H H.H and (h) (h) the end of training. j) Alignment of kernels and O at the end of training. The color in panes i-j is the color of the same model in panes a-c.
a) Insuriant norm
b) Fa aliemment to W2W
LeChon monaa idht
F) LeCam normal Int.
Juc: M8.137 (9..))))
1eOtn ACmmmmalnn 18-0000
Anc: 18 88% (DF: VY") 9.6
0.99
10
Unienn 1I-23M5
Anc: U8 99.5677 (99.3
0:56
09.0155 (8"YY"") 0.83
100
99:2377
Lft-1100
0.45
1,00
3) Uniform
155
Hr auuee ini.. LR=it (43
00.
Ane. 09 09.015, (er MY
0.87
0.39
1.0
NC1
d) XC2
NC3
0.5
-0.30
40.6
0.8
100
11.00
h)) He normal init
0.25
40.4
0.99
0.36
100
230
320
400
100
210 1preds
200
010
150
200 epeczh
350
420
Tearring
1) Features borms
}] Keruels alienment
25.2 B
TEEAEE le
52
GoCoo
Tefler
Ougr
Figure 12: ResNet20 trained MNIST with CE loss. See Figure 2 for the description of panes a-h.
i) Norms of matrices H, H. H2H, and (h) T at the end traming. j) Alignment of kernels and h at the end of training. The color in panes i-j is the color of the same model in panes a-e.
29

---

a) (((()
b) 2.,
c) 1
d) 160 1
n) Aligument
nar
50
03
13
02
12
Ave
200
200
Figure 13: NTK block structure of VGG11 trained on MNIST. LeCun normal initialization, initial learning rate 0.131. The kernel is computed on random data subset with 4 samples from each class. See Figure D for the description of panes.
a) Σ. (x())
b) Σ. (X)
d)
10 1r
Aligument
20
Aa
07
Rer
18
Pam
2.0
Pas
Figure 14: NTK block structure of VGGI1 trained on FashionMNIST. LeCun normal initialization, initial learning rate 0.049. The kernel computed on random data subset with 4 samples from each class. See Figure for the description of panes.
a) Σ. (x)
b) Σ. (2))
c) [9,
d) |0),1 .1r
+) Aligmment
Aor
1.5
02
01
48
3.0
Ave
TTP PE FF
hur
231
2003
Figure 15: NTK block structure of VGGI] trained on CIFARIO. LeCun normal initialization, initial learning rate 0.131. The kernel is computed random data subset with 4 samples from each class. See Figure I for the description of panes.
a) ((x))
b) Ex)(,()
[9u.]
d) |
c) Aligument
02
1.
25
100
Figure 16: NTK block structure of ResNet20 trained on FashionMNIST. LeCun normal initialization, initial learning rate 0.094. The kernel computed on a random data subset with 12 samples from each class. See Figure for the description of panes.
30

---

a) ((,(((
b) (x,
c) 1
d) [6 1
O) Aligument
03
55
02
12
2/
22
200
Figure 17: NTK block structure of ResNet20 trained on CIFARIO. LeCun normal initialization, initial learning rate 0.068, The kernel computed on a random data subset with 12 samples from each class. See Figure 1 for the description of panes.
a) (x,)
b) (X)
d) |0%,1*
Aligument
Par
07
01
Am
56
Aec
3:4
Au
Figure 18: NTK block structure of DenseNet40 trained on MNIST. LeCun normal initialization, initial learning rate 0.049. The kernel computed on a random data subset with 12 samples from each class. See Figure [ for the description of panes.
a) Σ. ((,()
b) Σ. (2)
c) [9,
d) |0)
+) Aligmment
Aa
02
3.0
01
48
2.4
Au
12
231
2003
Figure 19: NTK block structure of DenseNet40 trained on FashionMNIST. LeCun normal initializa- tion, initial learning rate 0.094. The kernel computed on a random data subset with 12 samples from each class. See Figure I for the description of panes.
a) (x)
b) Ex)l,()
d)
| 1
c) Aligument
Aur
02
1.6
Aur
[T FTT
2.0
100
Pu
2.0
Figure 20: NTK block structure of DenseNet40 trained on CIFARIO. LeCun normal initialization. initial learning rate 0.094. The kernel computed on a random data subset with 12 samples from each class. See Figure I for the description of panes.
31


Reviews:
[{"review_id": "wMnfIqgeZ9", "reviewer": "Reviewer_PDBS", "summary": "This paper connects NTK (neural tangent kernel) and NC (neural collapse) by assuming NTK has a block structure in the late training stage, meaning the kernel for samples within the same class is much larger than samples from different classes. The technical assumption additionally assumes that the gradients of classification heads are independent of each other, and the gradients of neurons of the penultimate layer are also independent of each other. The authors then claim that the kernel gradient descent of this NTK leads to NC. Empirical results provide the correlation between the block structure and NC.", "questions": "Can you verify the calculation I did above?", "limitations": "I do not see any potential negative societal impact.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is the first paper I have seen that builds the connection between NTK and NC. It is inspiring in concept with good novelty.\n\nThe presentation of the paper is overall good and easy to read. ", "weaknesses": "I believe the assumption of the paper is too strong because the authors not only assume the block structure of NTK but also assume the independence between gradients of neurons/classification heads for any input. Taking a closer look at Assumption 3.2, one can already derive strong results. For example, one can wrote $\\nabla f_k(x)$ as a function of $\\nabla h_s(x)$, $\\mathbf W$ and $h_s(x)$ by chain rule, and plug into $\\Theta$. Then it is not hard to find that \n\n-  $\\langle \\mathbf W_k, \\mathbf W_s\\rangle$ is identical for any $k\\neq s$;\n- $\\langle \\mathbf W_k, \\mathbf W_k-\\mathbf W_s\\rangle=\\gamma_d/\\kappa_d=\\gamma_c/\\kappa_c=\\gamma_n/\\kappa_n$;\n- $\\langle h(x), h(x')\\rangle =-\\kappa \\langle \\mathbf W_k, \\mathbf W_s\\rangle$ ($\\kappa$ is $\\kappa_d, \\kappa_c$ or $\\kappa_n$ depends on the relationship between $x$ and $x'$).\n\nThose are some implications I derived within an hour. There might be some flaws but the message is clear --- the independence between gradients is a very very strong assumption.\n\nAnother thing I want to point out in the proof is that the training may not be in the kernel regime. The authors assume after hundreds of training epochs, the empirical NTK satisfied the property, but it is not clear whether it changes afterward (meaning it is not in the kernel regime but feature learning regime, which does not support linearization of the model). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for rising the concern about the assumption of the independence between output/feature neurons, i.e., the part of Assumption 3.2 stating that $\\Theta_{k,k'}(x,x')=\\Theta^h_{k,k'}(x,x')=0$ for any $k\\neq k'$ and any $x,x'$. While the calculations in the review do not appear fully correct (we provide details below), we agree that the independence assumption is quite strong. In fact, as we show below, we can get rid of the independence part of the assumption with minimal changes to the paper.\n\n**Relaxation of the assumption on $\\Theta^h$:** Based on the review, we identified that we can relax our main assumption without significant changes to the paper in the following way:\n$$\\Theta^h_{k,k'}(x,x')=\\beta\\cdot\\Theta^h_{k,k}(x,x')\\quad \\forall k=1,\\dots,n,\\forall k'\\neq k $$\nfor some $0\\leq\\beta<1$. Note that our original assumption is the special case of this assumption with $\\beta=0$. Instead of requiring independence between different features, the relaxed assumption only states that the dependence between different features is weak, which is indicated by $\\beta$.\n\nThis relaxation amounts to a simple change in the training dynamics of DNNs with block-structured NTK. Indeed, in the proof of Theorem 4.1 (Appendix B.1), we obtain the following dynamics for $h_s(x_i)$:\n$$ \\dot h_s(x_i)=-\\sum_{s'=1}^n\\sum_{i'=1}^N[ \\mathbf W\\mathbf H-  \\mathbf b\\mathbf 1_N^\\top-\\mathbf Y ]\\_{s'i'}\\Theta^h_{s',s} (x_i,x_{i'})$$\nVectorizing each term of the outer sum, we get the dynamics of the whole features matrix:\n$$\\dot{\\mathbf H}=-\\mathbf A[\\mathbf W^\\top(\\mathbf W\\mathbf H+\\mathbf b\\mathbf 1_N^\\top-\\mathbf Y)][(\\kappa_d-\\kappa_c)\\mathbb I_N+(\\kappa_c-\\kappa_d)\\mathbf Y^\\top\\mathbf Y + \\kappa_n\\mathbf 1_N\\mathbf 1_N^\\top],$$\nwhere $$\\mathbf{A}:=(1-\\beta)\\mathbb I_n + \\beta\\mathbf 1_n\\mathbf 1_n^\\top.$$\n\nThus, adding the dependence between the feature neurons amounts to \"scaling\" the dynamics of $\\mathbf H$ by an invertible matrix $\\mathbf A$. The dynamics of $\\mathbf W$ and $\\mathbf b$ remain unchanged.\n\nWe notice that such scaling does not effect our proof of variability collapse (NC1). Indeed, applying a change of variables $\\tilde{\\mathbf E}:=\\sqrt{\\mathbf A}{\\mathbf E}\\sqrt{\\mathbf A}$ and $\\tilde{\\mathbf H}_2:=\\sqrt{\\mathbf A^{-1}}\\mathbf H_2$, we can carry on the same proof and show that $\\tilde{\\mathbf H}_2\\to\\mathbb O$ and thus $\\mathbf H_2\\to\\mathbb O$. Similarly, we can apply a change of variables $\\tilde{\\mathbf W}:=\\mathbf W\\sqrt{\\mathbf A}$ and $\\tilde{\\mathbf H}_1:=\\sqrt{\\mathbf A^{-1}}\\mathbf H_1$ in the proofs of NC2-4 to show the duality $\\tilde{\\mathbf H}_1\\propto\\tilde{\\mathbf W}^\\top$. From this duality and the assumptions of Theorem 5.1, we get the ETF structure of the features matrix and the duality $\\mathbf H_1\\propto\\mathbf W^\\top$. Thus, the statement of the main Theorem 5.1 remains unchanged under the relaxed assumption.\n\nWe will adopt the relaxed version of the assumption in the revision of our paper. \n\n**Relaxation of the assumption on $\\Theta$:** It is also possible to relax the assumption on the NTK $\\Theta$ in the same way:\n$$\\Theta_{k,k'}(x,x')=\\beta'\\cdot\\Theta_{k,k}(x,x')\\quad\\forall k, \\forall k'\\neq k$$\nfor some $0\\leq\\beta'<1$. This relaxation does not effect the dynamics derived in Theorem 4.1, which is the main target of our analysis. The assumption on $\\Theta$, in fact, only effects the convergence analysis in Section 4.1. The argument about the convergence rates still holds with the relaxed assumption, since the eigenvectors of the non-diagonal terms of the NTK $\\Theta_{k,k'}$ are the same as of the diagonal terms $\\Theta_{k,k}$.\n\n**Implications of the assumption:** Following the reviewer's reasoning and writing $\\nabla f_k(x)$ by chain rule as a function of $\\nabla h(x)$, $\\mathbf W$ and $h(x)$, we can derive the following two implications of our relaxed assumption:\n\\begin{align}\n    \\langle\\mathbf W_k,\\mathbf W_{k'}\\rangle&=-\\beta\\sum_{s\\neq s'}\\mathbf W_{ks}\\mathbf W_{k's'}+\\dfrac{\\Theta_{k,k'}(x,x')}{\\Theta^h_{s,s}(x,x')}\\quad\\forall k\\neq k',\\\\\\\\\n    \\langle\\mathbf W_k,\\mathbf W_{k}\\rangle &=-\\beta\\sum_{s\\neq s'}\\mathbf W_{ks}\\mathbf W_{ks'}+\\dfrac{\\Theta_{k,k}(x,x')}{\\Theta^h_{s,s}(x,x')}+\\dfrac{\\langle h(x), h(x') \\rangle+1}{\\Theta^h_{s,s}(x,x')}\n\\end{align}\nAs the reviewer correctly noted, this implies $\\langle\\mathbf W_k,\\mathbf W_{k'}\\rangle=const$ for any $k\\neq k'$ in the special case of $\\beta=0$ (but not if $\\beta>0$). The two remaining conclusions in the review do not seem to hold even in case $\\beta=0$. In particular, as far as we can see, neither the ETF structure of the class means (NC2), nor the duality between the weights and the features matrix (NC3) follow from these equations. Moreover, even variability collapse (NC1) does not in general follow from the assumption. \nA trivial counterexample where $\\mathbf H$ satisfies the above conditions but NC1-2 do not hold is given by the following configuration of the feature vectors: $h(x_1^{c_1})=(1,0,0,0)$, $h(x_2^{c_1})=(1/\\sqrt{2},1/\\sqrt{2},0)$, $h(x_1^{c_2}) = (0,0,0,1)$, $h(x_2^{c_2}) = (0,0,1/\\sqrt{2},1/\\sqrt{2})$.\n\n**Kernel regime:** There is empirical evidence that high levels of NTK alignment are achieved before the loss decreases to near-zero values (see Figure 3 in [1]), and that the kernel changes most rapidly in the early stages of training (see [1,2]). These observations justify the analysis of the dynamics with block-structured NTK before the terminal phase of training. Nevertheless, we of course do not claim that the NTK of real DNNs remains completely constant even in the end of training. Generalizing the analysis to non-constant NTK is an interesting but challenging future work direction. \n\n**References**\n\n[1] Atanasov et al. Neural networks as kernel learners: The silent alignment effect. (2021).\n\n[2] Fort et al. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. NeurIPS (2020)."}, {"review_id": "SIMOmW5exp", "reviewer": "Reviewer_LnGq", "summary": "This work provides a theoretical connection between two related phenomena in deep learning dynamics: the structural change to the empirical NTK during training, specifically its alignment with class labels; and the neural collapse, which refers to a set of behaviors exhibited by NNs trained on multiway classification. The authors asked whether the first phenomenon causes the second, and studied NN dynamics under a toy-model NTK that has a block structure.", "questions": "I wonder how much the analysis can be generalized to weakened assumptions about the kernels? For example, instead of assuming that the kernel has a block-like structure, can we assume that it has the form K=K0 + Kb, where Kb is the block-like structure proposed by the authors, and K0 is the eNTK at initialization? I think this is a much more realistic assumption about the dynamics of NN learning.\n\nI also welcome additional comments from the authors about the insights (of course, what counts as \"insights\" is different for different people!) that this work brings about the emergence of block structures / NC.", "limitations": "Please see weaknesses. There are no negative societal impact. I appreciate the author discussing the weakness of making the assumptions in 3.2", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Both neural collapse and NTK alignment are prominent empirical features of DNN learning (at least in some cases), and the authors provided a careful treatment of a relevant toy model that connects the two. The writing is clear despite the convoluted nature of the subject (speaking as someone who has written on this specific topic before). Assumptions, claims and proven results are easy to find and understand.", "weaknesses": "My primary concern is with the amount of insight that this work brings.\n\n(1) The central question posed by the authors, in line 37, is stated as \"does NTK alignment lead to neural collapse?\". But I do not understand the cause-and-effect relation implied here. It has been long understood that NTK provides a dual perspective to NN training dynamics. So much as the NTK does not \"lead to\" the reduction of loss in NN training, NTK alignment does not \"lead to\" neural collapse. \n\n(2) My second point is closely related to the first one. The assumptions about the NTK kernel and feature kernel in Sec. 2.2 are motivated by empirical observations of trained NNs and make ensuing analysis simpler. But I'm concerned that they are sufficiently strong to \"guarantee\" neural collapse -- these are very strong assumptions and already imply that the NN representation/dynamics are in a stage where things are \"collapse-y\". In other words, on a conceptual level, this work risks falling into a circular logic -- if we assume the NN is in neural collapse, we can derive that it is in neural collapse. This work does not touch on the topic of how the kind of block-like structure arises, which in my view is what \"causes\" neural collapse (Conceptually, I think the emergence of block-like kernels and NC are two sides of the same coin).\n\n(3) As the authors pointed out, an important implication of NC is the generalization benefits that it brings. The current analysis, however, appears to entirely focus on the training set. This is a limitation shared by previous work on empirical NTK dynamics, though.\n\n(4) A minor point, but I find Sec. 4.1 to be a little excessive. It is well known that if you do regression with a kernel, eigenstructure of the kernel determines how quickly different components of the target vector gets learned. There are plenty of prior work on the subject. Perhaps it is good to cite them and shorten this section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful evaluation of our work! Below are our responses to the reviewer's questions and concerns.\n\n**Does NTK alignment cause NC?** We fully agree with the reviewer that we do not show a causal relationship between NTK alignment and NC but rather explore the connection between the two empirical phenomena. We will make this more clear in the introduction of the revision. Specifically, we will avoid the misleading statement that NTK alignment \"leads\" to NC. However, there is empirical evidence that high levels of NTK alignment are achieved before the loss decreases to near-zero values (see Figure 3 in [1]), and that the kernel changes most rapidly in the early stages of training (see [1,2]). These observations justify the analysis of the dynamics with block-structured NTK before the terminal phase of training. Combined with these observations, our results also suggest that we could use NTK alignment (together with the invariant identified in our paper) to ``predict'' NC before convergence. \n\n**Are NTK alignment and NC two sides of the same coin?** Our results indicate that NC and NTK alignment are related but are not the same thing. Indeed, we show (both theoretically and empirically) that DNNs with NTK alignment do not always converge to NC. Moreover, our assumption on the dynamics invariant provides a necessary condition for the emergence of NC in DNNs with block-structured NTK (see Theorem 5.1 and the following discussion). Thus, we believe that NTK alignment is a more common phenomenon than NC, and this insight is one of the main contributions of our work.\n\n**Does NC follow trivially from the NTK block structure?** While the connection between NTK alignment and NC may seem obvious to some readers on the intuitive level, establishing such connection theoretically is not entirely trivial and, to the best of our knowledge, it has never been done before in the literature.\nIn the response to Reviewer PDBS, we explain why NC does not follow trivially from the assumption on the NTK block structure. Moreover, as we mentioned above, NC does not even always occur in DNNs with block-structured NTK. \n    \n**Broad implications:** Our work proposes to study NC through the lens of NTK alignment, which opens new research directions. Previous works on NC focus on the top-down approach (layer-peeled models) and fundamentally cannot explain how NC develops through earlier layers of a DNN, what are the effects of depth, etc. On the other hand, NTK alignment literature focuses on the alignment of individual layers, and recent theoretical works even quantify the role of each hidden layer in the alignment [3]. Therefore, we believe that the connection between NTK alignment and NC established in our work provides a conceptually new method to study NC. We will include a discussion section in the revision to explore the implications of our work in more detail.\n\n**Weakened kernel assumption:** The NTK at initialization is random (w.r.t the random initialization parameters) and has an approximately diagonal structure (see [4,5]). Therefore, a kernel of the form $\\Theta = \\Theta_{block} + \\Theta_0$ can be modeled as a certain block-structured kernel $\\Theta_{block} := a\\mathbb{I} + b\\mathbf Y^\\top\\mathbf Y + c\\mathbf{1}\\mathbf{1}^\\top$ plus an i.i.d. noise term. Under this model, the dynamics of the features' class means $\\langle h \\rangle_c$ can be viewed as an approximation of the expected dynamics for each class. Thus, assuming that the noise is centered and the sample size is large enough, the noise should not change the dynamics of the class means too much. However, quantification of the randomness in such a setting may be theoretically challenging and we propose it as a future work direction in our paper.\n\n**References**\n\n[1] Atanasov et al. Neural networks as kernel learners: The silent alignment effect. (2021).\n\n[2] Fort et al. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. NeurIPS (2020).\n\n[3] Lou et all. Feature learning and signal propagation in deep neural networks. ICML (2022).\n\n[4] Xiao et al. Disentangling direct from indirect relationships in association networks. (2022).\n\n[5] Seleznova et al. Neural tangent kernel beyond the infinite-width limit: Effects of depth and initialization. ICML (2022)."}, {"review_id": "scXIqOLIla", "reviewer": "Reviewer_Hrgu", "summary": "The main contribution of the paper \"Neural (Tangent Kernel) Collapse\" is the connection of the Neural Tangent Kernel (NTK) alignment and Neural Collapse (NC) phenomenon in deep neural networks (DNNs). The authors assume that the empirical NTK develops a block structure aligned with the class labels. They derive the dynamics of DNNs trained with mean squared (MSE) loss and identify three different convergence rates for certain components of the error. They also identify a hyperbolic invariant that captures the essence of the dynamics and use it to prove the emergence of NC in DNNs with block-structured NTK. Also, it provides large-scale numerical experiments to support their theory. Overall, the paper provides valuable insights into the dynamics of DNNs with block-structured NTK and the emergence of NC.", "questions": "1. This paper only focuses on DNNs trained with mean squared error (MSE) loss. Since Cross Entropy(CE) loss is a common choice for training classification networks, does this dynamics of DNNs and conclusions in this paper also hold for CE loss? \n\n2. Can we get some theoretical insights from this work to improve the design and training of deep neural networks?", "limitations": "The authors have addressed their limitations in chapter 7. There is no negative societal impact to be expected from this work", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Originality: Yes, this is the first work to connect NTK alignment and NC. Another new contribution is exploring the dynamics of DNNs with block-structured Neural Tangent Kernel (NTK).\n\n- Quality: High quality, the paper presents a thorough analysis of the dynamics of DNNs with block-structured Neural Tangent Kernel (NTK) and the emergence of Neural Collapse (NC) phenomenon. Also, they provide large-scale numerical experiments on three common DNN architectures and three benchmark datasets to support their theory.\n\n- Clarity: Yes, the paper is well-structured and clearly explains technical terms and concepts, making it easy to understand for readers.\n\n- Significance: Yes, for explaining the emergence of NC, most theoretical works adopt the unconstrained feature models. However, this paper provides a new point of view of NTK to explain it. It makes a step towards realistic DNN dynamics by means of the NTK.\n\n", "weaknesses": "The main weakness of the paper is that its fundamental assumption 3.2 is not justified enough. It assumes the block structure of the NTK kernel but it’s often not the case in real DNNs which is not well-trained.\n\nSome minor remarks:\n1. I think the subtitles for a) and b) in Figure 1 should be the **sum** over the classes or feature dimensions.\n2. Equation 2: missing part of the parentheses.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive evaluation of our work! Below are our responses to the questions.\n\n**Cross Entropy loss:** Generalizing our theoretical results to CE loss is challenging, since the dynamics equations with CE loss are more complex than in case of MSE even with block-structured NTK. In general, CE loss is difficult to analyze, so theoretical NC papers typically focus on unconstrained features models dynamics with MSE loss [1,2]. To the best of our knowledge, NC papers that consider CE only analyze global minimizers of the loss function and do not study dynamics equations [3,4,5]. However, we agree that it is important to verify whether our conclusions also hold for CE loss, and will provide additional empirical results for CE loss in the appendix of the revision.\n\n**Insights for design and training of DNNs:** While not all DNNs display NC, little is known about particular factors that determine whether a given DNN would converge to NC or not. Our work may shed some light on the importance of weights decay and batch normalization for the emergence of NC, which previous literature also conjectured (see [1,6]). In particular, our assumption on the invariant (which is necessary for the emergence of NC) has effects similar to regularization (see e.g. Appendix A.1 in [7]), while the zero global mean assumption is related to batch normalization. \n\n**References**\n\n[1] Han et al. Neural collapse under mse loss: Proximity to and dynamics on the central path. ICLR (2022).\n\n[2] Mixon et al. Neural collapse with unconstrained features. (2020).\n\n[3] Lu \\& Steinerberger. Neural collapse with cross-entropy loss (2020).\n\n[4] Zhu et al. A geometric analysis of neural collapse with unconstrained features. NeurIPS (2021).\n\n[5] Wojtowytsch et al. On the emergence of simplex symmetry in the final and penultimate layers of neural network classifiers. (2020).\n\n[6] Ergen \\& Pilanci. Revealing the structure of deep neural networks via convex duality. ICML (2021).\n\n[7] Tirer \\& Bruna. Extended unconstrained features model for exploring deep neural collapse. ICML (2022)."}, {"review_id": "LiRdvsNXI6", "reviewer": "Reviewer_UoyH", "summary": "The paper proposes a mechanism behind the empirical phenomenon of Neural Collapse in deep neural networks. The paper derives and analyzes the training dynamics of DNNs with MSE loss and block-structured NTK, identifying three distinct convergence rates in the dynamics.", "questions": "Please dicuss the relationship of your work with this phenomenon [cite1]\n\n[cite 1] Liu D, Wang S, Ren J, et al. Trap of feature diversity in the learning of mlps[J]. arXiv preprint arXiv:2112.00980, 2021.", "limitations": "The limitations of the paper include the focus on MSE loss and block-structured NTK, which may not be applicable to other loss functions and NTK structures. The paper also does not provide practical solutions to prevent or mitigate Neural Collapse in deep neural networks.", "rating": 5, "confidence": 2, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "The strengths of the paper include its theoretical rigor, the clarity of its presentation, and the large-scale numerical experiments that support the theory. The paper also provides a new perspective on the empirical phenomenon of Neural Collapse and identifies the conditions under which it occurs. ", "weaknesses": "The weaknesses of the paper include the assumption of balanced datasets and the lack of exploration of the effects of non-balanced datasets on the dynamics of DNNs with block-structured NTK. The paper also does not explore the effects of adding stochasticity to the dynamics considered in the paper. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive evaluation of our work! Below are our responses to the reviewer's questions and concerns.\n\n**Unbalanced datasets and stochasticity:** Although we do consider only balanced datasets, we believe that our analysis could in principle be generalized for unbalanced data. This would amount to considering a block-structured NTK with blocks of different sizes. Since such generalization would make the calculations more cumbersome, we propose this as one of the future work directions in our paper. We also refer adding stochasticity to the dynamics to future work. As we note in the response to Reviewer LnGq about a possible weakened kernel assumption, we believe that adding centered noise to the kernel in our model should not significantly change the dynamics of features class means. We would also like to point out that, although our theory does not include noise, the experiments certainly include some stochasticity.\n\n**Relationship to [cite1]:** We looked into the work by Liu et al. [cite1], which studies the two-phase phenomenon in the training of MLPs, where feature diversity decreases in the first phase and then increases in the second phase. We believe that this phenomenon may be related to the dynamics of NTK alignment during training. Indeed, increasing similarity between features gradients in the first phase could mean that the NTK values first become more similar across the whole dataset, and only afterwards the block structure emerges. Thus, a possible future work direction is to study and compare the NTK structure in the first and second phases of training."}, {"review_id": "bQwQtiuQyr", "reviewer": "Reviewer_MvSr", "summary": "Previous work has observed an increased alignment and emergence of an approximate block structure in a trained network's Neural Tangent Kernel (NTK) as well as the Neural Collapse (NC) phenomenon in the last hidden layer. The paper attempts to connect the two by showing that in an extreme case of perfect block structure in NTK, NC follows from the dynamics of squared error minimization on the kernel.", "questions": "See weaknesses", "limitations": "See weaknesses", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is extremely well-written and well organized. I did not check the details of the proofs in the appendix but main paper is correct as far I can verify. The relevant literature is well covered in the paper and even more so in the appendix. The invariance derived to establish the connection can be an interesting result on its own and its cooccurrence with NC is demonstrated in realistic experiments.", "weaknesses": "While I enjoyed reading the paper, I cannot come to agree with the main premise of the work and so I'm voting for rejection. The results rely on a perfect block structure which is unlikely to be achievable in a realistic setting. Even in a restricted setting where such structure will emerge, it will likely emerge along with NC at the end of training and it is hard to argue that it is a cause or driving force behind NC. I will elaborate on these comments (W1, W2) as well as other weaknesses. Overall, a block NTK would have been interesting as a theoretical model (like unconstrained features) and the derivation of NC from this model is insightful, but the current paper is pushing this model far beyond its domain of applicability by assuming that the NTK of the whole network has this structure.\n\nW1: Perfect block structure in NTK assumes the gradients (of an output) are perfectly orthogonal across data points from different classes. I think a necessary requirement for this would be that the data is already linearly separable at the input or the first hidden layer. Otherwise there is some inevitable interference between the classes in the first few layers where data is not linearly separable, making gradients correlated in those layers and preventing the block structure, no matter how long the network is trained. The paper demonstrates a nearly perfect block structure on the extremely simple task of MNIST and the extremely overparametrized network of ResNet20. The perfect block structure is likely far from realistic in larger tasks and smaller networks.\n\nW2: If the block structure was a property of the NTK at initialization then one could argue for a causal relationship from the block structure to NC. The problem is, both the structure and NC are emergent properties of training dynamics and it is not clear to me why the former would be causing the latter. Previous work arguing that NTK alignment facilitates optimization focuses on alignment already present at initialization or on a partial increase in alignment within an early phase of training. It may be reasonable then to argue that the rest of the training will be faster compared to training with a misaligned NTK. The perfect block structure in this paper is not established as an early phase phenomenon.\n\nW3: I'm surprised by the block structure in the individual hidden layer neurons in Fig 1 (b,d). While the output layer neurons are individually fit to the labels, the hidden layer neurons are updated as a whole through backpropagation. Emergent properties in this layer are then likely to be rotation invariant in general. The fact that each individual neuron is showing a pattern here could be a byproduct of the activation function (ReLU). Does the same pattern emerge with leakyrelu activations (with default parameters)? If not, the current results are limited to ReLU. (I may be mistaken here and perhaps the property is rotation invariant.)\n\nW4: The main result assumes centralized activations and the argument is that Batch Normalization (BN) to some extent enforces this property. Are the representations in the experiments extracted right after a BN layer or right after a ReLU layer? The former would partially satisfy this assumption but the latter would in fact ensure that this is not the case. The text refers to Fig 3 in the appendix but some of the numbers in that figure are ~10. Even if they were smaller, I'm not sure how they would support this assumption since it seems to me that it is not the absolute magnitude of the mean but its magnitude relative to the overall norm of the features that matters.\n\nW5: Previous work on NTK alignment (Baratin et al, Atanasov et al) generally shows a final alignment in the approximate range of 0.1 and 0.5 in different networks and tasks. The kernel alignments in the submission are consistently higher. I understand that different experiments in these works would simply result in different alignment values. Nevertheless the large gap suggests to me that there is a key factor (in training setup or centralizing the kernels or another aspect of the experiment) that differs between the submission and previous work which should be spelled out to prevent the reader from drawing premature conclusions across these papers.\n\n---------\n\n**Post-rebuttal:** I read the other reviews and responses. The rebuttals answered my comment about causality (W2) and so I raised the score. In particular, the response clarified that NTK alignment as measured in the submission is similar to previous work and that it is established as an early-time phenomenon (while NC is a late-time phenomenon). The intro should mention this before claiming a causal relationship. Also showing an example of early increase in alignment instead of the current example in Fig. 1 (e) helps with motivating the causal relationship.\n\nRegarding W1: I agree that the theory doesn't assume a perfect block pattern. My issue is what the title of this paper implies. Unless NTK shows a near perfect block pattern like Fig. 1 (a-d) in more general cases, calling its behavior \"collapse\" is misleading. The revision should either provide results like Fig. 1 (a-d) in more general cases or clarify that this figure is a rare and extreme case.\n\nOther reviewers argued that the new result has little significance beyond what what is known from unconstrained features model. I do not have the expertise to comment on this issue and leave this decision to other reviewers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful evaluation of our work and for raising their concerns! Below are our responses.\n\n**W1 orthogonality:** We believe that the reviewer may have misunderstood our main assumption on the NTK structure, since we do not assume that \"the gradients (of an output) are perfectly orthogonal across data points from different classes\". Our Assumption 3.2 states that $\\Theta_{k,k}(x,x')=\\langle\\nabla f_k(x),\\nabla f_k(x')\\rangle=\\gamma_n\\in\\mathbb R_+$ for any two elements of different classes $x$ and $x'$. The value of $\\gamma_n$ is allowed to be non-zero in our assumption, and it is of course non-zero in practice for the reasons that the reviewer mentioned. We would also like to point out that the reviewer's claim that we assume \"the NTK of the whole network has [block] structure\" is not correct. We only assume the block structure of the NTK in the last two layers. Therefore, the kernels in the earlier layers are allowed to take arbitrary form.\n\n**W1 linear separability:** Neither NTK alignment nor NC implies that the data is linearly separable in earlier layers. Since we only make assumptions on the kernel in the two outer-most layers, there is a lot of flexibility in the inner layers for DNNs to lift the (possibly non-linearly-separable) data into another high-dimensional space, where we can separate them.\n\n**W1 adequacy of the block structure assumption:** While the illustrative Fig.1 only shows the NTK block structure for ResNet20 trained on MNIST, our experiments report the kernel alignment values for three architectures (ResNet, VGG and DenseNet) and three datasets (MNIST, CIFAR-10 and FashionMNIST). The alignment is high in most of our experiments, so the kernels display an approximate block structure in each setting. This shows that our assumption is realistic for the considered dataset-architecture pairs. We note that our experiments intentionally cover DNNs that display NC in previous literature. While strong NTK alignment may not occur in other DNNs (e.g., smaller or less fit to the task), there is no evidence that NC would occur in such DNNs either.\n\n**W2 causality:** We fully agree with the reviewer that we do not show a causal relationship between NTK alignment and NC but explore the connection between the two empirical phenomena. We will make this more clear in the introduction of the revision. Specifically, we will avoid the misleading statement that NTK alignment \"leads\" to NC. However, there is empirical evidence that high levels of NTK alignment are achieved before the loss decreases to near-zero values (see Figure 3 in [1]), and that the kernel changes most rapidly in the early stages of training (see [1,2]). These observations justify the analysis of the dynamics with block-structured NTK before the terminal phase of training.\n\n**W3 rotation invariance:** Since the traced kernel displayed in Fig.1b is a matrix inner product $\\langle\\nabla_wh(x),\\nabla_wh(x')\\rangle=Tr(\\nabla_wh(x)\\nabla_wh(x')^\\top)$, it is of course rotation invariant w.r.t. $h$ (i.e., it does not change if we multiply $h(x)$ by an orthogonal matrix $A$). Since we assume $\\Theta^h(x,x')=\\nabla_wh(x)\\nabla_wh(x')^\\top = \\kappa\\mathbb I_n$, our assumption is also rotation invariant. We thanks the reviewer for this observation and will add a remark about rotation to the revision.\n\n**W4 batch normalization:** We extract $h(x)$ after ReLU in the experiments and we agree with the reviewer that this setup ensures that the global features mean is not exactly equal to zero in practice. However, our experiments show that the global mean is insignificant (at least an order of magnitude smaller) in comparison with the class means displayed in the same figures, which supports our assumption. We also examine what happens if we discard the zero global mean assumption in the discussion after Theorem 5.1, point (4). In this case, the dynamics does not have to converge to perfect NC.  \n\n**W5 alignment values:** Atanasov et al [1] report the alignment values on MNIST for 2-layer MLPs, which are much less powerful architectures than we used in our paper. Previous works show that the alignment increases with depth (see Figure 5 in [3], where the alignment on MNIST reaches 0.75 for a 5-layer MLP), and this increase is associated with better performance. Thus, it is expected that our models display better alignment on MNIST. On the other hand, we report alignment around 0.5 for ResNet trained on CIFAR-10  (see Figure 8 in Appendix C), which is more consistent with Atanasov et al [1]. We believe that the numbers in Baratin et al. [4] are lower because they measure the alignment of matrices of size $NC\\times NC$ created by concatenating the features $\\nabla f_k(X)\\in\\mathbb R^{N\\times P}$ and labels $\\mathbf Y_k\\in\\mathbb R^N$ over $k=1,\\dots C$, while all the other works (including ours) measure the alignment of the traced kernel $\\sum_k\\Theta_{k,k}(X)\\in\\mathbb R^{N\\times N}$. We note that we used standard architectures and training procedures, which are described in detail in Section 6. We also provide the code to check and reproduce our numerical results. \n\n**References**\n\n[1] Atanasov et al. Neural networks as kernel learners: The silent alignment effect. (2021).\n\n[2] Fort et al. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. (2020).\n\n[3] Shan & Bordelon. A theory of neural tangent kernel alignment and its influence on training. (2021).\n\n[4] Baratin et al. Implicit regularization via neural feature alignment. (2021).\n"}]

Author rebuttal:
We thank all the reviewers for their valuable feedback! Based on the reviews, we were able to identify several improvements for the paper, which we will incorporate in the revision. Below we summarize the reviewers' concerns, our responses, and the proposed changes to our paper. 

## Concerns about the NTK block-structure assumption
Reviewers PDBS, LnGq and MvSr expressed several concerns about Assumption 3.2 on the NTK block-structure. We summarize these concerns and our responses below:

- **Ortogonality of gradients between different classes:** Reviewer MvSr states that our assumption implies that "gradients (of an output) are perfectly orthogonal across data points from different classes". However, as we state in the response, we do not assume orthogonality between points from different classes, since $\Theta_{k,k}(x,x')=\gamma_n\geq 0$ in Assumption 3.2. Therefore, we believe that this concern comes from a misunderstanding of our main assumption.
- **Independence between gradients of different feature neurons:** Reviewer PDBS states that the assumption of the independence between different feature neurons is too strong, since it already implies strong results related to NC. Our response is two-fold: 1) We show that we can relax our Assumption 3.2 to remove the independence. Such relaxation does not change the paper's main Theorem 5.1 and requires only small adjustments of the proofs (change of variables). 2) We show that NC does not follow from Assumption 3.2.
- **NC follows trivially from the assumption:** Reviewer LnGq expressed a related concern that our assumption is "sufficiently strong to "guarantee" NC". However, our results indicate that NTK block structure does not guarantee convergence to NC, and our invariant assumption provides a necessary condition for the emergence of NC in DNNs with block-structured NTK. We also detail why NC does not follow from the assumption in the response to Reviewer PDBS.
- **Adequacy of the assumption:** Reviewer MvSr expressed a concern that block structure of the NTK is not realistic and implies linear separability of data in the earlier layers. In the response, we explain that linear separability in earlier layers does not follow from our assumption. Moreover, our numerical results and previous works on NTK alignment provide a large body of empirical evidence to justify our assumption. We believe that this concern may come from a misunderstanding of our assumption and its implications.


We propose the following changes in the revision of our paper to address the reviewers' comments and concerns:

- Based on the comments by Reviewer PDBS, we will adopt a relaxed version of Assumption 3.2 without independence of different features neurons. This improvement requires minimal changes in the paper's reasoning and proofs. 
- Based on the comments by Reviewer LnGq, we will modify the introduction to make it clear that NTK alignment and NC are not the same phenomenon and do not always occur together.  

Overall, we admit that our main assumption is a simplification of the realistic DNNs’ behavior. However, empirical evidence (including our numerical experiments) shows that our assumption can approximate well-trained realistic DNNs. Moreover, much like our assumption, perfect NC is also a simplification that never holds exactly in realistic DNNs. 

## Concerns about significance
Reviewers LnGq and MvSr expressed concerns about significance and implications of our work. We summarize these concerns and our responses below:

- **Causality:** Reviewers LnGq and MvSr share a concern about the causality between NTK alignment and NC. We agree with the reviewers that we do not show a causal relationship between NTK alignment and NC but rather explore the connection between the two empirical phenomena. However, we also reference empirical evidence that high levels of NTK alignment are achieved before the loss decreases to near-zero values, which justifies the analysis of the dynamics with block-structured NTK before the terminal phase of training.
- **Insights:** The primary concern of Reviewer LnGq is the amount of insight that our work brings. We respond that our results are 1) novel, 2) non-trivial, 3) open new research directions. Therefore, we believe that they are valuable for the community. 

We propose the following changes in the revision of our paper to address the reviewers' comments and concerns:

- We will modify the introduction to make it clear that we do not show causality between NTK alignment and NC. Specifically, we will avoid the misleading statement that NTK alignment "leads" to NC. 
- We will include a new section to discuss implications and broad impact of our work.
- To address the question of Reviewer Hrgu about generalization of our results to CE loss, we will add numerical experiments with CE loss in the appendix of the revision.

Overall, our paper provides the first (to the best of our knowledge) theoretical connection between NTK alignment and NC. Since both phenomena are 1) of high interest for the ML community and 2) hard to study theoretically, we believe that our results would be interesting for a wide audience.

## Conclusion

We believe that the proposed adjustments cover all the major concerns about the paper's technical contribution raised by the reviewers. These adjustments also do not require major changes of the paper's content and proofs. Moreover, while most of the reviewers agree that our results are novel and interesting, we hope that our response made the significance of our work even more clear. We are optimistic that the reviewers will consider our arguments and be open to a partial reevaluation of their scores. We will also gladly respond to any additional questions that may arise during the discussion period!

Meta-review: This paper establishes a connection between the Neural Tangent Kernel (NTK) and Neural Collapse (NC) by assuming that the NTK exhibits a block structure during the late training stage. This structure implies that the kernel values for samples within the same class are significantly larger than those for samples from different classes. The authors prove that the kernel gradient descent based on this NTK leads to the emergence of NC. Empirical results are presented to demonstrate the correlation between the block structure and NC.

The authors' responses have effectively addressed most of the concerns raised by the reviewers. Most reviewers have increased their scores and recommend acceptance, with the exception of one reviewer who appreciates the idea of connecting NC with NTK but expresses concerns regarding the block-structure assumption. As the block structure is approximately observed during training, and the connection between NC and NTK provides additional insights compared with previous approaches based on the unconstrained features model, I recommend acceptance. The authors are advised to address this concern by incorporating the discussions in the responses, along with other comments raised by the reviewers, in the final version of the paper.
Decision: Accept (poster)