Wide Neural Networks Trained with Weight Decay Provably Exhibit Neural Collapse

We thank the reviewer for recognizing the strengths of our work and for the detailed comments. We address each of them below.

W1: The analysis critically relies on the fact that the last two layers of the neural network are linear. I can definitely see this condition makes the problem a lot easier to analyze. I am wondering how hard it is to remove such restrictions.

Response: This is a great point. One important assumption in our work is that there is a linear head in the network, containing at least two layers. However, experimental evidence suggests that two linear layers are not necessary for collapse to happen. Removing this restriction is likely to be difficult and it provides an exciting future direction, which would consist of using the approach developed here, i.e., the NTK analysis and the separation of timescales, to prove collapse without two linear layers.

W2: It seems the analysis of Theorem 4.4 relies on the neural network in the NTK regime, as the pyramidal topology assumption has appeared in previous works such as (Nguyen & Mondelli, 2020). I don't regard this as a major weakness even if it turned out to be true that the networks are in the NTK regime given the contribution of this work, however, I do appreciate clarification on this.

Response: Well, yes and no. In the first phase, we are essentially in the NTK regime and use NTK-type techniques to guarantee convergence (technically the pyramidal network setup is not exactly the typical NTK regime, see the last paragraph of Section 1 in (Nguyen & Mondelli, 2020); however, this distinction does not play a major role in our analysis). Nevertheless, there is a second phase, where the weight decay starts to take effect which is not NTK-like. Though our control of this second dynamics is weaker, we can still guarantee that the network will become more balanced and remain interpolating. To be more precise, we have a separation of timescales as $\lambda$ gets smaller: the number of steps needed to reach balancedness is of order $\frac{1}{\eta \lambda}$ (weight-decay timescale). This is significantly later than the interpolation time which is of order $\frac{1}{\eta}$ (NTK timescale). Note that at the end of training we end up with low-rank weight matrices and feature learning (hidden features that are different from their initializations), which could not happen in a purely NTK regime, so the second phase plays an important role.

This strategy/dynamics of “NTK followed by weight decay” is to our knowledge novel and it represents an important theoretical contribution of our paper. We have added a paragraph (NTK regime and beyond) in Section 7 of the revision about this aspect.

Q: I am wondering whether the layer balanced-ness property after training can be proved as a direct consequence of [1] in the author's setting. [1] Du, Simon S., Wei Hu, and Jason D. Lee. "Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced." Advances in neural information processing systems 31 (2018).

Response: This result is indeed very closely related, though what our analysis requires does not seem to be a direct consequence, mainly because we consider gradient descent instead of gradient flow, and we use weight decay/L2 regularization. Nevertheless, we have added a citation of this very relevant work.

We thank the reviewer for the positive evaluation of our work and for the insightful comments. We address each of them below.

W: The authors don't provide a conclusion or discussion section, and it would have been useful to have comments from the authors about what they think are the weakest parts of their work and what work should be prioritized in the future. I think they could get some additional space by removing some of the details about optimization from Sec 4 since I don't think they're super important.

Response: This is a great point. Following the reviewer’s suggestion, in the revision we have added a final section titled ‘Discussion and Concluding Remarks’. There, we discuss our assumptions, we compare with earlier work by Nguyen & Mondelli (2020) (see also the response to Q4 below), we discuss how our analysis relates to the NTK regime (see also the response to Q5 below), and we conclude with a direction for future research.

To make space and for the sake of readability, we have simplified the statements of Theorems 3.1 and 4.4 using a big-O notation, as suggested by Reviewer 1hwn.

Q1: Are there no assumptions on the activation function needed for Theorem 3.1? None are stated in Section 3, although assumptino 4.2 appears later. It seems odd that a potentially discontinuous and horribly behaved activation function would be permitted (I imagine Lipschitz is required?)

Response: Theorem 3.1 concerns the linear head containing at least two layers, after the non-linear part of the network. For this reason, no additional assumption is needed on the activation function.

Q2: Assumption 4.2 seems more like a smooth leaky relu, which has appeared in prior work [Chatterjee arXiv:2203.16462, Frei et al COLT 2022]

Response: This is a good point, thank you for pointing out these works. Our assumption is in fact the same as in [Frei et al, COLT 2022] and it is similar to that made in [Chatterjee arXiv:2203.16462]. We have edited the revision accordingly.

