The Prevalence of Neural Collapse in Neural Multivariate Regression

We would like to thank you for the detailed review and helpful comments as well as for recognizing that our work “derives results similar to those of traditional neural collapse”.

We acknowledge that the analysis for NRC bears some similarities with traditional NC for classification tasks with MSE loss. In our proofs, we leveraged some existing results from UFM theory for MSE loss. However, the regression problem is substantially different from the MSE classification problem, resulting in a very different definition of neural collapse. In classification with balanced data, neural collapse refers to the convergence of the last layer features and weight matrix to an ETF geometric structure. In contrast, in multivariate regression, the feature embeddings converge to an n-dimensional subspace spanned by the principal components of the features, which converge to the subspace spanned by the row vectors of the weight matrix $W$. Moreover, in regression tasks, the weight matrix satisfies distinct properties that depend on the covariance matrix of the target variables and their eigenvalues. Therefore, the qualitative nature of the final results is substantially different from those for classification. In addition, the NRC properties provide valuable insights for model training in regression. As demonstrated below, we also present training with a fixed $W$ based on the neural collapse solution, which improves training efficiency without compromising performance. We anticipate that further exploration will reveal additional implications, which we leave for future work.

Questions:

• Indeed, the weights $W$ and features $H$ in the final regression layer have desirable properties, as highlighted by Theorem 4.1. Consequently, instead of learning the final regression layer, we can design the weight matrix $W$ as a random matrix such that $WW^{\top}=\Sigma^{1/2}$ as suggested by Theorem 4.1. By training the model with the weights $W$ kept frozen, we can significantly reduce the number of parameters and the computational complexity of training, while still achieving comparable performance in terms of both training/testing MSE and NRC metrics, as illustrated in Figure D.

• This paper focuses on deep regression with a final linear layer, where the network is highly overparameterized and capable of learning a wide range of representations. In lines 217-221, we discussed the resemblance of the UFM with standard (multivariate) linear regression. On a high level, one can regard the features $h_i$, $i=1,...,M$, as the inputs to linear regression. Then, the objective becomes the minimization of the squared loss between the predicted outputs $Wh_i+b$ and the targets $y_i$. In standard regression, the $h_i$’s are fixed inputs, whereas in the UFM, we are optimizing over the weights $W$, the biases $b$ and most importantly over the “inputs” $H$. As for GP regression, this is indeed an interesting research direction, but one that is certainly very rich and beyond the scope of the current paper.

• The norms of $W$ and $H$ have different training dynamics under classification and regression tasks. In classification tasks with CE loss and no regularization, once features can be perfectly separated (during TPT), the training objective will always decrease if we fix the direction of H and W and only increase their norms (similar to temperature scaling), with the optimum attainable only at infinity. However, with MSE loss without regularization, the loss function increases towards infinity as the norm of the predictor Wh_i approaches infinity. The MSE loss function therefore forces the norms of the predictors to be finite on its own without the need of regularization. The actual optimal solutions for W and H are characterized in Theorem 4.3 (no regularization case) which shows that there are actually an infinite number of finite-norm solutions.

• In Figure 5, the train MSE decreases when decreasing the regularization constants, which mirrors the decrease of the train MSE in Figure 4, when decreasing $\lambda_{WD}$. This observation aligns with modern machine learning practices, where training MSE tends to be lower with no or weaker regularization. Moreover, Figure 5 exhibits a phase change in NRC, with NRC being more pronounced with stronger or higher regularization constants. Following your suggestion, we have added the experimental results corresponding to no regularization to Figure 5; see Figure B in the rebuttal PDF for these additional results.

• Figure 4/A has been updated to include experiments on all six datasets, providing a comprehensive validation of the impact of regularization. We have also updated Figure B to include the Reacher, Swimmer, Hopper, and CARLA2D datasets. Due to space limitations, CARLA1D and UTKFace are not included in Figure B and will be added in the future revision.

We would like to thank you for the detailed review and helpful comments. We also thank you for recognizing NRC as a “new form of Neural Collapse observed in multivariate regression tasks” and that “this study extends Neural Collapse to regression, suggesting a universal deep learning behavior”.

Weaknesses:

Indeed, in line 601, in the matrix that you quoted, multiplication of $\sqrt{Mc}$ by $\mathbf{I}_n$ is needed. We have spotted this typo in our manuscript and corrected it.

