Bayesian Treatment of the Spectrum of the Empirical Kernel in (Sub)Linear-Width Neural Networks

Comment: It is not immediately clear how the integrals in Theorem 3.4 can be explicitly evaluated.

Answer: In the integrals in Theorem 3.4, the integrands are explicitly known as functions of the integration variables $\Lambda, \Phi, \Phi^*$ and the measures over the integration variables are also known: Marchenko-Pastur map for $\Lambda$ and isotropic multivariate Gaussian for $\Phi, \Phi^*$. Given this, it suffices to sample $\Lambda, \Phi, \Phi^*$ according to these measures and to compute the integrands to calculate the integrals. The least straightforward step is to compute the Marcheko-Pastur map and we gave more details about this step in l. 424-426 in our initial submission (l. 431-472 in our revised submission).

Comment: Notation for $\Phi^*$ in Theorem 3.4 was never defined.

Answer: The notation was defined in l. 255 in our initial submission (l. 259 in our revised submission).

Comment: Between lines 300 and 303, the authors claim that for small datasets the computations would not be too computationally intensive. However, in footnote 3 (p. 6), they claim that the limit can be approximated using large objects. It is not completely clear what has to be large or small for the evaluations to be feasible.

Answer: The footnote highlights that our results are theoretically established in the limit of large dimensions. However, because of the kernel matrix inversion, dealing with large-scale datasets in practice is challenging. To rest the theory, we must therefore constrain ourselves to a setting where the dimensions are large enough for the quantities to have converged to their theoretical limits but small enough for the computations to be feasible. This is easier to achieve for simple, synthetic, datasets where the convergence to theoretical limits is quick, and more difficult for more complex real-world datasets. This partly explains the discrepancies we observe between our theory and experiments. However, please note that this is a shortcoming that the entire community working on theoretical infinite limits (especially with kernels) has to deal with, not a particular weakness of our work.

Question: Section 4 appears to be incomplete. It is not evident from the writing which of the figures correspond to the two examples they mention at the beginning of the section.

Answer: We have refined the presentation of our experiments to improve readability. Please see Figures 1 and 2 in our revised submission.

Question: In the first experiment they mention, with the "linear teacher", the authors do not specify how the $\beta$ is defined.

Answer: We have added the value we took for $\beta$ l. 42 in our revised submission.

Question: For Figures 1 and 2, how many simulations are the authors using to obtain the error bars? Also, is there any uncertainty around the blue lines in the same figures? If so, it should be reported.

Answer: We are using 10 simulations to obtain the error bars. Please see the improvements to our Figures in our revised submission.

Question: Can the authors reconcile the differences between the simulations in Figure 2 between their proposed estimates and the blue lines?

Answer: We have listed the sources of discrepancies in l. 476-479 in our initial submission (l. 483-485 in our revised submission): "the fact that we used finite values of P, N, N0 (whereas the theory is only exact in the limit); the SUA may not be fully accurate in this configuration; the limiting spectral measure may not be nonrandom."

Question: In Theorem 3.3, the authors bring in the notation $\lim_{\infty} P_{N_0}$ without having defined [...]. Does $P_{N_0}$ correspond to $p_{X}$ for a specific dimension?

Answer: $P_{N_0}$ (and $\lim_{\infty} P_{N_0}$) was defined in the preliminaries in l. 131-133 in our initial submission (l. 133-135 in our revised submission).

Question: The authors emphasize the rectified linear unit (ReLU) activation function (deservedly so). However I wonder if this approach can also work for other activation functions?

Answer: Yes, our approach does work for other activation functions, as stated in l. 260-261 in our initial submission for instance (l. 266-272 in our revised submission): the result from El Harzli et al. 2024 is established under very general assumptions on the activation function, namely measurability and Lipschitz continuity.

Comment: The experimental results are presented poorly. .

Answer: We have improved the presentation of our experimental results. Please see Figures 1 and 2 in our revised submission.

Comment: There are numerous prior works around deep kernel processes and machines that would be worth discussing in the related work. While this work uses a very different theoretical approach, it ultimately addresses similar conceptual issues.

Answer: We thank the reviewer for sharing these references which we have incorporated into the related work discussion. Please see l. 47-49 in our revised submission.

Question: What do the authors suspect would happen for networks in the $N >>P$ regime with mean-field $\mu P$ scaling that preserves strong feature learning? In this case, theory predicts that the kernels are non-random but the preactivations can become non-Gaussian as was also shown in the Bayesian setting. Do the spectral universality assumptions break down?

Answer: The spectral universality assumption relates to the prior kernel, not the posterior (as discussed l. 317-332 in our initial submission and l. 328-338 in our revised submission). However, it plays a crucial role in calculating the posterior by enabling integration over the prior. In the light of this, we do not think that the spectral universality assumption would cease to apply with the (non-lazy) scaling. This observation is consistent with the calculations from van Meegen and Sompolinsky (2024) in their appendix (equations A68 and A69), where they establish that the mean predictor remains unchanged compared to the lazy scaling (which is the scaling that we have) and that the variance is renormalised by a different factor (which vanishes as $N \to \infty$). This is exactly the behavior we would expect if the spectral universality assumption held, as hinted by our Theorem 3.5. Furthermore, the renormalisation factor is highly dependent on the limiting spectral measure (in the lazy scaling, our Marchenko-Pastur map) as established by the proof of our Theorem 3.5.

Comment: What do the authors think of the spectral universality assumption for cases where the matrices pick up low rank spikes such as learning single or multi-index models. Do spiked kernels violate the kernel renormalization theory?

Answer: In the aforementioned paper, the authors demonstrated an equivalence between training a two-layer neural network using only the first gradient step and training a low-rank spiked random feature model. This can be interpreted as altering the (prior) data representation to reside in a low-dimensional space, which we believe significantly constrains the orthogonal matrices permissible for constructing the prior kernel. As a result, we anticipate that the spectral universality assumption would fail in this scenario. However, this does not impact the validity of the spectral universality assumption in the contexts considered in our paper, as a single gradient descent step is not expected to replicate the full Bayesian training of a network.

The line numbering has changed again in our new revised submission (in response to a feedback from Reviewer H2W1). Please see below the line numbers relevant to your review.

"l. 317-332 in our initial submission and l. 328-338 in our revised submission" -> l.332-342 in our new revised submission

I thank the authors for answering my questions and clarifying my confusion regarding posterior/prior kernel. I will revise my score. My one remaining recommendation would be to emphasize that the current theory may not be able to capture the spiked kernels which result from large steps of feature learning, which would break the universality assumption.

Additionally, in the appendix, I think the variance should be $\left< (\delta f)^2 \right>$ rather than $\left< \delta f \right>$, unless I am mistaken.