Neural Collapse Beyond the Unconstrained Features Model: Landscape, Dynamics, and Generalization in the Mean-Field Regime

Thank you for the valuable and detailed comments, we address concerns below.

1. Supplementary materials.

We will add this in the revision.

2. NC not surprising under MSE loss.

We politely disagree with the claim that “it is quite intuitive that the features concentrate at their class mean when the train loss is small under MSE loss.” It is true that small MSE loss implies that the output of the neural network is roughly the same for all data points in the same class. However, this is only a necessary but not sufficient condition for NC1 to occur. The technical subtlety is that, since the feature dimension is larger than the number of classes, there might be components of the features in the space orthogonal to the last linear layer. Without the balancedness condition, it is not obvious to show that these orthogonal components vanish at stationary points. Intuitively, these components make the weights more isotropic, thus they could potentially decrease the entropic penalty, but increase the L2 penalty, meaning that the tradeoff needs to be carefully considered. Thus, our results in Section 4.2 require a suitable characterization of the learned feature and careful manipulations to show that the orthogonal components are indeed vanishing with training loss and gradient norm.

3. Extension to CE loss.

Please see point 1 of our response to Reviewer g3N9.

4. Measuring rate of convergence.

We provide a convergence rate for the loss in Equation (24). This is already non-trivial given the non-convex nature of the problem and to our best knowledge, no explicit rate of convergence is known in the mean-field regime beyond the 2-layer case (which has a convex free energy). Characterizing the convergence rate of the gradient norm is harder, as it requires controlling the curvature along the whole trajectory.

5. Linear separability assumption in Theorem 5.3.

In general, it is not true that NC1 and vanishing test error co-occur under MSE loss without any assumptions on the data distribution. In fact, the occurrence of NC1 implies that the model overfits the data, and overfitting is not always benign without additional assumptions: classic benign overfitting results for simple models such as linear regression require a decay in the data spectrum and proper alignment between data spectrum and signal. In this sense, showing the co-occurrence of NC1 and vanishing test error may be regarded as a harder problem than benign overfitting, and the latter phenomenon has not been investigated in the mean-field regime to our best knowledge.

6. Comparison of our work to [1].

First, [1] operates in the NTK regime where the behavior of the weights is fundamentally different from the mean-field regime that we consider. More precisely, the proof of Theorem 4.4 in [1] distinguishes two phases in training: in the first phase, the loss becomes small and no feature learning takes place; in the second phase, the weight matrices in the linear part of the network become balanced and NC1 occurs. In contrast, our work operates in the mean-field regime where feature learning takes place from the beginning and, therefore, there is no need for a second phase of the dynamics: NC1 occurs as soon as the loss and its gradient are small. Our experiments on ResNet-18 and VGG-11 also indicate this relation (see point 1 of our response to Reviewer AaqF for more details).

Second, the balancedness condition in [1] is not just a technical requirement, but the key reason for the occurrence of NC1. In fact, the balancedness condition rules out the possibility of having any components of the features in the orthogonal space to the last linear layer. However, our results both theoretical (see Lemma 4.5) and experimental (see the end of Section 4.2) conclusively demonstrate that NC1 holds even if the network is not balanced.

7. Impact of weight decay.

Firstly, non-zero weight decay is necessary to ensure non-vanishing $W$ and features at all stationary points. Secondly, the weight decay affects the order of the empirical loss we converge to: when $\lambda_\rho = \beta^{-1}$, the empirical loss at convergence is of order $\beta^{-1}$. This is reflected in Corollary 4.9, as better NC1 guarantees (smaller $\delta_0$) require smaller weight decay $\beta^{-1} \leq poly(\delta_0).$

8. Extension to deeper/more complex networks.

We expect that our results can be extended to a deep network with many linear layers, see point 4 of our response to Reviewer AaqF. In contrast, analyzing deep networks with many nonlinear layers would require fundamentally different techniques and we regard it as an exciting future direction. Towards this goal, two possible strategies are (i) to characterize features at an $\epsilon_S$-stationary point, or (ii) to directly consider the evolution of features during training.

