Adaptive kernel predictors from feature-learning infinite limits of neural networks

Summary Clarification

Regarding the theoretical results, to avoid possible confusion, we decided to change the name adaptive Neural Network Gaussian Process Kernel (aNNGPK) into adaptive Neural Bayesian Kernel (aNBK) in the newer draft of the paper. Our feature learning theory for deep non linear Bayesian networks does not resort to any perturbative approach nor any Gaussian approximation of hidden layer pre-activation distributions, which are non-Gaussian in the regime where $\gamma_0 = \Theta_N (1)$ (Fig. 3(b) in the main text).

To the best of our knowledge, no prior work has attempted to solve the Bayesian feature learning saddle point equations non-perturbatively (Alg 1). We gave some intuition on how the fixed point equations for the kernels simplify in the $\gamma_0 \to 0$ limit using a perturbative approximation in Appendix A.5, but the theory and all the numerical results in the main text are not derived under this assumption.

Relations to Other Works

We thank the reviewer for pointing out these references, which are all worth to acknowledge in our "Related works" section. We had no intention of leaving them out on purpose, and we provided to expand and clarify how these works stand in comparison to our in the newer draft of the paper ("Related works" section) as well as next to the equations where they are presented.

Regarding Ref. https://proceedings.mlr.press/v235/fischer24a.html, we give a detailed explanation on how and why our work differs from the one of Fischer et al. in the "Relationship to Fischer et al" of Reviewer MHbU (which we invite the Reviewer to read).

The same can be said for Ref. https://www.nature.com/articles/s41467-023-36361-y, where the authors derive feature learning as a finite 1/width effect, being the predictor a Gaussian Process Regression (GPR) predictor. In the paper conclusions they say Central to our analysis was a series of mean-field approximations, revealing that pre-activations are weakly correlated between layers and follow a Gaussian distribution within each layer with a pre-kernel $K(l)$. In the paper analysis they add Remarkably, despite strong changes to the kernels and various non-linearities in the action, the pre-activation is almost perfectly Gaussian.

Again, our theoretical results are substantially different:

1. feature learning is not a 1/width effect in the parameterization we chose, and the kernels can exhibit arbitrarily large changes in their structure at infinite width (see Fig. (7) in the Appendix) when $\gamma_0 = \Theta_N(1)$.
2. The predictor has not the form of a GPR, but instead is a kernel predictor with task-adaptive kernel where the pre-activation distributions is highly non-Gaussian in deep non-linear networks ( as in Equation (7)). Feature learning in our setting has the $\Theta_N(1)$ effect of accumulating non-Gaussian contributions in single site density measure.

Regarding Ref. https://proceedings.mlr.press/v202/yang23k/yang23k.pdf, what the authors define as a "The Bayesian representation learning limit" in Sec. 4.3 consists in sending the number of output features (what they call $N_{L+1}$) to infinity, which has the effect of rescaling the likelihood term in the action by $N$, but they don't rescale the readout of the last layer weights as we did which leads to some differences in the resulting eqns.

As an example, consider linear network with $L=1$ hidden layer. Our equation is

$

\Phi^{-1} = C^{-1} - \gamma_0^2 \Phi^{-1} yy^\top \Phi^{-1}

$

This differs from the equation 128 of Yang et al which states

$

C^{-1} = \Phi^{-1} yy^\top \Phi^{-1}

$

which shows that our theoretical predictions are different in general, especially for small $\gamma_0$.

We can obtain their equation from our result by taking a very rich limit $\gamma_0 \to \infty$ and rescaling the definition of $\Phi = \gamma_0 \bar{\Phi}$ leading to the equation

$

C^{-1} = \bar{\Phi}^{-1} yy^\top \bar{\Phi}^{-1} + \mathcal{O}(\gamma_0^{-1})

$

However our equation holds for a variety of richness levels, interpolating between lazy and rich regimes.

What is new for GF/GD?

Please see response to reviewer MHbU on this topic.

Positioning of Paper in Bayesian NN Literature

We thank the reviewer for pointing out these issues. In the newer paper draft, we expanded the "Related work" section and provided a more detailed comparison to these relevant prior works.

Numerics Details

Thanks for this important question. We expanded Appendix F with details concerning the theory solvers and the numerical experiments on finite width NNs. We have an entire new subsection to discuss the computational costs (time for kernels, time for final outputs and memory) to solve the theory compared to end-to-end training. We discuss this in detail in the answer of "Methods And Evaluation Criteria" for Reviewer eXoJ, which we invite the Reviewer to read.

Analysis on the Cost of Solving Mean Field Eqns

We thank the reviewer for this important question, and in the new draft (Appendix F) we gladly added an extensive discussion on numerical methods and computational costs.

