Average gradient outer product as a mechanism for deep neural collapse

Thank you for your review. We address all your concerns below:

Given that the main results apply both to a non-standard kernel method and a non-standard training algorithm [...]

Our first motivation for studying DNC with Deep RFM is that, unlike the standard analytical approaches for neural collapse, the Deep RFM procedure is highly dependent on both the input samples and their labels. Specifically, the collapse process of Deep RFM depends on mapping with the AGOP of a predictor trained to fit these input/label pairs. Moreover, as we demonstrate in Theorem 4.3, the rate of this collapse throughout the layers is highly dependent on the specific structure of the data through its relationship to the parameters $\lambda_\Phi$ and $\lambda_{\mathrm{lin}}$. We also provide significant evidence that this AGOP mechanism is present in deep neural networks in Section 5 of our original manuscript, indicating that our results apply to more practical architectures.

Secondly, the ability of Deep RFM to recover interesting behaviors of deep learning has broad implications for our understanding of neural networks. In fact, the same Deep RFM model we study has recovered a number of other interesting deep learning phenomena including, e.g., grokking modular arithmetic [1], convolutional feature learning [2], and even low-rank matrix completion [3]. As you mention, Deep RFM is a simpler model than neural networks, in the sense that we understand its individual components: (1) a kernel regression step, (2) mapping with the AGOP, and (3) random feature projection. Therefore, our work additionally justifies Deep RFM as a simplified proxy to understand neural networks.

[1] Mallinar, Beaglehole, Zhu, Radhakrishnan, Pandit, Belkin, “Emergence in non-neural models: grokking modular arithmetic via average gradient outer product.” arXiv pre-print, 2024.

[2] Radhakrishnan, Beaglehole, Pandit, Belkin, "Mechanism for feature learning in neural networks and backpropagation-free machine learning models." Science, 2024.

[3] Radhakrishnan, Belkin, Drusvyatskiy, "Linear Recursive Feature Machines provably recover low-rank matrices." arXiv pre-print, 2024.

The authors assume that the gram matrix of the data is full-rank. This requires assuming that the number of datapoints is smaller than or equal to the input dimension (which subsumes assumption 4.1). Standard datasets violate this assumption.

The strong assumption that the gram matrix of the data is full-rank is needed only if one requires neural collapse in every layer of the network, starting from the very first one. In contrast, if we consider collapse starting at any given later layer of Deep RFM, then we only need that the smallest eigenvalue of the corresponding feature map is bounded away from 0. This in turn requires that the number of features at that layer is greater than the number of data points, a condition which is routinely satisfied by the overparameterized neural networks used in practice.

Are the models trained with SGD?

The neural network models are trained with standard SGD. We agree that this result is of independent interest, and we will emphasize in the main section that our training procedure is standard. Let us also clarify that Deep RFM is not trained by any gradient-based method, and is instead optimized through an iterated three-stage process: (1) solve kernel ridgeless regression, (2) compute the AGOP of the predictor from step 1, and (3) apply a high-dimensional feature map.

Thank you for your feedback. We note that we will significantly improve the presentation of our paper, and specifically Sections 4.2 and 4.3. Please see the global response for a summary of our changes. We now proceed to address individual comments.

I could not follow most of section 4.2 [...] We will clarify this. The two kernels $k_{\mathrm{lin}}$ and $k_{\Phi}$ have opposite effects. In particular, the closer $k_{\mathrm{lin}}$ is to a linear kernel, the more accelerated collapse will be for Deep RFM. In fact, a linear kernel induces NC in just one AGOP application. To see this, note that ridgeless regression with a linear kernel is exactly least-squares regression on one-hot encoded labels. In the high-dimensional setting we consider, we find a linear solution $f(x) = \beta^T x$ that interpolates the labels, and the AGOP of $f$ is $\beta \beta^T$. Since we interpolate the data, $\beta^T x = y$ for all $(x,y)$ input/label pairs and, hence, the data collapses to $\beta\beta^T x = \beta y$.

In contrast, having $k_{\Phi}$ far from a linear kernel accelerates collapse, and the non-linear activation function is needed so that $k_{\Phi}$ is significantly non-linear. In fact, if $k_{\Phi}$ were exactly linear and the original data were not linearly separable, then collapse cannot occur at any depth of Deep RFM. In that case, the predictor would need to contain a non-linear component in order to fit the labels exactly, deviating from the ideal case described above.