[1] Jacot, Sukenik, Wang, Mondelli, Wide neural networks trained with weight decay provably exhibit neural collapse.

Thank you for the detailed review. We address concerns below.

1. Lack of soundness in experimental design.

We followed the suggestions and did additional experiments, please download pdf in https://github.com/conferenceanonymous152/icml25

(a) For three-layer networks, Figure 3 shows the linear relation between log of NC1 and log of gradient norm/training loss throughout training. Different points correspond to different epochs of training (purple ones in the top-right correspond to early stages, and yellow ones in the bottom-left correspond to late stages), and we average the results of 4 independent experiments. The plots indicate a polynomial relation between the quantities as predicted by Equation (17).

(b) For practical models, we show that ResNet-18 and VGG-11 trained on CIFAR-10 exhibit a similar behavior as three-layer networks. Figure 4 is similar to Figure 1 in the paper. The scatterplots in Figures 5, 6 plot the $\log-\log$ relation between NC1 and gradient norm/training loss, and they indicate a polynomial relation in the terminal phase of training (while the early phase of training may be a little different).

2. NC2-NC3.

To prove NC2-NC3, [1] uses that the number of linear layers is large. If we stack many linear layers on top of the feature-learning layer and run gradient flow, we expect that NC2-NC3, as well as NC1, can be shown. As for NC1, our approach could be adapted as follows. Let the network be $f(x; \rho, W_2, \dots, W_L) = W_L \dots W_2 h_\rho(x),$ and define $W^\top =W_L \dots W_2$. Then, the characterization in Theorem 4.2 still holds, since the condition of Definition 4.1 only involves the gradient computed over $\rho$. The lower bound in Lemma 4.3 holds and the upper bound becomes $\sigma_{max}(W) \leq O( (\log \beta)^L ) = o(\beta)$, which implies Corollary 4.4. Theorem 4.8 can be proved by letting $t_*$ be the minimal first hitting time for $W_2, \dots, W_L$ and $\rho$ and following a similar strategy. Finally, since approximate NC1 holds for $H_\rho(x),$ it also holds for $W^\top H_\rho(x)$, for any $W$ satisfying Lemma 4.3.

As for NC2-NC3, we note that $W_2, \dots, W_L$ are balanced for all stationary points $(\rho, W_2, \dots, W_L)$. Thus, if the condition number of $W$ is bounded by a constant independent of $L$ (as in Theorem 3.1 of [1]), one can show NC2-NC3 following the approach in [1]. We note that $W_1 := \mathbb{E}_\rho[a a^\top]$ is not balanced with $W_2, \dots, W_L$ due to Lemma 4.5 of our work. Thus, to show NC1 for a network with many linear layers, the results of [1] cannot be applied.

3. Major differences compared to previous work, additional entropic regularization term and question about line 112.

The entropic regularization term is not implicitly induced by SGD, but it is explicitly added in the form of a Brownian noise. In fact, we consider noisy gradient flow for training, i.e., mean-field Langevin dynamics. This is standard in the mean-field regime [2, 3], and noisy gradient flow is the limiting dynamics of noisy SGD with step size $\eta$ and any batch size (including full-batch) on networks of width-$N$ as $\eta \to 0$ and $N \to \infty$ [3].

What distinguishes our work is that we show NC1 for a network that (i) is trained end-to-end and (ii) exhibits feature learning throughout training. In contrast, previous work focuses either on UFM (which does not capture the effect of data) or on a dynamics reaching near-interpolation in the NTK regime [1] (which does not exhibit feature learning [4]). Furthermore, the occurrence of NC1 in [1] relies on the balancedness of all linear layers, which may not hold, see Lemma 4.5. Instead, our result only needs small training loss and gradient norm, which is commonly satisfied as long as the training dynamics converges. Finally, we show for the first time the co-occurrence of NC1 and vanishing test error for a class of data distributions. For additional comparison of our work to [1], please see point 6 of the response to Reviewer sgNN.