Regarding the cost of solving for Algorithm (1) compared to the network training, we can elaborate about the computational requirements, which is the time for kernels $\{\mathbf \Phi^{\ell}, \hat{\mathbf \Phi}^{\ell}\}_{\ell=1}^L$.

If our neural network has an input dimension $D$ fixed and a number $P$ of training points, our infinite width limit theory predicts that in order to solve for the kernels to get the network predictor, one has to solve the $\text{min-max}$ optimization problem as in Equation (8). So, in principle, we have an inner loop that solves for $\hat{\mathbf \Phi}^{\ell} (\mathbf \Phi^{\ell})= \arg \max_{\{\hat{\mathbf\Phi}^\ell\}} S(\{ \mathbf\Phi^{\ell}, \hat{\mathbf\Phi}^\ell \})$ with a $t_{\text{inner}}$ number of gradient steps, and an outer loop which, given $\hat{\mathbf \Phi}^{\ell}$ it solves for $\mathbf \Phi^{\ell}(\hat{\mathbf \Phi}^{\ell})= \arg \min_{\{\mathbf\Phi^\ell\}} S(\{ \mathbf\Phi^{\ell}, \hat{\mathbf\Phi}^\ell \})$ for a $t_{\text{outer}}$ number of outer gradient steps.

An estimate of the computational overhead per self-consistent (outer) update in the theory solver is approximately $\mathcal{O}(t_{\text{inner}} B P^2)$. Indeed, the most expensive operation in computing the action $S(\{ \mathbf\Phi^{\ell}, \hat{\mathbf\Phi}^\ell \})$ from Equation (6) is the Monte Carlo estimate of the non-Gaussian single site density $\mathcal Z_\ell$. At each inner step $t=1,\ldots t_{\text{inner}}$, we sample a batch of $B$ Gaussian vectors $\mathbf h_k$ and use these to estimate

$\mathcal{Z}_\ell = \left< \exp\left( - \frac{1}{2} \phi(\mathbf{h})^\top \hat{\mathbf{\Phi}}^\ell \phi(\mathbf{h}) \right) \right>$

$\approx \frac{1}{B} \sum_{k=1}^B \exp\left( - \frac{1}{2} \phi(\mathbf{h}_k)^\top \hat{\mathbf{\Phi}}^\ell \phi(\mathbf{h}_k) \right)$

This operation has a cost of $\mathcal{O}(BP^2)$. Because we have to repeat the same sampling for $t_{\text{inner}}$ steps, we get $\mathcal{O}(t_{\text{inner}} B P^2)$. Notice that, in the case of deep linear network, because the single site is Gaussian, the computational overhead is much smaller, and since one can explicitly solve for the conjugated kernels $\hat{\mathbf\Phi}^{\ell}$ as in Equation (15), the cost of each step is instead $\mathcal{O}(P^3)$.

For the full Langevin dynamics for $T$ gradient steps, the cost of training is instead $\mathcal{O}(T N^2 P)$, being $N$ the network width. This is much less than the $\mathcal{O}(t_{\text{outer}} t_{\text{inner}} B P^2)$ for the full theory if we suppose the number of gradient steps for convergence between theory and simulations ($T$) to be of the same order.

Despite this being the case in general, for which working on optimizing the computational overhead of the theory can be a future work focus, our Algorithm (1) still represents a step forward compared to Algorithm (2) as introduced in Bordelon et al 2022. In the case of Algorithm (2), solving for the DMFT analysis has a time requirement that scales as $\mathcal{O}( T^3 P^3)$ for the full dynamics, because in this case the kernels are $PT \times PT$ matrices. For optimizing this Algorithm (2), in Appendix E.1 we show a very simple and concrete example on how with weight decay the fixed point equations of DMFT at convergence can be solved without integrating the full dynamics, which we see as a promising future direction.

To get an order or magnitude, in our simulations: $N=1024$, $t_{\text{outer}}= 10^4$, $t_{\text{inner}}=2\times 10^2$, and $T=2 \times 10^4$, $B=2\times 10^4$ and sample size at most $P=10^3$.

Derivations could be clearer

We are working on a cleaner and exaustive Appendix A, especially concerning the saddle point for which we provided the derivation in the new Appendix A.

Experimental Details

We apologize for missing out the details concerning our numerical methods. We added a new section in Appendix F to help clarifying the settings.

More Numerical Details

We are currently working on releasing both the codes that solve for the theories, which we have all intention of doing. In the meantime, regarding the numerical methods explanation, we added two new sections (Appendix F.2/F.3) as mentioned.

Time complexity compared to known kernel