Section 4.3 is also slightly difficult to read [...] Thank you for your suggestions. We address your points one-by-one:

• We introduce $M$ because it is the key part of the parametrized kernel ridge regression which can be understood as the objective function of the RFM iteration. We emphasize that when the data (or the features) are full rank, dropping $M$ does come without loss of generality. Due to the character limit, we kindly refer you to our response to reviewer AtBe for details.
• We appreciate your interpretation of this model as an unconstrained kernel model. We agree with this interpretation and will call it in this way in the revision.
• Theorem 4.4 is well integrated in our paper as the kernel ridge regression loss can be understood as the quantity the RFM learning algorithm is minimizing. Thus, the result describes DNC as an implicit bias of the RFM iteration.
• For additional structural changes to Section 4.3, see our global response.

The message of section 5 [...] “Linear denoising” highlights the nature of the mechanism: by only applying linear transformations, we decrease the within-class variability of the data and improve DNC1 by only discarding class-irrelevant information (i.e., noise) via the non-trivial null-space of the weight matrix. See our global response for additional changes.

Why splitting into $\sigma U$ and $SV^T$? In the compact SVD (which is an equivalent re-writing of the full SVD), it is the matrix $V^T$ which has non-trivial null-space, while the matrix $U$ always has a trivial null-space. Note that the non-triviality of the null-space allows for the DNC1 metric to decrease. Hence, this is the correct split of the layer.

More generally, the main message of this section is not which part of the SVD decreases the DNC1 metric, but rather that the linear transformation of $W$ has a stronger effect on NC1 than the non-linearity $\sigma$. In fact, the most natural grouping of the layer - into the non-linearity alone and the entire weight matrix alone - gives equivalent NC1 values to our grouping, as the two standard metrics for NC1, $\mathrm{tr}(\Sigma_W \Sigma_B^\dagger)$ and $\mathrm{tr}(\Sigma_W)/\mathrm{tr}(\Sigma_B)$, are invariant to the rotation by the left singular vectors of $W$.

This conclusion also depends on the chosen metric... This is an interesting consideration, however we kindly disagree that Fisher linear discriminant cannot decrease due to a linear mapping. Consider the example in which the feature vectors of $X$ are not collapsed, but the differences between points of the same class lie in the null-space of $W.$ Then, $WX$ is perfectly collapsed and any NC1 metric would be identically zero.

While the two metrics are similar, we choose our metric because it directly quantifies how close the feature vectors are to their class means - a common and intuitive interpretation of NC1. Our metric is also frequent in the literature, see e.g. Rangamani et al., 2023 and Tirer et al., 2023.

Are the feature maps $\Phi_l$ and kernels $k_l$ unrestricted? The feature maps and kernels do not have to match. We will clarify this when we introduce the two objects.

What about the normalized features in 4.1? We normalize the features as neither the AGOP projection nor the random feature map have a native ability to rescale the data. Therefore, while we still see collapse in terms of the angles between data points, the dynamics of the data norms are more difficult to control in this setting.

Why are the setting different between sections 4.1 and 4.2? [...] Neural collapse would still be observed without normalizing the data, in the sense that all points of the same class will be parallel and all points of different classes will be perpendicular (or at angle $-\frac{1}{K-1}$). However, we are not guaranteed convergence of the data Gram matrix to $yy^T$ in general without the re-normalization, as the diagonal entries of $X^T X$ will not be identically $1$. In our analysis, it is sufficient to scale the Gram matrix at each layer with $\kappa^{-1}$. The main advantage of this scaling is that it is analytically much simpler than explicitly projecting the data to be norm $1$.

Minor notes. We will address all these points in the revision. As an example, we will rename $k_{\mathrm{lin}}$ and $\Phi_{\mathrm{lin}}$ to $\widehat{k}$, indicating the predictor kernel, and $\Phi_{\mathrm{map}}$, indicating the feature map applied to the data.

We thank the reviewer for their reply.

Yes, we only study neural collapse on the training set. Note that there are mixed empirical results about whether or not neural collapse transfers to the test set [1, 2, 3]. Moreover, to the best of our knowledge no work has studied the emergence of neural collapse at test time theoretically. Perfect collapse at test time would mean perfect generalization, therefore only weaker forms of test-collapse are available [3].