Questions:

Finally, let us now elaborate on the derivation of Eq. (24) and Eq. (25) in the Appendix. The proof of Theorem 4.1 leverages [Zhou et al., 2022a, Lemma B.1]. For clarity, we have restated their lemma in our notation, see Lemma D.1 in Appendix D. In order to derive $WH$, Eq. (23) in our work, Zhu et al., first showed that, if $w_i$ and $h^i$ denote the $i$-th column of $W$ and $H^T$ respectively, then, when $w_i\neq 0$ and $h^i\neq 0$ (if this is not the case, then $w_i=h^i=0$), we have that

with $\sigma_i\ge \sqrt{M c}$, for all $i=1,...,n$, where $\sigma_i$ is the $i$-th singular value of $\tilde{Y}$ and $u_i$, $v_i$ are the corresponding left and right singular vectors. Rewriting the equation above (see Eq. (22) in [Zhou et al., 2022a, Lemma B.1]) in matrix notation readily yields Eq. (24) and Eq. (25) in the Appendix. We will include this clarification in the Appendix.

We thank you for your very detailed and insightful review, and for all the positive comments you made regarding our work. We also fully agree that the presentation of the experimental results can be improved. During the rebuttal week, we have worked very hard, running additional experiments and reorganizing our figures to respond to your comments. We believe our experiments are now “comprehensive and consistent”. We promise to include all additional results below in the future revision.

We ran all experiments with 3 random seeds and plot the variance across seeds as the shaded area. As shown in all figures, there is little change across seeds, confirming that NRC consistently emerges in regression.

Concerning your question about 𝜸 in Figure 2: We ran all experiments for long enough to ensure the training enters TPT measured by $\mathbf{R^2}$ (See Figure A). After training, we extracted the $W$ matrix and identified 𝜸 that minimizes NRC3 for that specific $W$. This 𝜸 was then used to compute the NRC3 metric for all $W$ matrices during training, resulting in the NRC3 curves shown in Figure 2. Figure 3 visually shows NRC3 as a function of 𝜸 for the final trained $W$. Mathematically, we can show that if $\lambda_{WD}$ is sufficiently large, NRC3(𝜸), as given in the definition of NRC3, is convex and has a unique minimum. Empirically, our experiments show that 𝜸 w.r.t. minimum NRC is consistent across different seeds.

We very much liked your suggestion to examine the explained variance ratio. The results are striking as shown in Figure C: for all datasets, there is significant variance for the first $n$ components; for other components, there is very low or even no variance. We also examine different definitions of NRC1 with $n$ varying from 1 to 4 below to justify the choice of $n$ being equal to the target dimension.

Concerning your question about converting a univariate regression problem to a classification problem, and vice-versa. We believe this is an interesting direction for future work. Of course, if $n$ is large, it would become difficult to convert from regression due to quantization and the resulting classification dataset would likely be highly imbalanced. In classification, as you state, the dataset can be balanced or imbalanced, and there are NC studies for both. If we were to convert a univariate regression dataset to a classification dataset using quantization, then the dataset would be balanced if there are the same number of $y_i$ values in each quantile. As this property is likely unrealistic for most datasets, then we can conclude neural regression has a stronger connection to classification with imbalanced datasets.

We believe that the UFM theory provided in the paper helps to explain many properties of neural regression. In terms of efficiency, experiments in Figure D show that we can fix W to a random matrix such that $WW^{\top}=\Sigma^{1/2}$ as suggested by Theorem 4.1, and use this fixed matrix throughout training. This approach significantly reduces the number of weights to be optimized and results in similar performance and NRC metrics, paralleling a well-known approach in NC for classification.

The number of features in the penultimate layer $d$ is typically 256 or larger for deep learning. In multivariate regression, the target dimension $n$ is typically small in comparison. For example, in most MuJuCo environments, often used to study imitation and reinforcement learning, satisfy $n \leq 6$.

We agree that “NC is…a fundamental property of multivariate regression” is a bit of an overstatement. We will rewrite this to say “often occurs in neural multivariate regression”.

We ignored the plots for 𝜸 and NRC3 for univariate regression with $n=1$. For $n=1$, $\lambda_{max} = \sigma^2$, the variance of the 1D targets. This is just a scalar value for each of the $n=1$ datasets and does not seem to lead to an insightful plot.