The time for final outputs are instead: $\mathcal{O}(t_{\text{inner}} t_{\text{outer}} B P^2)$ for Algorithm (1), $\mathcal{O}(T^3 P^3)$ for Algorithm (2), $\mathcal{O}(P^3)$ for a kernel regression with a known kernel, and $\mathcal{O}(TN^2P)$ for the full network dinamics.

Relationship to Fischer et al

We thank the reviewer for pointing us toward Fischer et al 2024 which we have added a detailed comparison to in the main text, as well as expanded the discussion of Seroussi et al and Rubin et al.

Our work differs from Fischer et al in a number of important ways. These differences lead to different predictors and interpretations on what drives feature learning.

1. We study a different feature-learning scaling limit than Fischer et al., therefore the features learned by our networks differ than Fischer et al. We keep $P= \Theta_N(1)$, and take width $N\to \infty$ in $\mu P$ parameterization (where $f= \frac{1}{N} w \cdot \phi$ and likelihood scaled up by $N$). They study a proportional limit in NTK parameterization. These scalings behave differently (see Figure 7).
2. We arrive at kernel equations in different ways. In our setting, the deterministic saddle point equations for the learned kernels in this feature learning limit are exact at infinite width.

In our theory the saddle point equations for the kernels at infinite width are provided in Eqn (31), which involve non-Gaussian single-site densities and no linearization around the lazy solution.

In Fischer et al., the saddle points equations for kernels at fixed $P$ and infinite width are instead (with 1/$N$ corrections included) provided in Eqns 15-16, which are clearly different from ours and involve Gaussian single-site averages. In their equations feature learning contribution vanishes when $N\to \infty$ at fixed $P$; they are effectively computing weak corrections to the lazy limit and fluctuations of the kernels.

3. We develop a solver for the min-max problem (Alg 1) for the exact eqns, which we verify this with experiments. To our knowledge no prior work has attempted to solve the exact saddle point equations.

In our work, it is not the finite width fluctuations that drive useful feature learning as in Fisher et al, but feature learning can already be accessed in the $N \to \infty$ limit where kernels concentrate. This is at odds with one of the primary takeaways of Fischer et al who argue that "larger kernel fluctuations enable stronger feature learning."

That said, both Fischer et al and our work identify kernels and conjugate kernels as the key order parameters ($\{ \Phi, \hat\Phi\}$ in our work $\{ C, \tilde{C} \}$ in their work) and both papers express a large deviation principle for them (equation 8 in our work and equation 6 in their work).

We will cover this in greater detail in App B.

What is new in GD setting?

The key contribution in this part is recognizing that knowledge of the final NTK is sufficient to characterize the learned solution, but only if weight decay is included. This suggests that alternative solvers focused on weight decay fixed points could circumvent the need to simulate through the dynamics.

More details about Expts

Additional experimental details are provided in the new version of Appendix F.

Questions

1. If $\gamma_0 = \frac{1}{\sqrt N}$, then we recover the NTK scaling and corrections to this lazy limit can be exacted at small but finite $\gamma_0$ by considering a perturbation series for the NNGP (or NTK) kernel in powers of $1/N$. In Appendix B we discussed finite width effects in NTK parameterization. We stress that at fixed $\gamma_0$ and finite width, our theories (Fig 7) are different from these theories.
2. In principle, one no longer needs to compute the time integrals if weight decay is used (see App E.1). In practice, however, we have been mainly integrating through the dynamics to solve for the final weight decay solution, but fixed point solvers for GD+weight decay could potentially be accelerated compared to DMFT.
3. In general, for gradient flow, one needs to know the full history of the NTK over training.
4. $\mu$P scaling is becoming the industry standard in many SOTA models, including for transformers https://arxiv.org/abs/2203.03466, https://arxiv.org/abs/2407.05872, https://x.com/andrew_n_carr/status/1821027437515567510. Part of the attraction for practitioners is the consistency of dynamics across widths https://arxiv.org/abs/2305.18411.
5. Our Bayesian theory has no assumptions on the homogeneity of the activations function. We did experiments in this setting with $\tanh$, that we will gladly add in App A.

Instead, GF + weight decay for non-homogeneous activation functions no longer satisfies a representer theorem for the $\text{aNTK}$ kernel as in Equation (10).

6. For CIFAR we resized the images to $28\times 28$ pixels and then we convert them to grayscale.

7. The simplest case is $P=1$ and $L=1$ with $\phi(h) = \tanh (h)$, $p(h) \propto \exp \Big(-\frac{1}{2}C^{-1} h^2 - \frac{1}{2}\hat{\Phi}\tanh(h)^2 \Big)$ Since $\hat{\Phi} < 0$, $p(h)$ can develop a Mexican-hat shape.

