Neural (Tangent Kernel) Collapse

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.

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. W1: 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. W2: 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. W3: 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.) W4: 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. W5: 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. --------- **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. Regarding 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. Other 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.

Questions

See weaknesses

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

See weaknesses

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.

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.

Questions

Please dicuss the relationship of your work with this phenomenon [cite1] [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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

2 fair

Contribution

3 good

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.

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.

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). - 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. - Clarity: Yes, the paper is well-structured and clearly explains technical terms and concepts, making it easy to understand for readers. - 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.

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. Some minor remarks: 1. I think the subtitles for a) and b) in Figure 1 should be the **sum** over the classes or feature dimensions. 2. Equation 2: missing part of the parentheses.

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? 2. Can we get some theoretical insights from this work to improve the design and training of deep neural networks?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors have addressed their limitations in chapter 7. There is no negative societal impact to be expected from this work

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.

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. (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. (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). (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. (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.

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. I 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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Please see weaknesses. There are no negative societal impact. I appreciate the author discussing the weakness of making the assumptions in 3.2

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.

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. The 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 - $\langle \mathbf W_k, \mathbf W_s\rangle$ is identical for any $k\neq s$; - $\langle \mathbf W_k, \mathbf W_k-\mathbf W_s\rangle=\gamma_d/\kappa_d=\gamma_c/\kappa_c=\gamma_n/\kappa_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'$). Those 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. Another 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).

Questions

Can you verify the calculation I did above?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

I do not see any potential negative societal impact.

Discussion

Thank you for the detailed rebuttal. Regarding W1: In Section 2.2, is lowercase **w** the parameters of the whole network or the last two layers? The kernel defined in Eq 2 is then assumed to have a near perfect block structure in the first part of Assumption 3.2. Are the parameters in this kernel w? And is this the whole network parameters or just the last two layers?

Thanks for your response

I appreciate your feedback. I would keep my current score regarding the response.

Thank you

Thanks to the authors for their careful response. The response and other proposed edits would address my concerns significantly. I am therefore raising my score.

Comments on the implication and relaxation

Thanks to the authors for the clarification! My quick calculation indeed had a bug, so my second and third implications are false. I also verify the two implications derived by the authors. However, if we take a closer look at the correct implications, we can find that, without the newly proposed relaxation, - $\langle W_k,W_{k'}\rangle=0$. - $\langle W_k, W_k\rangle=\gamma/\kappa+(\langle h(x), h(x')\rangle +1)/\kappa$ where $\gamma,\kappa$ depends on the relation between $x,x'$. They are very strong assumptions. Even with the relaxation, given the relation between $x$ and $x'$, $\langle h(x), h(x')\rangle$ is still a fixed value. Overall, I think the assumption (even with relaxation) is strong because of too much symmetry, which is convenient for proving NC. On the other hand, although I understand the changes in the proof might be straightforward when incorporating the relaxation, it is very hard for me to verify it confidently without reading it with all the changes added.