Q3: The authors talk about (NC1)-(NC3), what about (NC4) from the original Papyan paper?

Response: This is another good point. In the original Papyan paper (as well as in follow-ups), it has been shown that the NC4 property follows from the first three. Thus, if one proves the first three, the fourth one is automatically guaranteed. We did not discuss this explicitly, since many neural collapse papers already addressed this [4] and NC4 is not usually discussed in the recent NC papers [1, 2, 3].

[1] Tirer, Tom, Haoxiang Huang, and Jonathan Niles-Weed. "Perturbation analysis of neural collapse." International Conference on Machine Learning. PMLR, 2023.

[2] Súkeník, Peter, Marco Mondelli, and Christoph H. Lampert. "Deep neural collapse is provably optimal for the deep unconstrained features model." Advances in Neural Information Processing Systems 36 (2023).

[3] Jiang, Jiachen, et al. "Generalized neural collapse for a large number of classes." arXiv preprint arXiv:2310.05351 (2023).

[4] Wu, Robert, and Vardan Papyan. "Linguistic Collapse: Neural Collapse in (Large) Language Models." arXiv preprint arXiv:2405.17767 (2024).

We thank the reviewer for appreciating our work and for the detailed comments. We reply to each of the points raised in the review below. We have also uploaded a revised version of the manuscript highlighting in red the main changes.

W1: My main concern with the paper is the writing. The theoretical statements are very difficult to parse this is especially the case for eqs 2-5 in Theorem 3.1. Can the authors provide a more interpretable version of this theorem (hide constants, big-O etc.) and delay the details to the appendix?

Response: This is a very good point. Following the suggestion of the reviewer, we have re-written Theorem 3.1 in a more interpretable form hiding constants and using a big-O notation. The precise version of this result is now deferred to Appendix B (cf. Theorem B.1).

W2: I have the same concern about Theorem 4.4.

Response: Similarly, we have simplified the statement of Theorem 4.4, deferring the original version to Appendix B (cf. Theorem B.2).

W3 (Minor): The proof of Theorem 3.1 in the Appendix could also be more clear. For example mentioning Weyl's inequality in the first inequality of (21).

Response: Thank you for pointing this out, we have made the intermediate steps of the proof more clear in the revision.

W4 (Minor): In the proof of Theorem 5.2 at the Appendix (line 1264) shouldn’t it be $\kappa(W_L)=\kappa(W_{L:L_1+1})^{\frac{1}{L_2}}$ not $\kappa(W_L)=\kappa(W_{L:L_1+1})^{\frac{1}{L_1}}$

Response: Thanks for noticing this typo. We have corrected it in the revision.

W5 (Minor): Same concern about $L_1$ vs $L_2$ in the statement of Theorem 5.2.

Response: Thanks for noticing this typo. We have corrected it in the revision.

Q1: What exactly is the role of these additional linear layers? Are they only required for proving NC2 and NC3?

Response: We also need them for our proof of NC1, although there we only require two consecutive linear layers. In particular we need the balancedness property to ensure that the features $Z_{L-1}$ are approximately in the row space of $W_L$ so that the formula $Z_{L-1}=W_L^+Z_L$ approximately holds.

Q2: Do these results suggest that neural networks with 1 non-layer and many linear layers can also exhibit neural collapse?

Response: Yes, our results imply that neural networks with 1 non-linear layer and many linear layers provably exhibit neural collapse.

We thank the reviewer for appreciating the novelty of our work and our theoretical analysis, as well as for the detailed comments. We reply to each of the points raised in the review below. We have also uploaded a revised version of the manuscript highlighting in red the main changes.

W1: Accessibility of the paper for a '15 mins' reader. The main results are formulated using (Abstract, l19) 'balancedness', (Abstract, l20) 'bounded conditioning' and other only later defined phrases and makes it hard to assess the attractivity/usefullness of paper unless reading it whole. It is recommended to rework (the Abstract at least, if no Conclusions are present) to make it clear and self-sustained.

Response: This is a great point and we have revised accordingly. In particular, we have edited the abstract defining balancedness and bounded conditioning, so that it is now self-contained, see l. 21-24 of the revision. We have also added a final section titled ‘Discussion and Concluding Remarks’, where we start by presenting the main message of the paper, then provide a discussion and conclude with future directions.