8. We suspect that Langevin noise acts as an entropic regularizer that reduces overfitting in small data regimes.

We thank their reviewer for their detailed feedback and suggestions, which we try addressing below.

How is the infinite limit analyzed?

Here we provided exact expressions for the adaptive kernel predictors that one obtains in various feature learning limits through rigorous analyses of the underlying limits at fixed dataset. We believe our usage of the term "analysis" is appropriate, however we are happy to consider an alternative suggestion. We disagree that this is a "simple" derivation. Similar analyses, for example, providing kernel limits in the lazy regime without a study of the statistical generalization error have been very valuable in the past.

Derivations Dependent on Prior Work / What is new?

We thank the reviewer for pointing out some of these deficiencies in our presentation. We didn't use any previous theoretical results for deriving the aNBK (adaptive Neural Bayesian Kernel, the old aNNGPK), because our theory is new (see "Relationship to Fischer et al" discussion in Rev. MHbU or "Relation to other work" in Rev. Za7T). The derivation of Equation (9) from (3) is a central result of Langevin dynamics and covered in textbooks such as (https://www.cambridge.org/core/books/statistical-physics-of-fields/06F49D11030FB3108683F413269DE945). We feel inclusion of this basic result in the paper is not necessary, however we will gladly provide a reference below. Regarding dependence on [Bordelon and Pehlevan, 2022], we have already included an Appendix Section D reviewing the necessary facts from that paper and pointed to it from the main text. We are working on extending this section and make it more clear.

Add more precise statement of results (Theorem/Proof style)

We thank the reviewer for this suggestion. We added statement results as Theorems in the newer draft version and improved the clarity of the assumptions required for our results. In particular, we improved the explanation and statement of equation 8 and provided a more detailed proof of this result (what was the derivation sketch).

Appendix A confusing... How do we get generalization?

Thanks for this point. We simply evaluate the predictions of our adaptive kernel predictors on a set of test points to measure/estimate the generalization error for that point. This is not meant to be an analysis of the data-averaged or worst-case generalization error. However, these predictions on test points will agree with neural network predictions with high probability over the random initialization of the network weights and the random Langevin noise in the case of Bayesian networks.

Execution Issues with Result Statement beyond Lazy Limit

We apologize if the results were not as clear as possible. We hope our responses above clarified our contributions. We are making changes in "Related works sections" to address these clarity issues and clarify our contributions.

Questions

1. $\delta$ is used without being introduced. What is this? Should it be there?

We will add that $\delta$ here is the Dirac delta function (a function with the property $\int dt' \delta(t-t') g(t') = g(t)$ for any $g$), which specifies that the noise in the Brownian motion is uncorrelated across time. To see the discrete time version of an uncorrelated process, we could define a collection of random variables $\{ \epsilon_t \}$ with covariance structure $\left< \epsilon_t^2 \right> = \sigma^2$ and $\left< \epsilon_t \epsilon_{t'} \right> = 0$ for $t \neq t'$.

The continuous time version of this is one where the covariance of the process $\{ \epsilon(t) \}$ is a Dirac delta function.

2. What is up with Eqn 18?

We thank the reviewer for pointing out the typo in Eq.(18). We solved the inconsistency and fixed the notation for the $\mathbf W$s.

3. How do we get eqn 5?

We apologize that this was unclear. We added a more detailed derivation of this in Appendix A.2.

In the case of $\sigma(s)=s$ we have

$f(x) = \frac{\beta}{\lambda} \sum_{\mu} \Delta_\mu \Phi(x,x_\mu).$

We can solve for $\Delta_\mu \equiv -\frac{\partial L}{\partial s_\mu}$

$\mathbf \Delta = \mathbf y - \mathbf f= \mathbf y - \frac{\beta}{\lambda}\mathbf\Phi \mathbf \Delta= \left[ \mathbf I + \frac{\beta}{\lambda} \mathbf\Phi \right]^{-1} \mathbf y$ as desired.

4. Why does $\gamma$ show up in the dynamics?

The raw learning rate has to vary with $\gamma$ to maintain constant scale updates to the output to the network (see https://pehlevan.seas.harvard.edu/sites/g/files/omnuum6471/files/pehlevan/files/princeton_lecture_notes.pdf).

5. Why do Dynamics Converge to the Posterior?

When $\beta = \Theta_N(1)$, the dynamics in Eq.(3) are known as Langevin, which are known to converge to a stationary density that has the form of Eq. (9). A useful and complete reference for this derivation can be found in https://www.thphys.uni-heidelberg.de/~wolschin/statsem23_6.pdf, Eq.(54), or in https://www.stats.ox.ac.uk/~teh/research/compstats/WelTeh2011a.pdf.