Summary
This paper studies deep neural collapse (DNC) in deep neural networks (DNN) through the prism of the neural feature ansatz (NFA) and deep recursive feature machines (RFM). It is comprised of several results:
- empirical evidence that DNC occurs in deep RFMs,
- a theoretical analysis of DNC in a high-dimensional RFM setting,
- a theoretical analysis of DNC in a kernel learning setting,
- empirical evidence that the mechanisms which lead to DNC in RFMs and traditional DNNs are the same.
Strengths
This paper shows that deep neural collapse occurs in a similar way in deep networks and deep recursive feature machines. It thus provides a simplified setting in which to investigate deep neural collapse, which is an important research direction to further our understanding of deep learning. Specifically, it shows that neural collapse can be obtained just by iterating linear regression problems, without backpropagating through a deep network.
Weaknesses
My main issue with the paper is its writing, which makes it quite difficult to read.
- The notations could be improved in several places throughout the paper (see minor points below).
- I could not follow most of section 4.2, despite being rather familiar with kernel methods and their behavior in high dimensions.
On a high level, I don't understand how a linear kernel could be the best setting for neural collapse. The text contradicts itself, as it simultaneously state that "if [$\lambda_k = 0$] [...], collapse will occur in just one layer] , but also that "this theory offers an explanation for why non-linear activation is needed". A linear layer can collapse within-class variability but also typically collapses class means together, and thus cannot induce neural collapse (see paragraph below).
On a technical level, $k_\Phi$ and $\lambda_\Phi$ are referred to before being defined, and I do not understand the roles played by $k$/$\lambda_k$ vs $k_\Phi$/$\lambda_\Phi$. Assumption 4.2 is also referred to before being stated.
- Section 4.3 is also slightly difficult to read.
I took me several tries to guess that $k_M(x,x') = \tilde k_M(x,x')\mathrm{Id}$, which should appear in the paper. The terms "input-level" and "output-dimension level" kernels should be introduced for non-specialists in multi-task kernel learning.
I also do not understand the point of introducing $M$ if it is dropped afterwards. Theorem 4.4 could simply be stated as "the optimal feature map for ridge regression is the one which already predicts the label: $\Phi(x) = y$". This result is not very surprising, and is not very integrated in the paper. I suppose that it is some kernel-level equivalent of the unstructured feature model, and suggests that weight decay might be instrumental in bringing about neural collapse? The normalization of $k$ should be restated in the definition of Problem 3 (otherwise the optimal loss is obtained when $k \to 0$).
- The message of section 5 could be presented more clearly. What I understood was that it argues that RFMs and DNNs achieve neural collapse through the same means. I suggest making this point before introducing RFMs (in particular, stating the NFA correlations). I also did not understand why this mechanism is referred to as "denoising".
My second issue is that I was not convinced by the claim that it is the right singular vectors and singular values which lead to neural collapse. By the same logic as lines 309-315, the right singular vectors do not change the DNC1 metric (with a "full" SVD where $U$ and $V$ are invertible). Similarly, if I were to divide operations in the network as $V^T\sigma$ and $US$ as opposed to $\sigma U$ and $SV^T$, I should see that it is now $US$ which is responsible for neural collapse (again with a full SVD). This conclusion also depends on the chosen metric for evaluating collapse. Why do the authors consider the ratios of traces of between- and within-class covariances, rather than the trace of their ratio (the Fisher linear discriminant)? It seems that it would reverse one of the conclusions of the analysis, since the trace of the Fisher discriminant ratio $\mathrm{tr}(\Sigma_W^{-1} \Sigma_B)$ is invariant to invertible linear transformations, and decreases under non-invertible linear transformations, so can only be improved through the non-linearity. If the conclusion of which part of the network is responsible for DNC depends strongly on the chosen metric, can we really ascribe meaning to the question? It seems to me that it is really the sequence of weights and non-linearity which _together_ induce DNC, and trying to separate their effects is not really possible.
Finally, Proposition A.1 was first published by Cho and Saul in _Kernel Methods for Deep Learning_, NIPS 2009. Besides, the expression of the kernel in eq. (5) can be simplified with algebra and trigonometry (compare with their eq. (6)).
Minor notes and suggestions:
- I suggest using a so-called "diverging" colormap (such as "bwr") in Figure 1 to clearly separate positive from negative correlations, and use the same range for both datasets.
- I suggest replacing "Gram matrix" with "(uncentered) covariance" to refer to $W^TW$, as weight matrices $W$ are generally decomposed in rows which correspond to individual neurons.
- The notation $||X||$ to refer to the vector in $\mathbb R^N$ of column norms of a matrix $X \in \mathbb R^{d\times N}$ is never introduced (and clashes with the usual convention that this is a matrix norm).
- Why is the last layer denoted $W_{L+1}$ instead of $m_{L+1}$?
- The choice of layer-indexing is confusing and seems inconsistent throughout the paper. Contrarily to what is stated in section 3.1, isn't $X_l$ the features after $l-1$ network layers? I suggest to denote the input as $X_0$ instead of $X_1$ to simplify the notations. Also, it seems that $M_l^{1/2} X_l$ should be referred to as $\tilde X_{l+1}$ rather than $\tilde X_l$ given the chosen conventions.
- Typo: missing a norm in the definition of $\bar H_l$ line 128.
-In section 4.2, I suggest defining activations before the kernels, e.g., $\tilde X_{l+1} = \kappa^{-1/2} M_l^{1/2} X_l$ and $X_{l+1} = \Phi_{\rm lin}(\tilde X_{l+1})$. I also suggest choosing a different notation for $k_{\rm lin}$ and $\Phi_{\rm lin}$ which are confusing as they imply linearity, and to avoid the awkward "non-linear feature map $\Phi_{\rm lin}$".
- Typo line 260: the output space of $k_M$ should be $\mathcal R^{C\times C}$.
- I suppose that $\lambda = \mu$ in section 4.3.
- In the caption of Figure 2, I suppose that "fully-connected" should be removed in the case of ResNet.
Questions
- Are the feature maps $\Phi_l$ and kernels $k_l$ unrestricted in Algorithm 1, or do they have to match in the sense that $k_l(x,x') = \langle \Phi_l(x), \Phi_l(x')\rangle$?
- What is the motivation behind considering _normalized_ features _before_ the non-linearity in section 4.1? Could the authors clarify the role of these non-standard conventions?
- Why are the setting different between sections 4.1 and 4.2? (relative to normalizations). Is neural collapse still empirically observed with the choices of section 4.2? It raises the suspicion that it could not the case, in which case Theorem 4.3 would not really explain the results of section 4.1 (e.g., because the asymptotic regime is not reached in practice).
Limitations
See weaknesses.
In its current state, I think that the paper is slightly below the acceptance bar and would require minor, if not major, changes before it can be fully appreciated by the NeurIPS community. I would be happy to raise my score if the authors address the points raised above.