We emphasize that neural collapse is a particularly structured form of interpolation that is not implied by low train loss, or overfitting, alone. In fact, there are many ways either neural networks or Deep RFM can interpolate the training data. Therefore, it is remarkable that the particular interpolating solutions found by these models exhibit DNC. Our paper, as well as all the other theoretical papers on collapse, study why among all interpolations for a given model class and dataset, the training procedure selects a solution exhibiting DNC.

You are correct, the non-linearity is essential to achieve the collapse for linearly non-separable data. The ReLU plays an important role, as your proposed metric shows. We will make it clear in our revision that what we mean by saying that the linear layers are responsible for reducing within-class variability is not that they can achieve collapse on their own, but that in the solved model, it is the linear layers that directly decrease the NC1 metric. Moreover, the ReLU plays a critical role in enabling collapse, especially when the data is not linearly separable.

[1] Xu and Liu. “Quantifying the variability collapse of neural networks.” International Conference on Machine Learning, 2023.

[2] Kothapalli. “Neural collapse: A review on modelling principles and generalization.” arXiv pre-print, 2022.

[3] Hui, Belkin, Nakkiran. “Limitations of neural collapse for understanding generalization in deep learning.” arXiv pre-print, 2022.

Thank you for your review. We address your concerns below.

I found the paper challenging to read.

We will make a number of changes to our presentation. Please see our global response to all reviewers for a summarized description of these changes.

Can the authors clarify if other metrics, such as the Neural Tangent Kernel (not the limit), would effectively predict this behavior? Or AGOP is unique in this aspect?

Prior work has shown that neural collapse can occur if the empirical Neural Tangent Kernel has a particular block structure dependent on the inputs and labels [1]. In Deep RFM, we consider a different setting, where we fix a kernel function to be independent of the data, i.e. the NTK in the infinite-width limit, and then understand how a linear transformation prior to the kernel evaluation can induce DNC.

[1] Seleznova, Weitzner, Giryes, Kutyniok, Chou. “Neural (Tangent Kernel) Collapse.” NeurIPS, 2023.

What are the practical implications of this work?

While the scope of this work is theoretical, there are a number of works demonstrating the practical applications of (deep) neural collapse and Recursive Feature Machines. More specifically, neural collapse has been used as a tool in several important contexts, including generalization bounds, transfer learning, OOD detection, network compression and robustness, see lines 85-90 and the references therein. Furthermore, RFM itself has been shown to be a practical tool for tabular data (Radhakrishnan et al., Science 2024) and scientific applications [1,2].

Our work also has broad implications for our understanding of deep learning. In particular, this work is part of a large body of research demonstrating that Deep RFM is a useful proxy to explain interesting phenomena in deep networks. In addition to neural collapse considered in this work, RFM has been shown to exhibit grokking in modular arithmetic [3], recover convolutional edge detection, and learn the same features as fully-connected networks on vision tasks (Radhakrishnan et al., Science 2024). These findings point toward Deep RFM and AGOP as powerful tools to improve our understanding of neural networks, which would likely lead to significant practical implications.

[1] Aristoff, Johnson, Simpson, Webber. “The fast committor machine: Interpretable prediction with kernels.” arXiv pre-print, 2024

[2] Cai, Radhakrishnan, Uhler, “Synthetic Lethality Screening with Recursive Feature Machines.” arXiv pre-print, 2023.

[3] Mallinar, Beaglehole, Zhu, Radhakrishnan, Pandit, Belkin, “Emergence in non-neural models: grokking modular arithmetic via average gradient outer product.” arXiv pre-print, 2024.

Thank you very much for your detailed review. In our response here, we will pay significant attention to improving the writing and organization of our work, especially Section 4. Please also read our global response where we discuss the changes in presentation in detail to all reviewers. We proceed by addressing each weakness and question individually.

Section 4 is doing [...]

Thank you for this suggestion, please see our global response for the changes we will make in our revision.

The opening paragraph of Section 4.3 is unreadable.

We will present our result differently. As an example of our improved presentation, we would rewrite the opening paragraph of this section as follows:

“We have shown so far that mapping with the AGOP induces DNC in Deep RFM empirically and theoretically in an asymptotic setting. In this section, we demonstrate that DNC emerges in parameterized kernel ridge regression: in fact, DNC arises as a natural consequence of minimizing the norm of the predictor jointly over the choice of kernel function and the regression coefficients. This result proves that DNC is an implicit bias of kernel learning. We connect this result to Deep RFM by providing intuition that the kernel learned through RFM effectively minimizes the parametrized kernel ridge regression objective and, as a consequence, RFM learns a kernel matrix that is biased towards the collapsed gram matrix $yy^T$.”