Concerning your question about which NRC metrics fail when regression parameters are zero or very small. Figure A and B show when weight decay approaches zero, NRC1-3 typically become larger, compared with NRC1-3 obtained with larger weight decay values.

Particularly, when there is no weight decay, we run more training epochs to verify the asymptotic behavior of NRC1-3 and test MSE in Figure A. We observe that NRC1-3 has a strong tendency to converge (There is a relatively small amount of collapse since gradient descent tends to seek solutions with small norms.), while the test MSE increases on small MuJoCo datasets. Theorem 4.3 provides some insight: when there is no regularization, there is an infinite number of non-collapsed optimal solutions under UFM; whereas Theorem 4.1 shows that when there is regularization, all solutions are collapsed. When there is regularization, we are seeking a small norm optimal solution, which leads to NRC1-3.

The main assumption in the UFM model – for both classification and regression – is that the neural network is capable of mapping any set of training inputs to any set of feature vectors. Based on this assumption, the UFM model leads to a new optimization problem for classification and regression. We believe the UFM model for regression makes the same logical connection that is made for classification.

We have included UFM experiments with no regularizer in Figure B. Similarly as in Figure A, NRC1-3 do not converge to very low values as is when there is regularization. As indicated in section 4.5, Figure 5/B is intended to validate the results under the UFM framework. To align with the UFM assumption, we remove ReLU in the penultimate layer such that the learned feature can be any real number instead of being restricted to positive values.

Thank you for your comment. The original paper used only one seed. In response to your review, we re-did all the experiments over all 6 datasets with 3 seeds, and the results are shown in the figures of the rebuttal. The shaded regions are generally narrow, showing little change across seeds for most (but not all) of the datasets.

We will include Figure C (and all of the other updated figures in the rebuttal) if accepted. Thanks again for your excellent review.

We would like to thank the reviewer for the detailed review and helpful comments. We also thank you for recognizing that the “study is novel,” “provides a new understanding of neural multivariate regression,” and that the “theoretical results and experimental results are solid and well-organized”.

We agree, however, that we didn’t sufficiently emphasize the significance and impact of the results. Indeed, neural collapse for the classification problem under MSE loss has been extensively studied. We will update our paper to fully describe our contribution with respect to earlier works on MSE loss for classification. Indeed in our proofs, we leveraged important existing results from the UFM theory for MSE loss. But the regression problem remains substantially different from the MSE classification problem, a difference which is highlighted by a very different definition of neural collapse. Classification with a balanced data set gives rise to an ETF geometric structure, whereas multivariate regression gives rise to a distinctly different structure defined by our NRC1-3 definitions which involve subspaces spanned by principal components and the covariance matrix of the targets. This new definition requires an entirely new empirical analysis on entirely different datasets. And although the mathematical UFM derivations for regression mirror those for classification, the qualitative nature of the final results - involving the covariance matrix of the targets and their eigenvalues - is substantially different than those for classification.

In terms of impact, the NRC properties offer valuable insights into model training for regression tasks. For example, we can fix the last layer to neural collapse solutions by setting the weight matrix $W$ to be a random matrix which satisfies $WW^{\top}=\Sigma^{1/2}$ as suggested by Theorem 4.1. By training the model with the weights $W$ kept frozen, we can significantly reduce both the number of parameters and the computational burden, while still maintaining comparable performance in terms of training/testing MSE and NRC metrics, as shown in Figure D. We will include this figure and the corresponding discussion in the revised version. Thank you for your suggestion. We believe that the empirical and theoretical results in this paper will guide the future design of multivariate regression problems.

Thank you for your suggestion to investigate R2 during the terminal phase of training. In Figure A of the supporting PDF we provide the $\mathbf{R^2}$ results, which confirm the network has entered the terminal phase of training.

We omitted regularization of the bias term for the UFM for two reasons. First, in practice, it rarely improves performance. Second, it allows for a cleaner presentation of the theorem statements and proofs. It could be easily added without changing the main takeaway points.

Thank you for your interesting comment about no regularization. We agree that gradient descent often exhibits an implicit bias towards a minimum norm solution. We have added the curve corresponding to $\lambda_{H}=\lambda_{W} = 0$ in Figure B, which indeed shows to be the case, with the NRC metric typically decreasing as training progresses. This is consistent with the observation in the layer-peeled paper, as you indicated. We will highlight this in the revision.