4. Connection between theoretical results and real-world networks.

See point 1.

5. Line 116.

$Law(a,u)$ means the joint measure of random variables $a,u.$

6. Assumptions (A3).

ReLU does not satisfy (A3), but we still expect our results to hold by taking the limit of a sequence of approximations to ReLU. We note that (A3) is purely technical (it ensures the well-posedness of the Wasserstein gradient flow) and standard in works [2, 3] considering the mean-field regime.

[1] Jacot, Sukenik, Wang, Mondelli, Wide neural networks trained with weight decay provably exhibit neural collapse.

[2] Suzuki, Wu, Nitanda, Mean-field langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction.

[3] Mei, Montanari, Nguyen, A mean field view of the landscape of two-layer neural networks.

[4] Yang, Hu, Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks.

Thank you for your valuable review and positive feedback. We address questions and concerns below.

1. Difficulty in extending results to CE loss, and expected qualitative results.

Thank you for the question. It is indeed interesting to extend our results to CE loss and we summarize two major difficulties below:

• The intuitive difference between MSE loss and CE loss is that small MSE loss implies that the output of the neural network is roughly the same for all data points in the same class, as the label of data points in the same class is the same. In our model, we stack a linear layer after the feature learning part, which implies that data points in the same class having approximately the same output is a necessary but not sufficient condition for NC1 to happen. To see this, suppose $h(x_1), h(x_2)$ are the features of two data points in the same class. Then, ||h(x_1)-h(x_2)||\_2^2 \geq \frac{1}{\lambda_\max(W^\top W)}||W^\top h(x_1)-W^\top h(x_2)||_2^2. Now, in order to extend our results to CE loss, one would have to rule out the possibility that different data points in the same class could have different logits, and it is unclear how to do so in our setting even when the loss is small. However, we believe that this is just a technical issue coming from the fact that we consider a linear classifier in the last layer.

• Another technical issue is that our results require the characterization of features in Equation (6). This only holds for MSE loss, and a different characterization would be needed for CE loss.

That being said, we do expect a similar characterization to hold for CE loss. As a first step towards handling the CE loss, we note that there are indeed papers in the mean-field regime training with CE loss [1,2] or training with general losses that include CE [3]. However, all those papers consider two-layer networks, [1] requires no regularization terms, and [2] focuses on the specific problem of learning parities.

2. The technical reason for considering the approximated model truncating the second layer in Section 5, and whether the following conclusions would be affected if the full model was considered.

Having an approximated second layer implies that there exists a unique global optimum for the Gibbs minimizer as in Equation (27) and such a global optimum is achievable by mean-field langevin dynamics [3]. Specifically, one needs that $\nabla_\theta \frac{\delta}{\delta \rho} \mathcal{L}_n(\rho)(\theta)$ is Lipschitz in $\theta, \rho$ (cf. Proposition 2.5 in [4]), and without truncation this cannot be guaranteed. However, we note that our results in Section 5 hold uniformly for large enough $R$, which means that we also expect the same conclusion for the original model which corresponds to $R = +\infty.$ Rigorously speaking, one cannot interchange the order of the limits (in the time $t$ and the truncation parameter $R$), but we believe this to be a purely technical issue.

3. Typo in the title.

Thanks for spotting this, we will correct it in the revision.

[1] Chizat, Bach, Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss.

[2] Suzuki, Wu, Oko, Nitanda, Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond.

[3] Suzuki, Wu, Nitanda, A. Mean-field langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction.

[4] Hu, Ren, Siska, Szpruch, Mean-field langevin dynamics and energy landscape of neural networks.

Thank you very much for finding that our paper advances the study of neural collapse and recommending a strong acceptance! We are happy to have a further discussion in case additional questions or comments come up.