Dropping the dependence on $M$ is not always equivalent to the original problem.

You are right that optimizing over all $k$ is not always equivalent to optimizing over all $M$. Thus, in general, this provides only a relaxation. However, if the gram matrix of the data is invertible (this is a realistic assumption if we consider layers $\ell>1$ of Deep RFM, where we are allowed to pick the feature dimension to be large enough so that this is satisfied), then the two optimizations have the same solution (where the min is taken to be an infimum) and, thus, the relaxation is without loss of generality.

We first show any $k$ realizable by our input-level kernel (under Euclidean distance on some dataset $R$) can be constructed by applying the input-level kernel to our data $X$ under Mahalanobis distance with appropriately chosen $M$ matrix. Let $k$ be a realizable matrix - i.e. any desired positive semi-definite kernel matrix for which there exists a dataset $R'$ on which our input-level kernel satisfies $k = \phi(d(R', R')))$, where $d(\cdot, \cdot)$ denotes the matrix of Euclidean distances between all pairs of points in the first and second argument. Our construction first takes the entry-wise inverse $r=\phi^{-1}(k).$ Since $k$ is realizable, this $r$ must be a matrix of Euclidean distances between columns of a data matrix $R$ of the same dimensions as $X$. Assuming the gram matrix of $X$ is invertible, we simply solve the system $R=NX$ for a matrix $N,$ and set $M=N^TN,$ which yields $k=\phi(r)=\phi(d(R, R))=\phi(d_M(X, X)),$ where $d_M(\cdot, \cdot)$ denotes the operation that produces a Mahalanobis distance matrix for the Mahalanobis matrix $M$.

We now give a construction demonstrating that the solution $k^*$ to problem $(2)$ is realizable up to arbitrarily small error using our input-level kernel on a dataset $R$ under Euclidean distance. We can realize the ideal kernel matrix that solves problem $(3)$, $yy^T$, up to arbitrarily small error by choosing $R$ according to a parameter $\epsilon > 0$, such that $\phi(d(R, R)) \rightarrow yy^T$ as $\epsilon \rightarrow 0$. In particular, for feature vectors $x_i,x_j$ of the same label in columns $i,j$ of $X$, we set the feature vectors $R_i, R_j$ for columns $i,j$ in $R$ to have $||R_i-R_j||=0$. Then, for $x_i,x_j$ in $X$ of different class, we set $R_i, R_j$ as columns of $R$ such that $||R_i-R_j|| > \epsilon^{-1}$. Then $k(R_i, R_j)$ is identically 1 for $R_i, R_j$ from the same class and converges to 0 for $R_i, R_j$ from different classes, as $\epsilon \rightarrow 0$, giving that $k = \phi(r) \rightarrow yy^T$.

For any choice of $\epsilon>0$, we can apply the procedure described two paragraphs above to construct $M$ that realizes the same $k^*$ using our input-level kernel applied to our dataset $X$ under the Mahalanobis distance. Therefore, the solution to problem $(2)$ can be constructed as the infimum over feasible solutions to problem $(3)$, completing the proof.

Is there any easy way to square this with the results on RFMs, that the right singular values align with $\nabla f$ terms?

This is a keen observation. You are correct that, although the NFA describes structure in $W^T W$, the correlation between NFM and AGOP is enabled by the alignment between the left singular vectors of $W$ and the gradients with respect to the pre-activations, as argued in [1]. This can be seen through decomposing the AGOP as $W^T K W$, where $K$ is the covariance of $\frac{df}{dx}$ and $x$ are the pre-activations. Then, the NFA holds provided that the right singular structure of $W$ is not perturbed through the inner multiplication by $K$, an alignment that naturally occurs through training by gradient descent. See [1] for a more detailed description of this phenomenon.

[1] Beaglehole, Mitliagkas, Agarwala. "Feature learning as alignment: a structural property of gradient descent in non-linear neural networks." arXiv pre-print, 2024.

It would be good to be explicit and put a subscript below the $\nabla$'s on the $\nabla f_\ell(x_{ci}^\ell)$

Thank you for this comment. We will fix this in the revision.

Why is Section 4.3 called “non-asymptotic analysis”?

We will rename this subsection “Connection to parametrized kernel ridge regression”.