W2: Quite a few assumptions are required in general. More over thay are added along the way, e.g., Assumption 4.2, $|\sigma(x)|\le |x|$, etc., (which is ok, but also narrows the applicability of the final result). Some of those are very technical (especialy those in Theorems, such as in Theorem 3.1, (2),(3), (4)) but as well Assumption 4.3, Theorem 4.4. (10) and more) and with paper specific notation to learn, e.g., $\lambda_{3\to L}$. It would help paper to have an thorough discussion on applicability/limitations of these assumptions. Perhaps at expense of shortening some 'proof sketches' refering them to SM?

Response: We agree with this and have revised accordingly. A detailed discussion on the assumptions is contained in the new Section 7 of the revision. In particular, we note that Theorem 3.1 provides a connection between neural collapse and properties of well-trained networks, i.e., approximate interpolation, approximate balancedness, bounded representations/weights and good conditioning of the features. As such, the result is rather general and requires no assumptions beyond the aforementioned properties. Theorem 4.4 then instantiates our framework for proving neural collapse to a class of networks with pyramidal topology (Assumption 4.1), smooth activations (Assumption 4.2) and for a class of initializations (Assumption 4.3). These assumptions are used only for the analysis of the first phase of the training dynamics, where the network achieves approximate interpolation. Thus, they could be replaced by any other set of assumptions guaranteeing that gradient descent reaches small training loss. Specifically, such guarantees are obtained by Zou & Gu (2019) for deep ReLU networks (with stronger requirements on over-parameterization but no assumptions on the topology) and by Bombari et al. (2022) for networks with minimum over-parameterization (under a requirement on the topology milder than Assumption 4.1). As concerns Assumption 4.3 on the initialization, we discuss in page 5 a setting where it holds. In addition, by following the argument in Appendix C of Nguyen & mondelli (2020), one readily obtains that Assumption 4.3 also holds for the widely used LeCun's initialization, i.e., $W_\ell^0$ has i.i.d. Gaussian entries with variance $1/n_{\ell-1}$ for all $\ell\in [L]$, as long as $n_1=\Omega(N)$.

To make space and for the sake of readability, we have simplified the statements of Theorems 3.1 and 4.4 using a big-O notation, as suggested by Reviewer 1hwn.

W3: 'Discussion and Conclusion' sections are not presented. This is most likely due to space constrained and will be added in case of acceptance (?). Yet, I find it impacts the paper quality negatively. Especially in a light of the previous (W2) comments, the discussion and conclusions could have brought a critical 'back to reality' summary. Idealy it would bring more intuition for results and their application. Yet, I find it very important to have such Discussion reviewed before publishing …

Response: This is an excellent suggestion, and we have now added the new Section 7 titled “Discussion and Concluding Remarks”. There, we discuss our assumptions, provide a comparison with earlier work by Nguyen & Mondelli (2020), discuss connections with the NTK regime and conclude with future directions.

Thank authors for a detailed revision and addressing raised concerns sufficiently. Especially discussion and clarifications (NTK + large learning rate regime as per original review) added to the paper are appreciated. Overall, amendments are significant and I raise my score to 6.

We thank the reviewer for appreciating our work. We answer the question below.

Q: In Figure 3 the authors' numerical results show that non-linear layers are increasingly linear, as the depth of the non-linear part increases. Could the authors provide more insights into how this observation relates to their theoretical results and the mechanisms driving this increased linearity?

Response: This experimental observation serves as an empirical justification for the usage of linear layers at the end of the network. The more linear the non-linear layers are, the better justified is our usage of linear layers not only as a theoretical construction, but also as a phenomenon that is reasonable in practice.

As for the mechanisms driving this, our intuition is that once the network extracts all the relevant features of the training data, it is best for it to just carry these features all the way to the end of the network. This seems to minimize the total $\ell_2$-norm of the weight matrices. For more intuition on this topic please also see [1], where the author analyzes this emergence of linearity in detail.

[1] Jacot, Arthur. "Bottleneck structure in learned features: Low-dimension vs regularity tradeoff." Advances in Neural Information Processing Systems 36 (2023): 23607-23629.