Adaptive kernel predictors from feature-learning infinite limits of neural networks
Authors: Clarissa Lauditi
Blake Bordelon
Cengiz Pehlevan
Venue: ICML 2025 poster (2025)
Keywords: Kernel methods; feature learning; Bayesian networks
Primary area: deep_learning->theory
TL;DR: A theory of feature learning for Bayesian networks at infinite width.

Abstract: Previous influential work showed that infinite width limits of neural networks in the lazy training regime are described by kernel machines. Here, we show that neural networks trained in the rich infinite-width regime in two different settings are also described by kernel machines, but with data-dependent kernels.  For both cases, we provide explicit expressions for the kernel predictors and prescriptions to numerically calculate them. To derive the first predictor, we study the large-width limit of feature-learning Bayesian networks, showing how feature learning leads to task-relevant adaptation of layer kernels and preactivation densities. The saddle point equations governing this limit result in a min-max optimization problem that defines the kernel predictor. To derive the second predictor, we study gradient flow training of randomly initialized networks trained with weight decay in the infinite-width limit using dynamical mean field theory (DMFT). The fixed point equations of the arising DMFT defines the task-adapted internal representations and the kernel predictor. We compare our kernel predictors to kernels derived from lazy regime and demonstrate that our adaptive kernels achieve lower test loss on benchmark datasets.

Adaptive kernel predictors from feature-learning infinite limits of neural networks
Clarissa Lauditi!
Blake Bordelon 123
Cengiz Pehlevan 123
Abstract
Progress in this endeavor is limited but promising. An in- fluential set of results that motivates this paper identifies infinitelywwide neural networks under certain initializations as nonparametric kernel machines (Neal, 1995; Cho & Saul, 2009a; Lee et al., 2018; de G. Matthews et al., 2018; Arora et al., 2019b; Jacot et al., 2020; Lee et al., 2020). This important because kernels are theoretically well well-understood (Rasmussen & Williams, 2006; Scholkopf & Smola, 2001), offering rich mathematical frameworks for analyzing the ex- pressivity and generalization properties of neural networks. A drawback of this identification is that it works in the lazy regime of neural network training (Chizat et al., 2020), i.e. when data representations are fixed at initialization and do not evolve during learning. However, sfute-of-the-rr deep networks operate in the rich, feature learning regime (Geiger et al., 2020; Vyas et al., 2022; Yang & Hu, 2022; Vyas et al., 2023), where they adapt their internal representations the structure of the data.
Previous influential work showed that infinite width limits of neural networks the lazy training regime are described by kernel machines. Here, we show that neural networks trained in the rich, feature learning infinite- width regime in two dif- ferent settings are also described by kernel ma- chines, but with data-dependent kernels. For both cases, we provide explicit expressions for the ker- nel predictors and prescriptions to numerically cal- culate them. To derive the first predictor, study the wargh limit of feature-learning Bayesian networks, showing how feature learning leads to lask-relevant adaptation of layer kernels and pre- activation densities. The saddle point equations governing this limit result min-max optimiza- tion problem that defines the kernel predictor. To study gradient
derive the second predictor, flow training of randomly initialized networks trained with weight decay in the infinitewwidth limit using dynamical mean field theory (DMFT). The fixed point equations of the arising DMFT defines the task-adapted internal representations and the kernel predictor. We compare our kernel predictors kernels derived from lazy regime and demonstrate that our adaptive kernels achieve lower test loss benchmark datasets.
Motivated by these observations, here, we ask whether in- finitely wide neural networks still admit nonparametric ker- nel predictors in the rich domain. And if so, what charac- terizes these predictors? Since in the rich regime data rep- resentations evolve dynamically according to the training dynamics, identifying the nature of predictors in this regime is crucial for uncovering the principles behind feature learn- ing. Further, as wider networks are believed to perform better on the same amount of data (Hestness et al., 2017: Novak et al., 2018; Kaplan et al., 2020; Hoffmann al., 2022), it opens the possibility of directly training infinitely- wide feature learning networks through their corresponding kernel machines for obtaining the best performing instances of a given model architecture.
1. Introduction
As neural-networkk-based artificial intelligence is increas- ingly impacting many corners of human life, advancing the theory of learning in neural networks is becoming more and more important, both for intellectual and safety reasons.
In this paper, we derive kernel predictors for infinitely wide multilayer perceptrons (MLPs) and convolutional neural net- works (CNNs) in the rich regime. They are kernel predictors in the sense that
Equul contribution 'John A. Paulson School of Engi- neering and Applied Sciences. Harvard University "Center for Brain Sciences 'Kempner Institute. Correspondence to: Clarissa Lauditi <clauditi @ aerddii@g.harvaat h.vaaard..ddu> Blake Bordelon <blake baaake_bordelone @ g. edu>, Cengiz Pehlevan <epehle- van @ seas, harvard.edu2.
P f(x) = (an
W=1
where a a nonlinear function, al are scalar coefficients, [ are training data, and x the test point. The kernel depends on the architecture of the network, and,
Proceedings of the 4dnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
importantly, adapts the training data, unlike the Neu- ral Tangent Kernel (NTK) (Jacot et al., 2020) and Neural Network Gaussian Process Kernel (NNGPK) (Cho & Saul, 2009b; Lee et al., 2018; de G. Matthews et al., 2018) derived from lazy infinite limits.
gradient flow dynamics the final learned network predictors behave as kernel predictors. We provide the DMFT fixed point equations that define the aNTK predictor.
4. We develop numerical methods to solve for our predic- tors.
To arrive at these predictors, we start from existing work characterizing the feature-learning infinite-width limits of deep neural networks under gradient flow dynamics. Previ- ously, Bordelon & Pehlevan (2022) adopted the dynamical mean-field theory (DMFT) (Martin et al., 1973; De Domini- cis, 1978; Sompolinsky & Zippelius, 1981; Arora et al., 2019a; Arous et al., 2004) approach to the training dynam- ics of deep MLPs and CNNs under the maximal update parameterization (μP) (Yang& & Hu, 2022; Yang et al., 2022; Mei al., 2018; Rotskoff & Vanden-Eijnden, 2022), which leads to a feature-leaming infinite-widde limit. This DMFT results in set of stochastic integro-differential equations in terms of summary statistics, which define the deterministic training time (t) evolution of the deterministic network pre- dictor f(x, t). Performing inference with this framework requires solving these equations over the training time using sophisticated Monte Carlo techniques. Further, the com- plexity and history dependence of the equations make the predictor interpretation challenging.
5. For kernels arising from deep infinitely-wide linear Bayesian networks, we solve the saddle point equations exactly via recursion, bypassing sampling strategies needed for the nonlinear case. We analyze the behav- ior of kernel-task overlap parameters across depth for whitened data.
6. We provide comparisons of our adaptive kernel ma- chines to trained networks and lazy NTK and NNGPK predictors for MLPs and CNNs. We demonstrate that our adaptive- kernes are descriptive of feature-learning neural network training on a variety of metrics, in- cluding test loss, intermediate feature kernels, and pre- activation densities. In addition, they outperform lazy kernel predictors on benchmark datasets.
1.1. Related works
Neural networks as kernel machines. certain initial- ization and parameterization schemes, taking the width of neural network to infinity leads the model to learn ker- nel machine, with kernel that depends only on the initial task insependeennd features which do not change during train- ing (Jacot et al., 2020; Lee et al., 2018). This "lazy train- ing" regime has been extensively studied, particularly in the context of infinitely wide networks (Chizat et al., 2020; Jacot al., 2020; Lee et al., 2018: 2020). However, neural networks outside of the lazy limit (in the so-called "rich regime") often perform better than their corresponding ini- tial NTKs (Vyas et al., 2022) and are not obviously related to kernel machines. Domingos (2020) argues that gradient descent training for any deep network corresponds to re- gression with a history-dependent "path kernel", though this definition does not satisfy the standard representer theorem where coefficients are only functions of the training data as in Equation (1). In general, the learned function of a network trained with gradient flow can always be written as an integral over the history of eтror coefficients on training data and a time evolving NTK. In the present work, how- ever, we are focused on when the final solution a network converges to satisfies history independent representer the- orem as in Equation 1. Some experimental and theoretical works have indicated that this is often the case, where re- gression with the final NTK (computed with gradients at the final parameters) of network provides an accurate ap- proximation of the learned neural network function during rich training (Geiger et al., 2020; Atanasov et al., 2021; Wei et al., 2022). A complete theoretical understanding of when
Our main observation is that we can analyze the network predictor at convergence ((x L. = oo) in two ways to obtain deterministic adaptive kernel predictors: (1) by interpreting the dynamies of gradient flow with added white noise as sampling the weights from a Bayesian posterior, (2) by studying the fixed points of DMFT equations for gradient flow with weight decay. Specifically, our contributions in this work are the following:
1. We study the noisy gradient- flow dynamics with weight decay in the rich regime for MLPs and CNNs. We iden- tify two novel infinite-width limits (Table 1, Figure 1) that lead to adaptive kernel machine interpretations of neural networks. We name the corresponding kernels adaptive Neural Bayesian Kernel (aNBK) and adaptive Neural Tangent Kernel (aNTK) for reasons that will be apparent below.
2. To analyze the first of these limits, we introduce a novel Bayesian interpretation of feature-learning neural net- works, where the posterior distribution at infinite width characterizes the network's 's state after training. Ana- lyzing this posterior in the infinite width limit using statistical mechanics methods, we identify a min-max optimization problem, arising from a saddle point ar- gument, that defines the aNBK predictor.
3. To analyze the second limit, we invoke the DMFT anal- ysis of gradient-flov dynamics (Bordelon & Pehlevan, 2022). We show that when weight decay is added
2

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
A Auo
A.p =0
Kee Kinal
Kust Kivel
rest Prediction KNIK
Thst Prodiction
fest pred ction
Test Prediction
Test Prediction NN
Prediction NN
(a) Gradient flow À =
(b) Gradient flow >> 0
(c) Noisy gradient flow
Figure 1. Test network predictors of two-layer MLP (width N 5000) trained with P = 300 data of two-classes of CIFAR 10 compared with the theoretical kernel regression predictors. Three panels are for different regimes of regularization À and temperature 1/8. In all the three cases of X, Kiii represent the network predictors initialization, while / correspond either aNTK at convergence for (a)((b) or aNBK that at convergence for (c). (a) When À = 0 the NN predictor of gradient flow without weight decay is not a kernel predictor; (b) instead, when À > 0 the test prediction is well captured by fANTK (Eq. (11)). (c) Bayesian empirical predictor is well-described by aNBK predictor at convergece. Here, the Bayesian predictor refers to the predictor obtained after thermalization when training with noisy gradient flow (i.c., Langevin dynamies). We expect these matches to be exact for adaptive kernels at infinite width.
and why the correspondence between final NTK and final network function holds is currently lacking.
can adapt due to feature learning). In this setting, feature learning acts as a 1/ 1/width effect in the posterior measure. Fischer al. (2024) derive a general large deviation princi- ple for the distribution of kemnels under the posterior, which holds for any a P/N, however, they focus on linear re- sponse theory for the kernel updates since feature learning induces O(N-1) corrections to the the posterior distribu- tion of kernels. Since the kernels do not concentrate in the proportional limit P, NN - os with P/N = a, Fischer et al. (2024) also compute how fluctuations in the kernels propa- gate across layers. They show that the linear response theory used to track fluctuations can also be used to compute fea- ture learning corrections when linearizing the action around the NNGP kernels. Their results are thus non-perturbative in a = P/N but perturbative in what we call "richness" 70- Similarly, (Seroussi et al., 2023) explore variational approx- imations of the hidden neuron activation densities at finite width. More recent versions of these proportional theories have begun to explore other parameterizations (Rubin et al., 2024a) in linear networks, resorting to cumulant expansion of activation densities (up to Gaussian order) in one-hidden layer non- networks.
Adaptive kernels. In Bayesian neural networks, several works have identified methods that describe learned net- work solutions beyond the lazy infinite-width deseription of NNGP regression. Some works pursue perturbative approx- imations to the posterior in powers of 1 1/width (Zavatone- Veth et al., 2022a; Roberts et al., 2022) alternative cumu- lant expansions of the predictor statistics (Naveh & Ringel, 2021). Others analyze the Bayesian posterior for deep linear networks, which are more analytically tractable (Aitchison, 2020; Hanin& & Zlokapa, 2023; Zavatone- Veth et al., 2022b: Bassetti et al., 2024) and actually also capture the behavior of deep Bayesian nonlinear student-teacher learning in particular scaling limit (Cui et al., 2023). Several works Bayesian deep learning have argued in favor of a propor- tional limit where the samples and width are comparable (Li & Sompolinsky, 2021; Pacelli et al., 2023; Aiudi et al., 2023; van Meegen & Sompolinsky, 2024; Baglioni et al., 2024; Fischer et al., 2024). the context of fully connected (MLP) linear neural networks, (Li & Sompolinsky, 2021; Pacelli et al., 2023) argued that, the mean predictor under the posterior is the same as regression with the lazy NNGP kernel, though with rescaled ridge and a predictor variance that depends on scale factor ((a) that both change as a function of a = P/N. Extending this result to deep net- works Pacelli et al. (2023) found that each layer has a scale renormalization constant Qa(.) while Aiudi al. (2023) showed that convolutional architectures are characterized by a matrix of scale renormalization constants Q (ax) that capture learned space x space that adapt to the data.
Closer to our approach, Yang et al. (2023) analyze rescal- ings of the Bayesian likelihood induce representation learning at infinite width. This approach has also been used to explain sharp transitions in the behavior of the poste- rior as hyperparameters, such as prior weight variance, are varied in large width networks (Rubin et al., 2024b). Alter- natively, some works have developed non-neural adaptive kernel algorithms, which filter information based on gradi- ents of a kernel solution with respect the input variables (Radhakrishnan al., 2022), which exhibit improvements in performance over the initial kernel and can capture interest- ing feature learning phenomena such as grokking (Mallinar
In the same NTK parameterization, Fischer et al. (2024) and Seroussi et al. (2023) developed theories of feature learning where kernels adapt more flexibly (each entry in the kernel
3

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
al., 2024).
width limit.
In our work, we study the infinite width NN + 0o limit of Bayesian networks at fixed P in the mean-field (μP) pa- rameterization setting, As consequence, feature learning is not driven by finite width and cannot generally be ex- tracted from linear response theory (or can be interpreted in terms of fluctuations since kernels concentrate). Con- trary to previous results (Li & Sompolinsky, 2021; Pacelli et al., 2023; Aiudi et al., 2023; van Meegen & Sompolinsky, 2024; Baglioni et al., 2024), in our theories, the kernels deterministically adapt to data. These kernels can exhibit arbitrarily large changes in their structure (rather than just scale or spatial renormalization) at infinite width, and these cannot be obtained from linear response theory, at odds with previous computations (Fischer et al., 2024; Seroussi et al., 2023; Rubin et al., 2024a) (see Figures 1 & 3). Addition- ally, our adaptive kernels at infinite width are deterministic quantities, while in previous scaling theories, one has to compute kernel fluctuation statistics to get non-parametric predictors in the feature learning regime. To the best of our knowledge, we are also the first to solve for these adap- tive kernels in the Bayesian setting without resorting any perturbative (in richness "oo) approach nor any Gaussian approximation of hidden layer pre-activation distributions, which are Gaussiass at infinite width.
3. A Adaptive kernel limits of training dynamies
Next, we study the rich training dynamics of the NN defined as in Eq. (2) to arrive at adaptive kernel predictors. In particular, we study the infinite time limit / + oc of the noisy gradient-flow dynamics
d0(t) = - - - AB ((t) dt + 2 de(t) (3)
for the collection of weights 0- Vee (w)) w((1,, loss function L()) and for ridge X, ve [L.. Here, y2 = ensures the feature updates to be 1₂(1) in the infinite width limit (Bordelon & Pehlevan, 2022), and de is Brownian motion with covariance structure (dee())e()))) = 8(t - I (being 5 the Dirac delta), whose contribution to the dynamics can be switched off by tuning the tempera- ture T 1/ + O. When the Brownian motion is on, this dynamics can be interpreted as sampling from a Bayesian posterior, which will be detailed below. When it is turned off, this is gradient-flow dynamics with weight decay.
It is well know that infinite limits of this training dynamics in the lazy regime (%o - 0) lead to kernel machines defined by the NTK (Jacot et al., 2020) and NNGPK (Cho & Saul, 2009b; Lee et al., 2018: de G. Matthews et al., 2018). Here, we show that there infinite limits in the rich regime that also lead to kernel predictors, but this time these kernels adapt to data. The order of limits for width, time, tempera- ture, and feature learning strength parameters {N, t, 3 to get either the already known lazy (NNGPK, NTK) or novel adaptive kernel predictors is shown Table 1. The latter kernels correspond either to the infinite width limit of NN at convergence (i.e. 1 + oc) that learns with (aNBK) or to the infinite time limit of an infinitely wide NN learning with gradient flow and weight decay (i.e. B - oo) (aNTK).
2. Preliminaries
We start by describing our setup. Here, for ease of presenta- tion, we discuss the fully connected MI.P setting, referring to Appendix D for the case of CNNs.
For an empirical training dataset D {z Up 1p of size P. input vectors F E IR. and labels yas we define the output of MLP with L layers
0
Jn ( ( W/ (4) -(()))
IVN
y 6+1 = VNe w o(h,), h/ = VD
W (0) Z.p.
NNGPK
aNBK (ours)
where ((-) represents a homogeneous transfer function in its parameters s(al) = a"<('' (00) hí, E R. at layer l is the pre-activation vector and W E Ne+1NN is the matrix of weights to be learned. We initialize each trainable param- eter as a Gaussian random variable W ~ (N 1) with unit variance, in such a way that in the infinite width limit Nr = N + oc, Vl E { the pre-activations at each layer will remain e (1). At the same time, we scale the network output by a factor 7 2o V in order to study feature learning. This allows to interpolate between lazy limit description of NNs when 70 - 0. and a rich regime descrip- tion when To- = (1). This parameterization known as maximal update parameterization (μP) (Yang & Hu, 2022), allows feature learning by enabling preactivations to evolve from their initialization during training even in the infinite-
lim - lim lim
lim oz lim { Ox=
8 3= Ox(1)
NTK
aNTK (ours)
lim lim---o lim limg-+0 t= Ox
lim lim limg-
7o Ox(1)
Table 1.
Limiting orders for (N,t, ro in the dynamics (3) to get either (known) static kemness NNGPK & NTK (left column), or (new) adaptive kernel predictors (right column).
Next, we give the functional forms of our novel aNTK and aNBK predictors and briefly discuss their derivations. Details are given in Appendix A.
4

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
3.1. aNBK
Mingard et al., 2020; Naveh et al., 2021) over the trainable parameters 0 given the dataset D, which can be interpreted as a Bayesian posterior with log-likelihood C()) and Gaussian prior of scale X-1/2
As described in Table if we take first the & + oc limit, and then the N - oo limit with temperature (1/3) and feature learning strength (ro) fixed, we arrive at an adaptive kernel predictor (see Fig. 1)
L Σ
p(0 (D) z exp B)
e
1|
(9)
(
1 A1 1=1
P A S (m)
JaNuK( =
(4)
The distribution of Eq. (9) can be studied in the overparam- eterized thermodynamic limit where the width at cach layer N - oc while P = x1)) When / - oo, the posterior is dominated by the set of global minimizers of the loss for the training data (i.e. solutions l that minimize [(0)). We are interested in the partition function Z of Eq. (9), which is function of the order parameters { 6 since it has the form Z x S S d d exp (-NS(Ø0) $6))
with A, = DC being the error signal, 8/ = No w(().
Bap ((h()) the pre-readout as in Eq. 2, and $ = ( p(h²) )) the train-test feature kernel at the last layer L. In this notation, the kernel matrix element is given by $L uu = (Z,,. In the squared loss case with linear readout (ar(s s), the predictor fxx) becomes a kernel regression predictor of the form
(+)
where S(Ø*, Φc) is the Bayesian action given in Equa- tion (6). When N is large, p(( is exponentially domi- nated by the saddle points of S. and looking for the values of { 1 which makes S locally stationary means solving the min-max optimization problem in Algorithm 1. In the lazy learning limit tt - 0. the dual kernels vanish Φ + 0 and we recover Gaussian preactivation densities, consistent with the NNGP theory (Lee et al., 2018). A full account of this derivation is given in Appendix A.
SANBK (2)) - (x)) T Tb + An
y,
(5)
3
where (a)) = (1 (for derivation, see Ap- pendix A.2).
In Appendix A, we show that the kernel gh is given by solution to a min-max optimization problem that involves the data-adaptive kernel x- E RP',, intermediate layer adaptive-kernels & E IR PxP and dual adaptive-kernel vari- ables &' E PxP. Here, we present this min-max problem for the squared loss and linear readout for simplicity, see Appendix A for the full expressions. First, we define the action:
We present a numerical algorithm to calculate the aNBK Regression Predictor from data in Alg. 1.
Algorithm aNBK Regression Predictor
Input: Dataset D {z yu ) with covariance 4. fp 1 Tz Ty. hyperparameters { Ae,, step size 8. Output: Predictor [(x) for any test point E (Prest).
S(1
1
Σ
Tr
T' +2
20
I
$
y
2
3
1=1
Generate: Initial guesses for { $0)/z.
L-1
g6 = (('(((o)) "00 ( 1). and t' =0 AE.
Σ
In
Z, (5 8-1 $9.
(6)
while Kernels do not converge do Define action the action S of Equation (6) as differen- tiable function of {Ø', Φ using importance sam- pling to estimate the functions Z,.
(=1
where the functions Zz Zz-
are defined as
2/ 1- Φ0 f
/
dd' exp
_ (h²) (
Solve the inner optimization problem
2
S(( ((1) $
1 o0hn) (occ)TToo()) $2o(h()
max
cxp
with gradient ascent Perform gradient updates on feature kernels
with base case 40 = F. Ty. Then, the saddle point that dominates the distribution is
a
- 5 a Φh - Φb
S({ (([ (
{ { 6.14 = arg min max S(( $')). (+') (.')
(8)
Compute:
For a test point x, v'(x) (o(h' (x:))((()))) (z)))(h"))pt.))) ((t..))) with importance sampling and p(h, h) from Eq. (7) (see A.2).
Derivation sketch. Taking + 0o at fixed temperature 3 = F/2 and finite width N in Eq. (3) converges to a sta- tionary distribution (Kardar, 2007; Welling & Teh, 2011;
return faNax(a ()) as in Eq. (5)
5

---

Adaptive kernel predictors from feature-lcarning infinite limits of neural networks
3.2. aNTK
the A 0 (disccused in (Bordelon & Pehlevan, 2022)) and À > 0 cases, the last of which leads to the kernel predictor as in Eq. (IO). In the case of gradient flow with weight decay, since the predictor dynamics can be written by the chain rule diu deu dan, we can just track the dynamics of the output pre-ace areaactivttion
The second order of limits we study is when - oo in Eq. (3). We are interested in the infinite time limit t + of the dynamics when N - Jo, which leads to predictor
dap de
(
P
1
faNTK( -
Σ
A,
(10)
ds 4
where again Au =
ac
the output pre-
dt
K KINTK (t, (1,1)aa())
ALRSP
(14)
3sp"
a=l
activation and KINIK lim, 01,(1) 01.(12
oim Ee Gft ju (fo t) the adaptive Neural Tan-
if we suppose dspx o 0-. How to arrive at this formula can be found in Appendix F. Here, at the fixed point the out- put pre-activation is sp = λy. E, A, (t, t), which recovers the kemne predictor of Eq. (10).
gent Kernel. The gradient kernel (t, = (t)
_
(t) represents the inner products of gradient vectors (t) which are usually computed with back- (t)( = No ohe (t)
propagation.
In principle, for any given value of X, in order to have an estimate of K at convergence, one has to simulate the stochastic non-Markovian field dynamics as shown in Algorithm 2. When À > O, the contribution from initial conditions (1), e (1) (see Eq. (13)), exponentially sup- pressed at large time, while the second term of Eq. (13) contributes the most only when the system has reached con- vergence. This is not true if we switch-off the regularization X, in which case the contribution from & (4), v (1) persist late in training, since without the weight decay term, the initial conditions prevent the dynamics from converging to a fixed kernel predictor.
For squared loss and linear readout, a homogenous acti- vation function o(-), the final predictor has the form
fx. = k. NTT (x.)
+ A [TTNT
y. (11)
The factor i appears in the expressions from the weight decay contribution to the dynamics (3), since we consider degree homogeneous network as specified in 2.
In the infinite width limit N + co, the neurons in each hidden layer become independent and the o, G kernels can be computed as averages (Bordelon & Pehlevan, 2022)
(t, t) = (0(h), (t)))((()) Df f
(12)
In Fig. I we clearly demonstrate this, by comparing the pre- dictor of two-layer MLP trained on a subset of CIFARIO with the theoretical predictor of Eq. (11). For the first case when x = O, the network predictor at convergence is not the kernel predictor aNTK. Instead, when X >> O, the net- work dynamics is well-described by Eq. (11). We refer to Appendix E.1 for the case of CNNs.
(t,(t) (t)
where (.) denotes the averages over stochastic process for pairs (1 =/]=== which obey the dynamics
(t) = c (t)
+ 70 I, d'' (MC-) (1(
Algorithm 2 aNTK Regression Predictor
(t) -c)r((((t
Input: Data oo, and hyperparameters {70,Ae): }
+ +0 70 1'' dt' c Σ (0(
Output: Predictor faNtk(:)) for any test point æ E (Pie))-
Generate: Initial guesses for (Ø, " $", GL+1 11"
(t) (())(1))))
Draw S samples for random fields at initialization ef g.(() W (O))(((()))))) and v..((1) = 1 W
(13)
VN (0)g n.
2+1 (1)
VN
where (t). (t) are stochastic processes inherited from the initial conditions, which become suppressed at large
while Kernels do not converge VE E (L).Vs Vs E (S) do
Implement the non- non-Markovian dynamics of Eq. (13)
times. These equations are one way (but not the only way) to converge to a set of fixed point condition for the final fea-
Compute new "0','")
Compute KaNTK- = limm El_ wn (1, (t,t)
tures and final predictor (see Appendix F). Alg. 2 provides pseudocode for calculating the predictor in this setting.
return faNrk()) as in Eq. (11)
As specified in the Appendix E, we restrict the readout activa- tions to those with a' (s,,) * O, otherwise, the gradient signal does not backpropagate through the network, preventing convergence to kernel predictor.
Derivation sketch. A derivation to obtain the DMFT dy- namics as in Eq. (13) can be found in (Bordelon & Pehle- van, 2022). Here, we want to stress the difference between
6

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
107
10a
10
2.0
5.0
107
10.0
10
107
10
10a
EiginiLys!
Theory 1024
10a²
107
2.5
5.0
7.5
10.0
12.5
15.0
10.4
104
10a
104
Vo
(a) Kemel-Label Overlaps
(b) Overlap vs "/a
(c) Low Rank Spiked Kerelss
Figure Linear networks with whitened data are determined by set kernel-label overlap matrices. (a) The overlap variables ca
increase exponentially with & with rate that depends on Solid lines taken from Langevin dynamics on N = 1024 network. (b) The alignment of final layer et, as a function of to and exhibits three distinct scaling regimes. (c) Examples learned kernels (at each layer E) in depth ê E {L 83, 4.0 and finite width N networks compared to the N oc theory.
4. Infinitely-wide feature learning deep linear networks
To illustrate this spike effect in the learned kernels in the rich limit, we specialize to whitened input data K I. We show in Appendix A.3 that the equations get simplified even further, since now the kemnlls &' only grows in the rank-one yy direction. By defining some set of scalar variables (cz, d which are the overlaps with the label direction y Φby = Gt and y & Ō2 we find
Deep linear networks (h) = h) provide a simpler frame- work for analysis than their nonlinear counterparts (Saxe al., 2014; Advani & Saxe, 2017; Arora et al., 2019a; Aitchison, 2020; Li & Sompolinsky, 2021; Jacot et al., 2022; Zavatone- Veth et al., 2022a), yet they still converge to non-trivial feature aligned solutions. In deep linear net- works, preactivations remain Gaussian at each layer when P= x1) for the Central Limit Theorem (CLT), and this greatly simplifies the saddle point equations to algebraic formulas which close in terms of the kernels for both aNBK and aNTK theories. This means that
JOL ct - (B-1 c1,
Aí {L}
(16)
which means that there is an exponential dependence of the overlap on the layer index (full derivation can be found in the Appendix) as Fig. 2(a) shows. From Eq. (16) we derive the scalings for lazy, large depth, and large feature strength limits. For the last layer overlap cr. these
solve for the adaptive kernels in both cases without any refined sampling strategy. The deep linear case of gradient flow dynamics can be found in (Bordelon & Pehlevan, 2022). Here, we report the solution to the saddle point equations that define the kernels in the Peature-learning Bayesian setting, special- izing to the regression problem (the generic loss case can be found in Appendix A)
~ + LnD
2 +
21/(L+1) CL ~ To
To - 0o, L fixed
(17)
Lid
~ In(Lnd)
L 00, To fixed
f-1
I
=0
vL= L,...,L
11-1
Ar-1
which closely match the theory in their respective regimes plotted in Fig. 2(b). In Fig. 2(c) we show examples of learned kernels for a L 8 network and no 4.0 matching the finite width N = 1028 network trained with Langevin dynamics.
e+1
[
$c41)
=0
Ae 1,
L
de
de
$
70
2
I
TL
-1
yy
1
44
-1
0.
d
8
d
(15)
5. Numerical Results
Here, as a consistency check, it is easy to see that in the lazy limit "o + 0 the dual kernels &' = Vl {L}, and consequence all the kernels te will stay equal to the data covariance matrix T consistent with lazy learning. However, for the rich regime where To > 0 the &' kernels do not vanish and alter the fixed point kernels &f with target dependent information in the form of low rank spikes (by definition of 44
Two-layer MLPs. In Fig. 3(a) we compare test losses of lazy feature learning kernels for two-layer MLP trained on P subset of two classes of CIFAR 10 in regression task. The green curve is the performance of NNGPK, Orange is the aNTK, red the aNBK. There is a gap in performance between the lazy predictors and the adaptive feature learning predictors. However, when sample size P is small, feature
7

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
ANIK
Lazy
K5055
K-
10-20
MNGPK
ANTX
0.35
<NBC
Test Loss
ANDK
0.25
Saueeria
0.15
200
400
600
800
1000
P
(a)
(b)
(c)
Figure 3. Feature learning theories outperform lazy predictors for a two-layer MLP trained with Squared Loss (SI.) on two classes of CIFARIO (airplane automobile). (a) Test losses as function of sample size Solid lines refer to theories, dashed lines to numerical simulations on a N = 5000 network. Blue is the NNGP lazy predictor, green is the deterministic NNGPK kemel predictor; orange is aNTK with feature learning strength to 0.3; red aNBK predictor with the same 7u. (b) Kon-Gaussian pre-activation densities as a function of "to for (top) aNTK and (boftom) aNBK. Black dashed lines are theoretical predictions. (c) Learned feature kernels of the adaptive theories closely match their relative finite width network trainings and evolve with ?oo
learning in a data-limited scenario can let the model to overfit on test points and lazy learning can be beneficial a small window of
1=1
t = 2
t=
==
aNBK
In this plot, also include the Neural Network Gaus- sian Process in blue (Neal, 1995; Lee et al., 2018; de G. Matthews et al., 2018), where for a number of pat- terns P the solution space is sampled from the posterior of Eq. (9) by taking = x (1). Here the mean predictor is
1=1
1= 1=4
Langevin
equivalent to the NNGPK predictor, however there is also a variance term, which comes from the fact that we are averaging over all possible random weights.
In the rich scenarios, derive the preactivation distribu- tions p(h) function of "oo (Fig. 3(b) top and bottom) at convergence. At initialization, p(h) follows N(0, z)). However, as learning proceeds and features are learned, the densities accumulate non- non-Gaussian contributions which are identified by our theories (e.g. Eq. (7)). Fig. 3(c) shows that there is clustering of 100 data points by category in the feature space defined by the adaptive kernels.
Figure 4. Theory vs empirical kernels at each layer for SHI MLP with 6(h) = lanh(h)) learning / = 50 patterns of MNIST with y = {+1)² labels. Alignments AA( {
1: ff
| cpp
between theory and empirical kernels are A($* $ (mp) 97%. A(# x2n)) 81%, A( $ 73% and A( 2 89%.
non-Gaussian single-site density of Eq. (7) from the layer index L.
Deep MLPs. In Fig. 4 we show simulations of a Bayesian == Tanh MLP and compare to our infinite-width pre- dictors. In order to study how feature learning propagates through depth, we plot the theoretical adaptive kernels of the aNBK theory versus the empirical kernels of Langevin dynamics thermalization as function of the layer in- dex l. Coherently with the deep linear case in 4, the last layer feature kernel aligns first with the labels, since the clustering of features by class is more evident in this case. There a small difference between theory and experiments because solving the min max problem when L > 1 makes the convergence of harder, given the dependency of the
5.1. CNNs
DMFT for infinite-width CNNs under gradient- was pre- viously derived (Bordelon & Pehlevan, 2022), which we solve numerically here for the first time. In Fig. show comparisons of DMFT kernel predictors convergence for two-layer MIP and two-layer CNN with kernel size k = 8 and stride 8. Black dashed curves, which are the theories, closely match the full colored lines, which are the network predictors on = 1028 width networks. CNN outperform the MLP at large sample size P for the same "oo-
8

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
Future work could focus on reducing solver computational costs by developing more efficient optimization techniques.
DM ye yon2 MPP
6.20
6.5
Test Loss
0.00
Acknowledgements
We thank Enrico M. Malatesta and Hugo Cui for useful and engaging discussions. B.B. supported by a Google PhD Fellowship and NSF CAREER Award IIS-2239780. C.L., is supported by DARPA Award DIAL-FP-OSH and
6.5
600
1000
P
The William F. Milton Fund from Harvard University. C.P. NSF CAREER Award IIS-2239780, DARPA Award DIAL-FP-038 a Sloan Research Fellowship. and The William F. Milton Fund from Harvard University. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foun- dation establish the Kempner Institute for the Study of Natural and Artificial Intelligence.
is supported by NSF grant AMS-21341
Figure 5. Test Loss function of sample size P for DMFT theorics at convergence: two-layer MLP vs two-layer CNN trained on P animate//nanimate data CIFARIO. Dashed lines are for theory, full-colored curves for empirical kernels.
See Appendix E.1 and Figure 5 for simulations of adap- tive convolutional kernels derived from the feature-learning setting.
Impact Statement
This paper presents work whose goal is to advance the the- oretical understanding of deep neural networks. There are many computationallalgoritmie consequences of our work, none which we feel must be specifically highlighted here.
6. Discussion
In this paper, we develop a theory of non- non-parametric fea- ture kernel predictors for MIP and CNN architectures in yP/mean-field parameterization. By analyzing gradient flow dynamics with weight decay and/or white noise, we iden- tify two distinct infinite-width adaptive kernel predictors: aNBK, representing a Bayesian description of DNNs, and aNTK, corresponding to the fixed points of gradient flow with weight decay (Bordelon & Pehlevan, 2022). Unlike static NNGP (Neal, 1995; Lee et al., 2018) and NTK (Ja- al., 2020) predictors, our kernels adapt data. The feature learning strength, controlled by 70. recovers lazy training when "to + 0 and enables richer representations for 7oo > 0 (Bordelon et al., 2024). We study their impact across architectures and benchmark tasks on real datasets.
References
Advani, M. S. and Saxe, A. M. High-dimensional dynamics of generalization error in neural networks, 2017. URL https: /arxiv////ravvv//a org hbspgg/bs/11100 03667.
Aitchison, L.
Why bigger not always better: on finite and infinite neural networks, 2020. URL htRps:t:://aae org/ abs / 91 01..... 08013.
Aiudi, R., Pacelli, R., Vezzani, A., Burioni, R., and Ro- tondo, P. Local kernel renormalization as a mecha- nism for feature learning overparametrriee convolu- tional neural networks. Nature Communications, 16. 2023. URL https::/aap semanticscholar. org/ CoopooII::20011::::
We also analyze infinitely wide deep linear networks in the feature-learning regime, where the saddle point equations simplify. Assuming white data covariance matrix, the order parameters reduce to scalar overlaps ce between ker- nels and labels at each layer f, exhibiting an exponential dependence on L. This implies that for fixed Yo. deep linear networks align last-layer kernels first and propagate align- ment backward. We derive scaling laws for these overlaps the lazy, large-width, and large-depth regimes.
Arora, S., Cohen, N., Golowich, N., and Hu, W. A conver- gence analysis of gradient descent for deep linear neural networks, 2019a. URL https: http///xvvvvv//////////ee org/abs/ 1810. 02281.
Arora, S., Du, S. S., Hu, W., Li, Salakhutdinov, R., and Wang, R. On exact computation with an infinitely wide neural net, 2019b. URt.ht:ss:/arrxiv.rrr
Our numerical results show that our adaptive kernels outper- form NNGP and NTK at large sample size and match the performance of trained NN in the feature-learning regime even in moderate widths (e.g. N = 5000). Our theory pre- diets non-Gaussian pr-acctivatio densities at convergence and data-clustered feature kernels, whose alignment with label covariance increases with Y0-
1904. 11955.
Arous, G. B., Dembo, A., and Guionnet, A. Cugliandolo- kurchan equations for dynamics of spin-glasses, 2004. URL https: / aaariiivvv. org/ abs/ abs/math/ 0409273.

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
Atanasov, A., Bordelon, B., and Pehlevan, C. Neural networks kernel learnrr:: The silent alignment ef- fect, 2021. URL http:ps//axxxw..rg////// 00034.
Cui, H., Krzakala, F,, and Zdeborová, L. Bayes-optimal learning of deep random networks of extensive-widtt In International Conference Machine Learning, pp. PMLR, 2023.
Baglioni, P., Pacelli, R., Aiudi, R., Di Renzo, F., Vezzani, A., Burioni, R., and Rotondo, P. Predictive power of bayesian effective action for fully connected one hidden layer neural networks in the proportional limit. Phys. Rev. Lett., 133:027301, Jul 2024, doi: 10. 103/PhysRevLet. 133.027301. URL https: /link. aps owg/ doi/ 133. 027301
De Dominicis, C. Dynamics substitute for replicas in systems with quenched random impurities. Phys. Rev. B. 18:4913-4919 Nov 1978. doi: 10.1 10.1103/PhysRev.. 18.4913. https: h///lin..ar. wtwps.oog//o///////////
19913.
10. 1003 1tt.000//yyyseeeese
de G. Matthews, A. G., Rowland, M.. Hron, J.. Turner, R. E., and Ghahramani, Z. Gaussian process behaviour in wide deep neural networks, 2018. URL pictps://arxo.
Bassetti, F., Gherardi, M., Ingrosso, A., Pastore, M., and Rotondo, P. Feature learning in finite-width bayesian deep linear networks with multiple outputs and convo- lutional layers, 2024. URL https: /arxiv.org/ 03260.
org/ abs 1804 11271.
Domingos, Every model learned by gradient descent is approximately kernel machine. arXiv preprint arXiv:2012.00152,, 2020.
Bordelon, B. and Pehlevan, C. Self-consistent dynam- ical field theory of kernel evolution in wide neural abs/
Fischer, K., Lindner, J., Dahmen, D., Ringel, Z., Krämer, M., and Helias, M. Critical feature learning in deep neural networks, 2024. URL https: http//axvvvv./.// abs 2405. 10761.
networks, 2022. URL 2205.09653.
Bordelon, B. and Pehlevan, C. Dynamics of finite width kernel and prediction fluctuations in mean field neural networks. Advances in Neural Information Processing Systems, 36, 2024.
Geiger, M., Spigler, S., Jacot, A., and Wyart, M. Disentan- gling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experi- ment, 2020(11): 113301, November 2020. ISSN 1742- 5468. doi: 10. 10.1/88/17 5445//2b2255///bbceee //dx.doi... org. 1ttp:////1111001/122---55550 55c-/22/-566////ccdeeee abc4de.
Bordelon, B., Atanasov, A., and Pehlevan, C. How feature learning can improve neural scaling laws, 2024. URL https:: //axiiv. org/ abs//b////// 220. 17858.
URL http:
Hanin, and Zlokapa, Bayesian interpolation with deep linear networks. Proceedings of the National Academy of Sciences, 120(23):e2301345110) 2023.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations Python+NumPy programs, 2018. URL. http:// http:///ithus com/ jaijjax-m.. m0mmmmm/10 jax.
Hestness, J., Narang, S., Ardalani, N., Diamos, G., Jun, H., Kianinejad, H., Patwary, M. M. A., Yang, Y., and Zhou, Y. Deep learning scaling predictable, empirically, arXiv preprint arXiv: 2017.
Chizat, L., Oyallon, E., and Bach, F. On lazy training in differentiable programming, 2020. URL https:// arxiv. oww. abs/://bb/aaae wtw.l..org/abs/1112222 /wtp://11......111
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arX1v:2203.15556 2022.
Cho, Y. and Saul, I.. Kernel methods for deep learning. In Bengio, Y. Schuurmans, D., Lafferty, J., Williams, C., and Culotta, A. (eds.). Advances in Neural Information Processing Systems, volume 22. Curran Associates, Inc., 2009a. URL ptttps::/prrcceeeiiee ree oc cc/ppppee_fflle/ppppprr///////// [r9ieeeeaafeaa577927778800222220 pdf.
neurips.
Ingrosso, A., Pacelli, R., Rotondo, P., and Gerace, F. Sta- tistical mechanics of transfer learning in fully connected networks in the proportional limit. Physical Review Let- lers, 134(17):177301 2025.
Cho, Y. and Saul, L. Kernel methods for deep learning. Advances in neural information processing systems, 22, 2009b.
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks, 2020. URL https://a www.arwwww.ww.....././/ee.. 1106.. 07572.
10

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
Jacot, A., Ged, F., Şimşek, B., Hongler, C., and Gabriel, F. Saddle-to-saddle dynamics in deep linear net- works: Small initialization training, symmetry, and spar- org/ sity, 2022. URL https: http://awww....ooaa/////0000 15933.
1806579115. URL htwphhttt//wwwppsssssss. /otp
1ass/1..103/paas htt.://///paa...181599999 1806579115.
Mingard, C., Valle-Pérez, G., Skalse, J., and Louis, A. A. Is sgd a bayesian sampler? well, almost, 2020. URL https: aatpaaixvvv.o....... wwt.///gaaa.//ooooee
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models.
Naveh, G. and Ringel, Z. A self consistent theory of gaus- sian processes captures feature learning effects finite cnns. Advances Neural Information Processing Sys- tems, 34:21352 -23644 2021.
arXiv preprint arXiv:2001.08360 2020.
Kardar, M. Statistical physics of particles. Cambridge
University Press, 2007.
Naveh, G., Ben David, O., Sompolinsky, H., and Ringel, Z. Predicting the outputs of finite deep neural networks trained with noisy gradients. Physical Review E, 104 (6), December 2021. ISSN 2470-0053. doi: 10.1103/ Soysrree.104.06440 10. 1..11033 www1111//hheeeeee. 104. 064301.
Lee, J., Bahri, Y., Novak, R., Schoenholz, S. S., Pennington, J., and Sohl-Dickstein, J. Deep neural networks as gaus- sian processes, 2018. URL htt.ttss./acxxiv.ooo /ttpxv..xxi.g/////// abs/1711.00m 00165.
URL http: / dox.dooi.rrr///
Lee, J., Xiao, L., Schoenholz, S., Bahri, Y., Novak, R., Sohl-Dickstein, J., and Pennington, J. Wide neu- ral networks of any depth evolve linear models un- der gradient descent * Journal of Statistical Mechan- ics: Theory and Experiment, 2020(12): 124002, Decem- doi: 10. 10.1448//742-54688 abc62b. URL http:: //o:://dxxxie doi. www.oogg//.... /10. 1088/ abc6bb.
Neal, R. M. Bayesian learning for neural networks. PhD thesis, CAN, 1995. AAINN02676.
Novak, R., Bahri, Y., Abolafia, D. A., Pennington, J., and Sohl-Dickstein, J. Sensitivity and generalization in neural networks: an empirical study, 2018. URL https: // arxiv. www./as///////////////./../e..a/ abs /10...11...00 08760.
ber 2020. ISSN 1742-5468.
Pacelli, R., Ariosto, S., Pastore, M., Ginelli, F., Gherardi, M., and Rotondo, P. A statistical mechanics frame- work for bayesian deep neural networks beyond the infinite-width limit. Nature Machine Intelligence, 5(12): 1497-1507, December 2023. ISSN 2522-5839, doi:
Lewkowyez, A. and Gur-Ari, On the training dynamics of deep networks with 1.2 regularization. Advances in Neural Information Processing Systems, 33:4790-4799, 2020.
Li Q. and Sompolinsky, H. Statistical mechanics of deep linear neural networks: The backpropugating ker- nel remormalization. Physical Review X, 11(3), Septem-
10. 10.1088/42256-023-00006 URL htw.thttpt///e /ddiidi//dxd..... org 10..0..133//22220 -20-000---00--000--0
Radhakrishnan, A., Beaglehole, D., Pandit, P., and Belkin, M. Mechanism of feature learning in deep fully con- nected networks and kernel machines that recursively learn features. arXiv preprint arXiv: 13881, 2022.
ber 2021. ISSN 2160-3308. doi: 10.1103/physrevx. 11.031059. URL http::/tt//// htt.:ttpp//xdoooooo http//xxxooooo///////o//////ooot o0o.i..rg/10.1111 htt.::::/....oi..o///1...111111. PhysRevX. 11. 031059.
Li, Q. and Sompolinsky, H. Globally gated deep linear networks. Advances in Neural Information Processing Systems, 35:34789-34800 2022.
Rasmussen, C. E. and Williams, C. K. Gaussian Processes for Machine Learning. The MIT Press, 2006.
Mallinar, N., Beaglehole, D., Zhu, I.., Radhakrishnan, A., Pandit, P.. and Belkin, M. Emergence in non-neural mod- els: grokking modular arithmetic via average gradient outer product. arXiv preprini urXiv:2407.201 2024.
Roberts, D. A., Yaida, S., and Hanin, B. The principles of deep learning theory, volume 46. Cambridge University Press Cambridge, MA, USA, 2022.
Rotskoff, G. and Vanden- E. Trainability and accu- racy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 75(9):1889-1933) 2022.
Martin, P. C., Siggia, E. D., and Rose, H. Sta- tistical dynamics of classical systems. Phys. Rev. A, 4:423-437, Jul 1973. doi: 10. 10.1103/PhysRev.. 8.423. URL https::///// /ttes://iikk.ee aps.kkapps..ooo orwpo//oo/o//1/ 1103/PhysRevA. 8. 423...
Rubin, N., Ringel, Z., Seroussi, I., and Helias, M. A unified approach to feature learning in bayesian neu- ral networks. In High-dimensional Learning Dynam- ics 2024: The Emergence of Structure and Reasoning, 2024a. URL https: http://pppppeeieeeeee net/ forum? id=7m09.22V2R..
Mei, S., Montanari, A., and Nguyen, P.-. A mean field view of the landscape of two-layer neural net- works. Proceedings of the National Academy of Sei- ences, 115(33)):77665555)) 2018. doi: 10.1073/pnas.

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
Rubin, N., Seroussi, I., and Ringel, Z. Grokking as first der phase transition in layer networks. In The Twelfth International Conference on Learning Representations, 2024b. URL httpt:://// http://///ppeeeeeeeeeeeee hete://w.eeeeeeeeeeeeeeeeeeeeeee Eorum? id=3ROGsTX313
Yang, G. and Hu, E. J. Feature learning in infinite-width neural networks, 2022. URL URLhttps://arxive
Yang, G., Hu, J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, J. Tensor programs Tuning large neural networks via zero-shot hyperparameter transfer, 2022. URL https: /arxiv. org/ abs /abb/22033 634
Saxe, A. M., McClelland, J. L., and Ganguli, S. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks, 2014. URL https: //arxiv.
org/aba abs/ 1312. 6120.
Zavatone-Vethh A., Canatar, A., Ruben, B. S., and Pehlevan, C. Asymptoties of representation learning in finite bayesian neural networks*. Journal of Sta- tistical Mechanics: Theory and Experiment, 2022(11): 114008, November 2022a. ISSN 1742-5468. doi: 10. 10.1088/1742-5468/aceee URL http: dx.doi.. 0xg/10.1088/1442-55555 58000///817772-55/8888889999
Scholkopt, B. and Smola, A. J. Learning with Kernels: Support Vector Machines, Regularization, Optimization, and Beyond. MIT Press, Cambridge, MA, USA, 2001. ISBN 0262194759.
Seroussi, I,, Naveh, G., and Ringel, Z. Separation of scales and a thermodynamic description of feature learning in some cnns. Nature Communications, 14(1):908, 2023.
Zavatone-Veth, A., Tong, W. I.,, and Pehlevan, C. Con- trasting random and learned features in deep bayesian linear regression. Physical Review E, 105(6):06411, 2022b.
Shan, H., Li, Q. and Sompolinsky, H. Order parameters and phase transitions of continual learning in deep neural networks. arXiv preprini wrXiv:2407.10315 2024.
Sompolinsky, H. and Zippelius, A. Dynamic the- ory of the spin-glass phase. Phys. Rev. Lett., 47: 359-362, Aug 1981. doi: Lo. 47.359. URL https://link aps. pass.erg//oo/11 1103/PhysRevLLet..4..33
van Meegen, A. and Sompolinsky, H. Coding schemes in neural networks learning classification tasks, 2024. URL https:/ arxiv. org/ abs/ 2406.1666666 16689.
Vyas, N., Bansal, Y., and Nakkiran, P. Limitations of the ntk for understanding generalization in deep learning. arXiv preprint arXiv:2206. 100 2022.
Vyas. N., Atanasov, A., Bordelon, B., Morwani, D., Sainathan, S., and Pehlevan, C. Feature-learning net- works are consistent across widths at realistic scales, 2023. URL https: / aaaiiviiivv..//.v///aaaaa/aaaavv org/ abs 2305. 18411.
Wei, A., Hu, W., and Steinhardt, J. More than a toy: Random matrix models predict how real-world neural representa- tions generalize, 2022. URL. https::/p//........... /ttp///xxvvv.///// abs/2203. 06176.
Welling, M. and Teh, Y. W. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th International Conference International Conference Machine Learning, ICML'11 1, pp. 681 1688, Madison, WI, USA, 2011. Omnipress. ISBN 90899778330 450306195.
Yang, A. X., Robeyns, M., Milsom, E., Anson, B., Schoots, N., and Aitchison, L. A theory of representation learning gives a deep generalisation of kernel methods. In Inter- national Conference on Machine Learning, pp. 39380- 39415. PMLR, 2023.
12

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
A. Multi-layer deep Bayesian MLPs
As mentioned in the main text, we would like to study feature learning when the solution space sampled from a posterior that is Gibbs distribution with a likelihood ((: D) and Gaussian prior 2 ||01 and by assuming width N₁ = NA e {L}. We can do this computation for generic loss and non-linear activation functions as Eq. (2) shows. The representer theorem for Bayesian network is indeed independent the activation choice. Here, the posterior takes the form
L
DD) = Z exp - ((;2
Σ 8=
0 11?
(18)
2
being 0 = Vec (we ,y.,,,,,, w( the collection of weights, D= { Ve the dataset with P patterns, and / = 1 the temperature inverse. Again, we are interested in the infinitely overparameterized limit where + 00, P O2 (1). Here, when 3 - xc the posterior becomes the uniform distribution over the set of global minimizers 0 E orggiin ((0). In
T
this setting, one needs to rescale the loss function L - 72C with - 7ov N in order to avoid for the Gaussian prior to dominate over the likelihood when N + oc, suppressing any interaction with the learning task D. From the normalization factor in Eq. (18), the partition function reads
L Z = /
2 I (-0" Eizo 2 1|1
a E=
(19)
and since we consider the dataset as fixed, we wish to integrate out the weights and move to a description in the space of representations. This can be done by simply enforcing the definitions of Eq. (20) through the integral representations of some Dirac-delta functions
/
II M.R
2+1 dh
dsp
II n.l
h f+1
W' o(h())
II (sH P
( o(h))) )
(20)
VN
OoNN
where A= = diag ( A1I, ,I, A) getting
(on w(ck)) Nxc
ds d3
( En Ene -1t-ittt-4410. Ee
Naatt E, 8p3p
2 =
2x
2n ()oNN
E=U
(wiatt wontt Ae
Enn &
$ No0 En (((""
x e
do d
d3
Eno Ee $/ $100 NN En snbu+ 4 o s, -NN E, (()))
X
2m N N t=1
N)
x N Et in f 1.60ll
(21)
In the last expression, we introduced the adaptive feature kernels as
Φf pv 1 N (h()) ((h))
(22)
and, again, we enforced their definitions in Z with some conjugated variables - Both ( $1 will become deterministic quantities in the N + oo limit.
The single-site density Eq. (21) is given by
dh dh,
+ En Nuhyy
Env by 2e-1 bo- 1 Eno (l..). o(hb)
Z
2T
11
n1 dh
- E Ay
(pe-1
hw - ((hn) d(he))
(23)
24-1
(pe-1)
2n det
Ae
13

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
At each layer, this decouples over the neuron index because we supposed the hidden layers having the same width dimension N for é ..... It and represents the normalization factor of non-Gaussian pre-activation density distribution p(h))
where the non-Gaussian part is proportional to & while the Gaussian contribution has covariance that is the feature kernel at previous layer Ø²-1
In the large N limit, we can collect the partition function of Eq. (21) in the compact form
dob dbe 2mN-1
ds ds NS(( $ (
Z
(24)
(7
2T
JAU
being S(( $L )im) the intensive Bayesian action
5 5= 3
ΣΣ
4 D Dú 7o Σ
SySp
SH
2 3 Sy +
10
Σ
((s))
Em
(25)
d
At infinite width N, this partition function exponentially dominated by the saddle points of S. We thus identify the kernels that make S locally stationary (õS 0) by the equations
as
0
Ae = I,...,L L
(26a)
O
as
AA = 1, 1
(26b)
OΦe
as
(26c)
1.L
0s
0
(26d)
06p
as
(26e)
0sp
where the last two saddle points fix the output pre-activation given the loss function L. If we explicitly write down the partial derivatives, we get
Φ = ((h)T) (LE
(27a)
Φ (D')
-e((0))-111 11 ( VE E [ 1}
(27b)
4
X 88
(27c)
dL
8 870 0s
(27d)
8 =
(27e)
70d1
from which we get a kernel predictor on unseen test point æ, since
f(x) = . ( Σ
P
A
A (2)
(28)
A p=1
being A, 0C 8s, the pattern error signal for that given loss.
A.1. Regression problem
In this specific case, where the form of the loss function is known -=
EP- (vn - fa) and the readout is linear, i.c. f(sp) 8p Vμ E (P), we can integrate over the output pre-activations and conjugated parameter (sp, Sp). After
14

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
integrating, we obtain
L
d
d
Z
per
4 ((m)
2T
JAM
t=1
(29)
L
de
- 2i 1+ ((00))+/((((+)
y g-ND 2 22
2x
MA É=1
where the important quantity to be extremized in the limit N + oc is the intensive action
S( 6 - (
( I
+
2 y
T
L-1 Σ
2n
15
(30)
t=1
Here, the saddle points which render the action locally stationary 65 = with respect to these 2L matrix order parameters can be collected as
Φ (0(h') (oh))))
ve=1....,4
(31a)
Φ = (D')-
(
1e= 1-1
(31b)
$ T
$b 21 yy ( x
1 $4
)
(31c)
d
X
being 4. m, my the data matrix covariance in this notation.
Notice that the last layer dual's kernel &% vanishes in the lazy limit "o - 0 and so do all the dual kernels at previous layers dh puv = 0, while for non-negligible to we see that the each hidden layer features are non-Gaussian from Eq. (23). Details on the numerical implementation of the numerical solver for Eqs. (23) can be found in Sec. G.2.
In general, solving the min-max problem for deep networks is hard, because solving max (S is the action defined +
in Algorithm 1) means to find some values for the dual kernels ( which tilt the non-Gaussian measure in Eq. (7) each layer E. In Fig. 4 of the main text we solved the theory for Algorithm by initializing the theoretical kernels {Φ, } with lazy guess explained in Sec. G.2. However,
easier strategy for convergence is to initialize the solver Algorithm with the empirical NN kernels { { Jé= at convergence obtained from Langevin simulations and perturbed with multiplicative Gaussian noise. In this way, one only needs to solve for the dual variables. This warm start allows 4'y faster convergence. This is what we did to produce Fig. 7. Here, we plot the kernel alignment A(Φ', yy ") Tuw T' which is the cosine similarity between the kernel &² and the label covariance yy" T each layer. We show that by increasing feature strength Yuu the alignments increase. As expected, the last layer kernel aligns first with the labels as "Yo increases, followed by the previous ones. Fig. 7(b) shows that the perturbed kernels of the theory retrieve the final NN kernels as locally stable fixed points.
In Fig. we show that the non Gaussian preactivation distribution, depending on the activation function chosen, can develop a multimodal profile, and so a Mexican-hat profile for high values of feature learning strength Yo-
A.2. Generalization error
Knowing the posterior distribution makes it easy to compute the test error on a new (unseen) example (æo. Mo), which is defined as
(=o Vo)
((yo fo (x 2)) (~)P)
(32)
where the sampling measure corresponds to Eq. (18). Here, we can include with an index μ = 0 the test pattern contribution, and just compute
(x vo)
d d d
P d
Enozo Ete N. fentt 2x
No I @ &naur 237,
No. a (o.
(8o)²
x
Z
2x N
2
Nv) -1
1==0 1
A=0
Epp 4
N2 Sp-o
Puru- ENN
>ーー (v" P x (yo =0)²
X e XL
(33)
15

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
0.5
0.5
2.0
2NBK
p(h)
012
0.0
(a) Kernel for ((h) tanh(h).
(b) Preactivation density for 6(h) tanh(h).
Figure 6. (a) Bayesian theory (aNBK) and empirical (Langeevin) adaptive kernels with feature learning strength and for P 50 patterns 0/1 classes of MNIST. the Gram matrix data. (b) preactivation distribution is general ann-ausssian at value feature strength To. For activation ((h) = tanh(h) show that phh) can develop Mexican-hat profile.
[=1
Alignments
Theory
021
EXp
02
0.8
14
V0
(a) Kernel alignments
(b) Hidden layer kerelss
Figure 7. Bayesian L 5 ReLU MLP trained on P = 1000 data CIFARIO. Colored curve are experiments, dashed lines are the predictors calculated from Algorithm 1. (a) Kernel alignments A(Φ yy ) at each layer feature strength. (b) Theory empirical kernels at each layer. Here, theory refers to the intermediate layer kernels calculated via Algorithm 1.
where the single-site action contains now all the possible interactions with the test point at each layer é in the test-test kernel Φ (²))))²) and the test-train kernel d = (o(ho))(h,,))
$
/
P
dhg i Di-o
bn &1 x hw- hu- f h3
xx-1
41 -E
hoh
2n
(34)
===
Die ((oh)))((h))
16

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
Exploiting the saddle point equations, we realize that the dual kernels concerning the test point $ 600 0, and that
Be = (0(W))(())))
vC= 1,...,I. I. (itttin5t1)
(35a)
He m
=t= = -....,,--I
(35b)
2d
BL
8188 0
(35c)
2dL
(1++
4e 1 ..... L
(35d)
d
$
P =0
(35e)
1
Σ
T
(35f)
ÀL
8p.
Up
8p)
(35g)
8p Nr.
Tbe Su
(35h)
This allows to rewrite the single site density in much simpler form, where the non nan-Gassiian contribution includes just the train points, while the Gaussian part has (P' + 1) x (P + 1) covariance matrix Φ-
P00
T
4
p.
II dhp
1 Epuzo hp (8' hw- Eiezi @hhn) (o o(ho)
2 2=
(36)
(zi-1 A22
2T det
This means that once we solved for Eq. (31) we can marginalize Eq. (36) to get p(h and hence the test-train vector kernel D The test error expression is
Eg (20o vo)
(10
1
L d
I
y)
(37)
d
3
meaning that the predictor f(x) = - (xx) Eqs. (31).
4
+
3
is again kernel predictor, with adaptive kernels from
alternative way obtain the predictor in the specific regression setting from Eq. (28) and when (s)) is by solving for A,
8L. In this case
Ozp
3 EE B
D Up fn.
Uu
(38)
and by solving this equation for E R.P we obtain
A
y.
(39)
Lastly, combining this with the original Eq. (28), we get
f(x) = T
I+
(40)
X
as is it in the main text Eq. (5).
17

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
A.3. Deep linear case
In the deep linear case (hh)) h², the action of Eq. (30) gets simplified because the single-site density is now Gaussian, and this leads to
S( 6
1
En((e
Σ
70 y +
TTT
I
+
-1
y
1 _
In det (I+ Φ
$()
(41)
d
2 8=1
x-11
with the following saddle point equations
$
$
Φ
1
=0
Ae L...., L
(42a)
1--1
Ae-1
$
/++1
[
T
-1
=0 0
Ae = 1 1...., --1
(42b)
Xe
Xe
T + 70
I
$4
-1
4 -1
yy
A = 0.
(42c)
AL
B
X
3
In principle, this is a closed set of equations, which can be iteratively solved as mentioned in Sec. 4. If we choose input data that are whitened, with g0 K and label norm ly = 1, the equations can be simplified even further, since feature kernels at each layer &e only evolve the rank-one direction yy" T, This allows to define the variables {ce, &e} which are the overlaps with the label direction y
y 2 b'y=ce y y E In this setting, the reduced saddle point equations become
(43)
Ci
C1 = 1+ cE+1 = 1 ce C2+1
(44)
ce-1
ē = (0 + cL) Eτ = 1 ci ce+1
(45)
which are 2L scalar equations for the overlaps c at each layer. We note the following conservation
7P2z
cτêe c1+140+1 = (8- + cc)2
= =((e))
(46)
which implies that
Ci
C4+1 c (1 (1-x(er)) = - -1-x(cr))) cr.
(47)
ct+1 + ceX(eL)
Since we have C1 T4x//01
C1 =1=1-((c)) hence we find
(1-
(48)
which means that here an exponential dependence of the overlap on layer index. In practice, we can solve Eq. (48) for the last layer overlap er since x(ez.) and then move backward in computing all the previous layer overlaps.
We can also extract the following small or small L asymptotics. Precisely, when L is fixed and 70 0 we recover perturbative feature learning regime
L3² el ~ ~ , +
CL 1+ CL
(49)
where correction are 0(7²L) Similarly, at large with fixed L. which stands for shallow but very rich regime, the overlaps scale as
c= c=(1+>cc = e 2wr22//100 21/(1++1
(50)
while, alternatively the large L asymptotics for a very deep network have the form
L
2+1/L = = = clnc IA²
(51)
1n(
which we show to be predictive in 3 of the main text.
18

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
NNGP aNBK Expt
y=
0.5
10
y-
2=
0.20
5=
p(h)
p(h)
10
Y- 10.0
0.15
Y= 20.0
10
10
-5
0
5
10
15
20
h
h
(a)
(b)
Figure 8. (a) Bayesian two-layer MLP trained whitened covariance matrix $ = I on P = 4 train points. Feature learning (7o leads to a non-Gaussian pre-activation distributions. Black-dashed curve the lazy NNGP when (0, Φ"); red curve is the aNBK theory when is sampled from Eq. 7; blue is the empirical pre-activation distribution of N 5000 network trained in the rich regime. (b) Non-gaussian pre-activation distrihution as a function of feature learning strength 7y for Bayesian 2-layer ML. trained with Squared Error (SE) on classes MNIST dataset. Here sample size 100.
A.4. Non-gaussian pre-activation density
In the non- non-linear case, as Eq. (23) shows, when "o > O the pre-activation density at each layer is Con-Gausiin,, with "o entering in the saddle point equation for the dual kernel at the last layer & (see Eq. (31)). This means that, once we have {$p, $ b } from the solver (Alg 1), we can evaluate Eq. (23) with importance sampling and compute p(h,) for given pattern μu. Fig. A.4(a) we show that, while in the lazy regime where no feature learning enters, the hidden layer pre-activation of the NNGP predictors are Gaussian, this is not the case in our setting.
A.5. Perturbative approximation
In the %o - limit, we recover the static kernels of NNGP predictor (Neal, 1995; Lee et al., 2018). Corrections to this lazy limit can extracted at small but finite Tt- In order to do so, we can expand each macroscopic variable 4(%) in power series of 7o. such q (0) (2) 26 + .... and compute the corrections up O(7)). First of all, notice that leading order in o
$ -
- @ + --1 yy (Øf ---
(52)
where we set each d for clarity of notation. For each dual kernel at previous layer L = 1,...,L L - we have instead a recursion
1 In Z( $(+1)
(53)
Oep
where non- perturbativell
dh
exp
(-2 ((
1
(1))
-1 [hh) ((1)
(54)
2
2
02
Oh exp
( f+1
(0h))
(55)
2
(-bo(h))sttdd
exp
19

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
being Oo the Gaussian average with covariance $ With little bit of algebra the numerator can be written as
exp (-2o(h)(o(1)))
(-
((h (-200))))())) exp
Bma ( 1 (-oom))ooch))
exp
Oh
Thu
Σ
o pn Pus
a6
(klhy)d(ha))ss exp (-2o(n)dushh))
Suv
Σ
(((h)))(hh exp (- - (-20(1)
(56)
pa.
Under the leading order approximation, we find the following relationship between successive layers
82 Thah T 6(h) T (5c110(h)
te
(57)
The entries of this Hessian matrix can be computed in terms of derivatives of the activation function
0 a Σ ((ha)
(h( + 6(h))ō( ( ]vvo]
(58)
ohy Oh aB
=2(((()))(hw)) Dn 26 (Ö(bx) Σ _o(hs) $w0
)
(59)
all of these Gaussians can be evaluated at the unperturbed NNGP kernels Gaussian densities. Once these &' matrices have been computed, the &f matrices can be asymptotically approximated
1 (()) (0(h) क5 (d())TT'(hh)) ((i))) #6 (o(()))TT doo())
&
(60)
2
which agree with the perturbative treatment of (Zavatone- Veth et al., 2022a) obtained in the NTK parameterization which is similar to our small % expansion (neglecting finite width fluctuations).
B. Expanded Related Works
In this section provide additional details comparing contributions with previous works. We break these up into (1)
scale-renormalization theories (2) Adaptive Kernel Theories in the NTK scaling (3) and rescaling the likelihood.
B.1. Scale Renormalization Theories
Many theories of proportional limits of Bayesian neural networks predict that the mean predictor has the form
(s(x)) = ((a)
[b + 3 (31(((() Σ th (2 xp) [Q(a) 44 + +3-11] -1 m yo = Σ (x, xrr)
-1 yu yv: (61)
JAM
where Q(a) E RR is a scalar that is determined self consistently as a function of a = P/N and T x²) is the final layer's NNGP kernel under the prior (Li & Sompolinsky, 2021; Pacelli et al., 2023). We rewrote this expression in the last line to emphasize that this is equivalent to a data dependent ridge with the original NNGP kernel. In the limit of low temperature 1 2o (zero regularization). these effects have no impact on the mean predictor. This theory does not capture adaptation in the entries of the kernel.
Extensions of this theory to CNNs, gated linear networks, and transfer learning result in slightly more order parameters which compute gate x gate overlaps for gated linear networks (Li & Sompolinsky, 2022), or space x space overlaps for CNNs (Aiudi et al., 2023), or task x task overlaps for transfer/continual learning (Shan et al., 2024; Ingrosso et al., 2025). While these theories are slightly more flexible, they do not allow for adaptation for the kernels to adapt within each block (gate, spatial location, task) but just reweight the blocks.
20

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
B.2. Adaptive Kernel Theories in NTK Scaling
Both Seroussi et al. (2023) and Fischer al. (2024) derive mean field actions for the kernels and conjugate kernels for Bayesian networks under the posterior. Seroussi et al. (2023) variational Gaussian approximation to the preactivation density which finds the best Gaussian density which approximates the true posterior density for h. Fischer al, explicitly describe how the joint distribution of kernels obeys large deviation principle at large N (or in the proportional limit
P. P.N + do with P/N=
1 N
In p(((())) max S({ $())
(62)
(41)
where S is action. However, in their theory the single-site distribution for hidden layer & have the form
(
Pe(h) oc exp
Σ AV
h
[ 1
1
Σ
((h))o(h))) (ee)
(63)
2N
JAr
which reveal that the non- non- -ou-Gaussianity is explicitly suppressed with respect to N. As they focus on large width networks N > 1, they expand the single-site densities at leading order
|1-
[1-
1 Σ 142
(h) sJ exp
(64)
2N
which enables computation of Gaussian averages with a perturbed density.
Some key differences between approach and this approach is that
Our theory does not allow for linearization of the preactivation densities unless the richness hyperparameter "o is explicitly made small. We therefore do not rely on linear response theory to solve our equations but rather confront the min-max problem directly.
Our theory results in deterministic kernels rather than random matrices for the kernels T The kernels arising from a proportional limit are random matrices. Reducing the noise from random initialization results in better performance as discuss in Appendix C.
B.3. Feature Learning At Infinite Width By Rescaling Likelihood
An alternative approach, in analogy to mean-field/pP scaling in the dynamics case (Mei et al., 2018; Chizat et al., 2020; Geiger et al., 2020; Yang & Hu, 2022; Bordelon & Pehlevan, 2022), one could instead reparameterize the network definition so that feature learning persists as - oo even if P is kept constant. This approach was pursued for Bayesian networks by Yang al. (2023). However, they do not reparameterize the definition the network predictor so their saddle point equations for the kernels are not in agreement with ours. To see this concretely in the case of hidden layer linear networks. For direct comparison, the posterior kernels satisfy the following saddle point equations in our theory and theirs at zero temperature
f Φ = C-1 - 2 "yy" + 0= Φ "y Φ
Ours
(65)
Theirs
where the second equation is provided in Equation 128 of Yang et al. (2023). We obtain their result from result by taking an ultra-rich limit to and rescaling the kernel as & & which gives
C-1 + 4
(66)
We thus see that our scaling allows for more intermediate behaviors between the ultra-rich To - oo and the lazy learning To + limits. Further, Yang et al. (2023) utilize different spline-based numerical method to approximate the preactivation density under the posterior, whereas we utilize importance sampling.
21

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
0.40
MNGP
N-500
N-0000
N-2000
aNBK
N0-00
Test Loss
N-1000
0.30
N-2000
a
N--50
N- 1000
0.25
N--000
0.20
0.15
200
400
600
800
1000
P
Figure 9. Finite-width effects in NTK and yP parameterizations. Blue curves are for different values, showing that NNGP predictor not consistent across network widths. We compare predictions finite width N networks (as (Pacelli al., 2023)) to the NNGP infinite width limit. Finite width effects are instead negligible in μP paramcterization, where finite N networks trained with Langevin consistently on the theory full orange and red curves.
C. winite-width effects
In principle, in pP, the kernels / and predictions fa at width exhibit (Bordelon & Pehlevan, 2024). Because these fluctuations generically add variance the predictor (see (Li & Sompolinsky, 2021; Pacelli et al., 2023)), they increase the test loss compared to the large N limit as show in Figure 9. Here, we compare the finite width light blue effects in IHL fully connected architecture (parameterized with NTK parameterization) at finite width and finite = P/N = N (1), with orange and red curves, which are the finite width effects in our mean-field parameterization. Light blue curves this plot are the predictor performance as obtained in (Pacelli et al., 2023). Light orange and light red curves are the finite width effects our aNTK and aNBK predictors respectively.
VN fluctuations around their limiting values
When N is comparable to either parameterization, the kernels &' should actually be thought of as random matrices with significant deformations to their spectra compared the N - oc limit.
D. Deep Bayesian CNNs
In this section, we describe the Bayesian posterior of a Deep convolutional neural network (CNN) for infinitely many channels. Here, we need to add an index a for each weight W order to account for the filter value at spacial displacement from the filter center at each layer, while St is the spatial receptive field at layer é. The L hidden layers of the CNN can be expressed as
h
WD 2..ot.j
(67a)
pia
VD
j=1
DESD
h jia
We (h
(67b)
VN
best
Σ
whe ch (iaa)
(67c)
70 N
From these definitions, the partition function turns out to be
1-1 II
Z
a 25b
II du &ToN2 E, (v.- TOr EN En wht () Eig Enn (w
(68)
E= agb
22

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
By imposing the pre-activation definitions with the of the integral representation of some Dirac delta functions as did in Eq. (20), integrate out the weights contribution and just move in the space of representations, and get
Z
dh E+
dh t+1
p De Exio Rusa &ctthett+t En Sn8p
IIIII
dw 156
II
dwine ToN Eu(ym~sp)
2TT
2T
136
t=
-tEn 8p ( Ei Do w (hi
: E.E.,
11 +4 En (w Enso Madd IN E Esest w Wi mni
X 6
(69)
1-1
d
d
d d
2 =
je
JAu
EiEl Expraot of mnnee ue +N Enuo Th $1.
$ pur
2T
2x
(=1 oz'
JMM
ds ds
troo N En 928 20 Exp - = - Ene 31 3 t
2
dh dhi
pia ue Eiin Enpo 1i pu Eten Eno Ens Moo Nyy praaal - -Eeml Emniat Benad (Es och
2x
8=1
Enu ti.
Eo pih (o
)M
dip
d pv.aa
d&L d uu
II
II
ty
N NEEt Enveed o ao B an +N Eno t 4%
2x
2π
Zmm1 1 paz
En
1
2l
4"+N
X
100
(70)
With saddle point equations
Te pu.aa!
H (0 Vp= = 1,.... T.
(71a)
L pu.aa'
2X pa Bear)
= 1, L
(71b)
$ L1 Av
Σ (h (o()))))))
(71c)
$L
7o
2
I
Day/
+
+
(71d)
2
X
B
pa
99
E. DMFT review
In this section, we briefly recall the dynamical mean field theory derivation gradient flow dynamics for multi-layer fully connected neural network. As clarified in the main text, we are interested in fully connected feedforward network with L layers, each of width N, defined as
1 fp A
L+1
w
((h)))
1+1
W (h
h
W (0)
z
(72)
VVN
VD
where each trainable parameter W is initialized as a Gaussian random variable W ~ N(O, 1) with unit variance. Here, the gradient updates for the weights (!) and the output function fy are given by
dw'(1) dl
7 (1))()()))))))))())))))))))))))) (I di
df
Σ a=1
K
NTT pa (t, (t,t)
AKff
(73)
N
23

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
captures the backpropagated gradients flowing from the downstream layers. Instead, the term ((h (t) involves the activations of the E-th layer and reflects the forward pass contribution to the weight update.
where (t)
8f, represents the pattern error signal for a pattern μa, and (f((t) vN ah (1) (hg(1)
8c
(t)
As specified in the main text, for the representer theorem to be valid in this case, we restrict to K degree homogeneous network, whose output scales as f( = a" f(0). Here, the predictor dynamics is governed by the driving force of the error signal propagated through the adaprive Neural Tangent Kernel K (t, tt, = 0h, Oh) tnr)) This quantifies the 80
00
interaction between parameter gradients for outputs of pattern pairs (1, at times (t, t'). The homogeneity factor k. appears
in the second term from the weight decay contribution. Some quantities of interest to define are the forward and gradient kernels at each layer
1 Φ mu (1, (') = (h( ((EE(1)) (1, No gE(()
1 N
(74)
which allows to rewrite K (t, t') Ei=o Ght (t, ( t"). If we take care of the initial conditions over the weights, in the DMFT limit where N, + 00 with 70 N,, we can determine the final kernels by solving the field dynamics
(t,(t) = = (t) + 7o ! dt'e -M-1))
Σ (t))) 2 (t,
(75)
<{(()) = e My) + 70 S." dt' e (A(
In Eq. (75), both pre-activations n and pre- mre-gradient signals 2/1 3 W gf+1 decouple over the neuron index and VN
factorize over the layer index, and the contribution from initial conditions (t) J VN track of the entire history trajectory, and computing at each step produces of kernel matrices that have dimension PT x PT. This scales cubically in both sample and time dimensions ON (PaT*), allowing in principle to solve for the field dynamics when P,T' = ON A sketch of an iterative algorithm procedure can be found in Algorithm 1, where given an initial guess on 1 one can compute KTNNTK and solve for the predictor dynamics of Eq. (72) once drawn a number S of 1f=1
and (i) 1 (0) is exponentially suppressed at large time t. Simulating a stochastic process like Eq. (75) requires keeping
(1)ss ~ (0, G and solved Eq. (75) for each (h (!), =f.nnn())) (4) sketch of the solver can also be found in the main text Algorithm 2. As can be noticed in Figures 10, solving for the field dynamics gives good agreement with simulations. Precisely, once knowing the predictor, one can study how both train and test losses updates during the training dynamics and for different values of y. Fig. 10(b) (top left panel) shows that the lazy learning regime when y = 0 does not allow the network to interpolate on the training data and leads to higher test loss compared to the rich cases with 7 > 0. Increasing has also the effect of speed up learning, while for that given sample size value (P- = 200 here) there exists an optimal degree of feature learning (i.e. 7) concerning the test loss. The specifics of the learning task can be found in the figure caption. Fig. 10 also shows the non-Gaussianity of the pre-activation and pre- rreggradeent distributions once the system has thermalized, at the end of training. As mentioned above, because of the weight initialization as W,Y ~ (N, 1), both (h, =} are Gaussian distributed when the training starts. Then feature learning has the effects of accumulating non non-Gaussian contributions as the training proceeds.
samples (1) m= ~ (O, Φ6-1),
E. DMFT for Convolutional Networks
In Fig. 11 we show how our DMFT theory can be extended two-layer CNN which is trained on CIFARIO images on a subset of P = 100 (left) P 1000 (right) points. When data sample is small, for any value of % there is an optimal early stopping time for the best Test I Loss performance.
The theory to get the predictor is an easy extension to the multi-layer fully-connected setting and can be found in (Bordelon & Pehlevan, 2022). However, we recall here for the sake of clarity how the field dynamic equations get modified in this setting. Again, in our notation a is the spatial displacement from the center of the filter at each layer where E S the spatial relative field at layer f as in Eq. (D). The are-activation definitions still remain the same as is in Appendix D. In the same way, the gradient signal are now defined as
0f
ahf+1
Sy = No
Σ b
p.b
(76)
2h 0h1+1 n.b
ah!
24

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
2110
130
Figure 10. The training dynamics of two layer ReLU MLPs trained with weight decay. The richly trained networks achieve lower training lest errors equal levels regularization. Our theory can reproduce preactivation and pregradient densities each setting.
P= 100
P 1000
Train Loss
Test
4:5
6:0
17
10
167
Steps
Steps
Figure 11. Training dynamics for layer CNN for varying richness to CIFAR-10 images. The images turned into patches hefore computing cross-spatial correlations in the data. The training dynamics for infinite width networks (black) is compared to training finite width networks. For small training set sizes, richer training result in faster overfitting, while this effect is less severe when there is more data.
25

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
while the weight dynamics per filter is
d (1)
70
Σ
et Aungnn (hppatt) (W)
e+1
(77)
dt
Given that, the stochastic dynamics for the pre-activation and pre- pre-gradient signals (similar to Eq. (75)) becomes
h h (1) =e= - A (1) ( dl'e- Σ
( ww.at Ei (t))
(t))
7o
v,b,e
(78)
z/aa e -a (fa (t) + +0 { dt'e dt'c-X(t-t)
Σ (
Myb.c
where again f/12 (t) VN (O) (t)) and (t)a(t) vN (0) as it is for the fully connected dynamics, the contribution form initial condition get suppressed and the fixed point predictor a kernel predictor, being the feature and gradient kernels now
(t))gttt(()) (t) from the initial conditions. At large time t,
E pa (t, t'
N (h pa (f)) (ot Gra.eht, (1)
1 N oia (t) gin (t))
(79)
pa t) and the predictor again at convergence The Neural Tangent Kernel is instead KONTK (t,t) Ees &f (t, Gftl
f(x)
I xÀL Σ, KuNKN (xx, xx).
F. Fixed point structure of GD with Weight Decay
In what follows, we will be interested in the infinite time limit of a dynamics such that in Eq. (75). Prior work on the dynamics of 1.2 regularization in the kernel regime revealed that training wide network for infinite time leads to collapse of the features and network predictor to zero (Lewkowycz & Gur-Ari, 2020). However, if one instead adopts a jP scaling, then it is possible to have a non-trivial fixed point at infinite width as the feature learning updates and regularization updates are of the same order (Bordelon & Pehlevan, 2022). This is because from Eq. we realize that in the setting where À > 0, not only the initial contribution of the fields dynamics are suppressed at large time t, but also the second terms contribute the most when the system has equilibrated, leading to a predictor which is kernel predictor
[K + [xx]] 'y (xx) = k(x.)) T
(80)
Because of the simple interpretation of DNNs in this regime, we wish to say something about the fixed point structure of the field dynamics, which are
70 Σ (m 2 = 20 Σ (E X
(81)
X
In what follows, we specialize to simple solvable cases to gain intuition for these fixed point constraints. In general, these constraints imply a set of possible joint densities over h z as well as determine the final feature and gradient kernels and the predictor.
F.1. Two Layer MLP with whitened data
Let's consider a single data point with a white covariance matrix K = 1, label 1 and a transfer function d(x) Re). In this setting, the dynamics of training for the pre-activation and pre-gradient signals are
d di M(1 M(0)= = 7))()(())) -Nh(t) _=(00= = =00((()0(((()) -xt)))
(82)
At the fixed point, the following conditions are satisfied
h 70 Ad(h)); == 70
X Ag(1).
(83)
26

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
In principle, there are infinitely many solutions these equations, and combining them gives the following constraint on
d A (h)) x
2 A (h).
(84)
X
Since we know that here ((4) max(0, ft), this means that the following two constraints on the pre-activation density must be satisfied
vh < O phh) - 0, A
(85)
Lastly, we have the equation that fixes the value of the pattern error signal through the predictor definition, which
A 1
(2σ(h))
70 (1²)
(86)
70
implying that the second moment of h must give
(h²) 70 X.
(87)
The solution here is correct if to X. Otherwise (12) = 0 and p(h) = 5(h). We can verify that this is true by comparing the pre-activation density of a two-layer MLP trained until interpolation with the theoretical predictions of the fixed points. As shown in Fig. 12, there a sharp phase transition in the limiting density right when effect of feature learning here to kill the left-side of the distribution and adjusting the variance of h > 0 such that the constraint of Eq. (87) satisfied. Another way of saying this is that in the infinite time limit t + so, the {70, Au - limits do not commute. In the first case of Fig. 12. when lim -0,0000 we get stable non-Gaussian behavior for the p(h). In the second case of Fig. 12 we collapse when limx,y-oo and nothing is learned by the network. In the same way, in the limit where fix t (1) and study the ridge-less limit of lazy network (%o 0), we recover the Neural Tangent Kernel predictor, and consequently the Gaussian pre-activation density at initialization.
A. Precisely, when To < X, the
10a
10"
10-1
10"1
Densty
10:2
10a
vp
10-3
10-3
10
-0.5
0.0
2.0
Figure 12, Pre-activation densitics of two-layer MLP trained with GD and weight decay at different times. Ligther colors represent the end of training. Dashed blue line is the thcoretical prediction from the fixed point in the infinite time limit.
F.2. Linear case
In principle, the constraints one gets by looking fixed points of Eq. (81) fix the first two moments of the pre-activation density distribution, but not enough determine the full marginal p(h). Indeed, this remains history dependent and one in principle has to track the entire update dynamics in order to get the full description. One possible way of understanding this is by looking at the simplest, linear case. Here, we initialize the weights of two-layer MLP to be Laplace distributed. Here, we expect the distribution p(h) to Gaussian distributed we train with weight decay until interpolation. Fig. 13 shows that training with weight decay to the fixed point does not recover a Gaussian single site density. But does have the properties demanded by the saddle point equations, which are again A A/2o; (h²) = 7o
27

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
10
Denaity
10'1
10"1
10'4
Figure 13. Pre-activation density of a two-layer MLP trained with P = 1 with a white covariance matrix. Dark colors correspond to carly time training. being the weights initialized Laplace distributed at t Light colors coincide with the end training, when the system has thermalized, Blue dashed line is the theory prediction from the fixed point equations.
NTK Lazy NN, lim
NTK
Lazy NN, y--0.01
0.4
Lazy Yo-0.03
NN, Y=00.05
Lazy NN, y>=0.09
0.3
0.3
LO55
Loss
0.2
0.1
0.0
10
10
10a
10
10P
15
15
10>
10*
Steps
Steps
Figure 14. Weight decay in the lazy training regime cause model "unlearn" and reduce its output after the features start to decay. Provided tu is sufficiently large compared to A, however, the final predictor will still be nontrivial, unlike the zero predictor obtained in NTK parameterization (Lewkowyez& & Gur-Ari, 2020).
G. Algorithmic implementation and numerical details
In this section, we provide more details regarding the numerical methods used for solving both the adaptive kernel theories and for training the finite width N neural networks.
G.1. Data pre-processing
As clarified in the main text, we focus on regression problems on two classes (respectively labeled as 1/11) for both MNIST and CIFARIO datasets. For MLPs architectures, the input patterns have 28 28 pixels except when comparing the MLP test prediction with CNN we did in Figure Indeed, for all the experiments except the one in Figure 5, we resize CIFAR10 images to 28 x 28 pixels by reducing their spatial resolution and then converting them to grayscale. In each case, we normalize the pixel values to have consistent scale, and then we flatten the images as a D = 784 input vectors. For CNNs on CIFARIO 10, we preserve the spatial structure of the 32 x 32 pixel images rather than flattening them. Specifically, after normalization, each image is partitioned into non-overlapping patches of size 8 x 8, resulting in a multi-dimennionnll array that maintains the local neighborhood relationships. This representation enables the convolutional layer to apply a K 8 spatial filter directly, thereby effectively capturing localized features relevant the regression tasks.
G.2. Theory solver: Min-Max optimization for Bayesian DNNs
In this section, we provide more detail on the solver for the aNBK algorithm. We use automatic differentiatiinn (via JAX) to compute gradients of the action S(( @b, $ (Equation 6) with respect to the order parameters ( &6) (Bradbury et al.,
28

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
2018). A key challenge for this is estimate the single site moment generating functions 2e as a differentiable function of both variables &e and क To do so, importance sampling, by recognizing that, conditional on o the 2g be expressed as an average of a nonlinear function with respect to Gaussian with covariance dc-1
exp (-
1 ( ((ih))T''00))
(8))))
Z₁
1~N(0,000-)
B
1 (ha()) $ ((s))s))
exp
(88)
B k=1
In our code, in practice, each of the vectors hg are i.i.d. draws from (o, (²).
). by setting de -11E {
At each step of the iteration scheme for the min-max solver, resample new batch of B vectors h3 and use these to estimate Z₁, which provides fresh samples at each iteration of the algorithm.
Our solver proceeds by alternately optimizing over the conjugate kernels (Φ') (via gradient ascent in the inner loop) and the kernels (b1 (via gradient descent in the outer loop). In practice, the inner loop runs for linner ~ 200 steps gradient steps, after which the outer loop updates { JE=1 given { JE= for maximum number louter 20000 of iterations. The code implement this scheme as follows:
computes the initial lazy guesses for kernels ( 61)2_ given E RPPP. At layer & = 1, this involves generating a samples = 20000 number of Gaussian vectors h TR. Px (by drawing from (O, o0) after computing the square-root of Φo) and then computing an empirical kernel via 4 samples of (( (h²) E PxP. The same iterative procedure applied at each layer, given &' somppss
Kernel Initialization: A function
a data covariance matrix " = 1 xx T
(oh(-1)> For the conjugated kernels, p' xxxP Vl
{L}.
samples
Single-Site Estimation: A function single_site computes the log- normalzzed moment generating function by tak- ing an average over the batch of B samples as described above: log 2, logg B Σ exp
1 (hs))
T (bo(h))))
Action Function: The overall action S({ ( $61 (Ez) is computed through a function act i on as in Equation (6).
Gradient Updates: The gradients of S(( $1)/ with respect &6 and कe are computed using JAX's automatic differentiation and are used to update the corresponding variables via gradient ascent/descent (GD) detailed above. Two separate learning rates (or update steps), up-step-Phi ~ le - 5 and up-step-hatPhize le - 4, are used to control the outer and inner updates, respectively.
G.3. Theory solver: DMFT dynamics
In this section, we detail the implementation of our DMFT dynamics solver, which is used to simulate the gradient flow dynamics with weight decay (given a regularization parameter À at each layer {L}). Minor variations to take into account spatial positions are implemented for the CNN case. This solver is based on the DMFT derivation described in Section E and captures the evolution of the pre-activation fields h, and gradient signals zμ via Monte Carlo sampling. For simplicity, we restrict to the case L = 1 (i.e., two layers architecture).
The implementation proceeds as follows:
Fields initialization: an initial set of samples= 20000 Gaussian pre-activation fields (0)) E R Pxsampless is generated by drawing Monte Carlo samples from (0, K*), where K = //=0 1 ax: E R PxP is the Gram matrix of the P data. In the same way, an auxiliary array of length samples= 20000, =00 gaamples is generated by
sampling Gaussian random variables (0) A0 1). From that, the gradient field g0)) E R Pxsamples is computed as g = ((h(0)) @ z(0) (Equation (13)).
Initial kernels and error signal: Having h())) E R sxsmples the feature kernel is evaluated as $<=1(0) = being <(-) = max(0,) -) the ReLU activation function. The initial gradient kernels are instead: (0= sample (0) E PxP, 2 E PxP Vt (T).
sumples (b(()))((())) E gg
29

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
According to Section E, the initial error signal for P patterns A
E IRP is computed A
y 70 somples (((()))z(0) because we focus on MSE loss L =
Field dynamies: at each step for T 20000 steps, the fields and the error signals are updated according to the dynamics
h( + 1) h(1) + tho (g(()) @ A (1)) 1 (1) Xh(1)
(89a)
z(t + 1) = xtt + m7o (o(h(t)) O (t)) - Az
(89b)
A (1) y 7osamples ((((1))2())
(89c)
Loss: Training loss [(t) are computed as the mean squared error over the corresponding At)), since C(t) = 1/2 4,t)3
Return: The function solve_dynase rrg returns: list of losses over iterations L(t), the target vector y. final states (h(T), xT)) for the fields and the final kernels { ( (T/=((T), (T)))(==(TT)) From that, the aNTK kernel is computed
as K-NNTK(T) Tr( Tr($#/000 (TT + (T).
In our theory solver, learning rate 1j = 1 x 10 samples 20000, and we span over P, 7o values.
G.4. MLP Experiments
For designing the model, we construct a deep multilayer perceptron (MLP) in JAX. The architecture built as a sequence of weight matrices with imput dimension D 784 (except for Figure 5 where the input dimension is D 1024) and hidden layer width N 1024 for number L + of hidden layers, each randomly initialized from a normal distribution such that W3 (0, 1). Each hidden layer is followed by ReLU non-linearity after being normalized according to the mean-field scaling as in Equation (2). The last layer returns a scalar output that serves the regression prediction.
For training, we either employ a variant of deep Langevin dynamics, or a gradient descent with weight decay dynamics according to Equation (3). The MLP is trained to minimize a regularized mean squared error (MSE) loss (rescaled by the factor N) in both cases. For Langevin training, the otimizer we use opt axtax. noinnniiinn iioossy-e sydddd- - designed to inject Gaussian white noise into the gradient updates at each iteration. This noise is drawn from a zero-m rero-mean Gaussian distribution whose variance is controlled by both the learning rate 7 and the inverse temperature parameter 5 in Equation (3). This white noise plays critical role in approximating the posterior distribution (Equation (9)) over the network parameters.
For both Langevin and gradient descent dynamics, we use a weight decay contribution proportional to À as it is for (3) (in all our experiments we use x = 1 for each layer, except when we compare with CNN test loss, where À = 1 x 10-3). For Langevin, we average the T = 20000 steps fluctuations every 1000 steps after t > 5000 epochs. We use a learning rate 7 = x 10-1 and an inverse temperature 3 = 50. For gradient descent, we train until convergence for T = 20000 epochs and we use a learning rate 17 = 1 x 10 Both the experiments are performed by varying the sample size P' and the feature learning strength To.
G.5. CNN Experiments
We implement our CNN using Flax's Linen module in JAX. The model architecture and training strategy are designed to capture localized features from the animate/inanimate regression tasks on CIFARIO images. The CNN consists of single convolutional layer with kernel size of 8 8 and stride equal to the kernel size. This choice effectively splits each 32 32 input image (with 3 color channels) into non-overlapping patches. We set the number of filters to N = 1024, and the convolution weights are initialized using a normal distribution with unit variance. To ensure that the variance is appropriately scaled, the convolutional output is divided by v3 x 8 A ReLU activation is applied afterwards. This activated output is then flattened and fed to a dense layer that outputs a single scalar prediction; the dense layer's 's weights are further scaled by a factor of N) to maintain the mean field scaling as in Equation (2).
30

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
The CNN is trained using gradient- -based optimization strategy over T 20000 epochs, over MSE loss. Before each parameter update, the gradients scaled by factor of 1 N and regularized with weight decay term proportional À as it is in the dynamics of Equation (3). The key hyperparameters for training the CNN are chosen as follows: learning rate 1 x 10-3, regularization À = 1 x 10 and the experiments are performed by varying the number of training sample P and the feature learning strength Yo-
G.6. Computational overhead
Time for Kernels Algorithm (1): the most computationally expensive part is the Monte Carlo estimation of the non- Gaussian single-site density. At each inner step, we sample a batch of B Gaussian vectors (as clarified in G.2) and estimate the single site desity function, an operation costing roughly O(BP²) per inner update (since it involves P P matrix operations). Repeating this for finer inner steps gives an overhead of approximately (timner P2) per outer update.
In contrast, when training finite-width network via full Langevin dynamics, the end-to-end cost scales O(TN2 P), where N = 1024 is the network width. For comparable T (we use T Louter 20000 for both Langevin and theory solver), this is typically much lower than the overhead of the full theory solver, (fauterr B 22), especially when the sample size P is are large.
For deep linear networks- where the single-site density is Gaussian- the overhead of Algorithm (1) is significantly reduced because the conjugate kernels can be solved explicitly (see Eq. (15)), leading to a cost of ((P)) per step.
When comparing to DMFT-based analysis (Algorithm (2)), the full dynamics require handling PT x PT kernels, leading to time complexity of O for the complete dynamics. In our work, Algorithm (1) already represents an improvement over Algorithm (2). In Appendix E.I, we propose an O(TP$) version of Algorithm (2) by leveraging weight decay to directly solve the fixed-point equations.
To get an order of magnitude, in most of our simulations fouter = 104, tineer = 2 x 10?, and T = 2 x 104, B = 2 x 10 and sample size at most P 10
Time for output predictions Generating the final predictions using Algorithm (1) requires approximately O(tomartinne operations, the procedure involves repeated sampling and kernel evaluations over the training dynamics. By comparison, Algorithm (2) incurs a heavier cost of O(T*p3), while simple kernel regression with known kernel only requires O p8)) time. End-to-end training of a finite-width network scales as O(TN2P) P) (for both Langevin and gradient descent with weight decay dynamics), which is significantly more efficient in practice when P is large.
Memory for kernels For Algorithm (1), the memory footprint is primarily dictated by the storage of the kernels ( & Jé=1- V')=, each requiring roughly O(P²) space. In contrast, Algorithm (2) have memory cost O(P2T?). Meanwhile, finite-width network training typically demands memory proportional to (N2).
H. Glossary
Here we provide more explanation of our choice of terminology for the various kernel predictors. We use the letter the end of a name indicate that this predictor is kernel method with no variance from random weight prior or initialization. We use the prefix letter "a" to indicate an adaptive kernel method where the feature kernels adapt the structure of the learning task.
NNGP: the Gaussian process with (1()) variance for the network outputs under both the prior and the posterior. The mean of this process is a kernel method with the matrix &' in the lazy limit.
NNGPK: the mean kernel predictor for lazy training with the initial final kernel neglecting. This corresponds to N + oo first followed by no + 0 in our parameterization. In this limit, there is variance of the predictor under either prior posterior.
NTK: a kernel method for the initial neural tangent kernel at infinite width without any randomness or variability in the predictor from initialization. This is the predictor obtained in the NTK parameterization if the initial output of the model is subtracted off (ie a centering operation is performed where f(0, æ)) - f(0, x) - f(00. x)).
31

---

Adaptive kernel predictors from feature-learning infinite limits of neural networks
aNBK:
the adapted Bayesian kernel method in scaling limit. This corresponds regression with the adapted 4 kernel.
aNTK:
the adapted NTK kernel in our feature learning scaling limit. This corresponds to regression with the adapted K kernel.
32


Reviews:
[{"review_id": "jHlduZo7T2", "reviewer": "Reviewer_upAJ", "summary": "This paper proposes two ways of taking the infinite-width limit of fully connected and convolutional neural networks trained by noisy gradient flow to arrive at kernel solutions in what the authors call the \"rich regime\", in which feature learning takes place, which contrasts with the \"lazy regime\", in which the weights stay in a neighbourhood of the values at initialisation. The first is the adaptive Neural Network Gaussian Process Kernel (aNNGPK), and the second is the adaptive Neural Tangent Kernel (aNTK). Derivations are given for these two solutions, as well as algorithms to compute them. No statistical results are given.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank their reviewer for their detailed feedback and suggestions, which we try addressing below. \n\n### How is the infinite limit analyzed? \n\nHere 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.\n\n### Derivations Dependent on Prior Work / What is new?\n\n 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. \nRegarding 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.\n\n\n### Add more precise statement of results (Theorem/Proof style)\n\nWe 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).\n\n### Appendix A confusing... How do we get generalization?\nThanks 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.\n\n\n### Execution Issues with Result Statement beyond Lazy Limit\nWe 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. \n\n\n\n### Questions\n\n1. *$\\delta$ is used without being introduced. What is this? Should it be there?*\n\nWe 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'$. \n\nThe continuous time version of this is one where the covariance of the process $\\{ \\epsilon(t) \\}$ is a Dirac delta function. \n\n\n2. *What is up with Eqn 18?*\n\nWe thank the reviewer for pointing out the typo in Eq.(18). We solved the inconsistency and fixed the notation for the $\\mathbf W$s.\n\n\n3. *How do we get eqn 5?*\n\nWe apologize that this was unclear. We added a more detailed derivation of this in Appendix A.2.\n\nIn the case of $\\sigma(s)=s$ we have\n\n$$f(x) = \\frac{\\beta}{\\lambda} \\sum_{\\mu} \\Delta_\\mu \\Phi(x,x_\\mu).$$\n\nWe can solve for $\\Delta_\\mu \\equiv -\\frac{\\partial L}{\\partial s_\\mu}$ \n\n$$\\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$$\nas desired. \n\n4. *Why does $\\gamma$ show up in the dynamics?*\n\nThe 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).\n\n5. *Why do Dynamics Converge to the Posterior?*\n\nWhen $\\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.", "claims_and_evidence": "In general, I think studying the NTK in the feature-learning regime is a really interesting problem, and I thank the authors for investigating this problem. However, based on my knowledge and evaluation, I think this current submission is not ready for acceptance. \n\nThe authors claim in the introduction that the behaviour of neural networks in the infinite-width limit is \"analysed\", from which I expected some kind of statistical generalisation error, but what the authors give is simply a derivation of the kernel regression solution that they converge to. Moreover, these derivations are heavily dependent on previous results in both cases, and the references to those results are very weak. For the aNNGPK, the derivation of (9) from (3) is not given in the paper, and it is unclear which result from the previous works the authors used to arrive at this derivation. For the aNTK, the authors cite [Bordelon and Pehlevan, 2022] heavily, to the extent that they use notations and quantities from that paper without saying what they are. For example, $\\xi^l_\\mu(t),\\psi^l_\\mu(t)$ are said to be \"stochastic processes inherited from the initial conditions\", but nowhere are these initial conditions given, nor are the expressions of $\\xi^l_\\mu(t),\\psi^l_\\mu(t)$ given explicitly. \n\nIn general, the mathematical presentation is weak, with many inconsistencies, abuses of notations being used without being introduced, and results being used without derivation or precise citation. \n\nMoreover, the \"derivation sketches\" are given without a precise statement of the results that they are supposed to prove. Why did the authors decide not to state the results as Theorems or Propositions, before giving the proofs?", "methods_and_evaluation_criteria": "I spent less time on the experiments section. I couldn't immediately spot any problems.", "theoretical_claims": "It is hard for me to verify the correctness of the theoretical claims, given the mathematical presentation.", "experimental_designs_or_analyses": "I spent less time on the experiments section. I couldn't immediately spot any problems.", "relation_to_broader_scientific_literature": "I think the lazy regime is pretty saturated by now, and also deemed less interesting by the community. I think contributions in the feature-learning regime are extremely valuable, but also rigorous results in this regime are very difficult to obtain. Even though this submission does not aim at rigorous theoretical results, the goals of this submission are still welcome, but I unfortunately cannot give a very high evaluation for its execution.", "essential_references_not_discussed": "There are of course many papers missing, but the literature in which this submission resides is vast and I wouldn't expect any paper in this field to be able to cite every paper in this field, so I do not count this against the submission.", "other_strengths_and_weaknesses": "In general, I think the presentation and structure of the paper could be improved a lot. For example, in Sections 3.1 and 3.2, there are \"derivation sketches\" but it is not immediately clear what these derivation sketches are supposed to prove. Also, see \"Other Comments Or Suggestions\" section for what is probably a non-exhaustive list of issues.", "other_comments_or_suggestions": "15R: \"analyzing expressivity\" -> \"analyzing the expressivity\"\n\n49R: For consistency with (1), $a^\\mu$ -> $a_\\mu$. \n\n74L: \"Further, complexity\" -> \"Further, the complexity\" \n\n72R: \"intermadiate\" -> \"intermediate\"\n\n126: \"at converge\" -> \"at convergence\"\n\n162R: (59) and (2) seem to be the same, but I think it would be much better here to refer to (2) than (59). I suspect this is due to a glitch in the labelling of the equations?\n\n189L: \"already know lazy\" -> \"already known lazy\"\n\n208L: Footnote 2 doesn't seem to exist?\n\n138R: In (2), from what's written on 150R and 170R, I think there is a square-root missing in $N_L$ in the denominator right?\n\n149R: If all the widths $N_l$, $l=1,...,L$ are assumed to be the same at $N$, then I suggest you just use $N$ everywhere. It is somewhat confusing, for example, to see both $N_L$ and $N$ being used on 170R and 171R. \n\n174R: Is the superscript $L$ missing from the $\\Phi$'s?\n\n192R: I think it would be much better to denote the dependence of $S$ on $\\boldsymbol{\\Phi}$'s explicitly, as done in (8) for example. \n\n209R: In (8), this is bad notation. The same symbols that are minned and maxed out on the right-hand side are used again on the left-hand side. \n\n224L: \"ie\" -> \"i.e.\"\n\n232L: \"means to solving\" -> \"means solving\" or \"means to solve\"\n\n222R: I think it would be more accurate to say that $K_{\\mu\\alpha}^{\\text{aNTK}}$ is the adaptive Neural Tangent Kernel (Gram) matrix, rather than saying it is the kernel itself. Same goes for the gradient kernel $G^\\ell_{\\mu\\nu}$. Also, for $\\Phi^\\ell(t,t)$, of course I can guess what this means but again this is mathematically bad notation, since, before, you used $\\boldsymbol{x}$ as arguments in $\\Phi$. Abuse of notations should be clearly stated if you're going to do it. \n\n613: A new sentence should start here?", "questions_for_authors": "175L: Here, $\\delta$ is used without being introduced. What is this? Should it be there?\n\n621: In (19), why is it that we have $\\lambda_0,...,\\lambda_L$, when in (18) we have just one $\\lambda$ without subscript? Also, for consistency of notation, e.g. with 613, all the $W$'s should be boldface. \n\nIt may be that I'm mistaken, but I tried to derive (5) from (4) using $\\sigma(s)=s$ and $\\mathcal{L}(s_\\mu)=(s_\\mu-y_\\mu)^2$, but I couldn't. Are you sure that the coefficients $\\Delta_\\mu$ are correct? Can you please spell out the derivations?\n\n$\\gamma$ is defined as $\\gamma_0N$, and then for the NTK regime you say you take the $N\\to\\infty$ limit first then $\\gamma_0\\to0$ in Table 1. But this will result in $\\gamma$ always being $\\infty$ surely? Do you mean that you let $\\gamma_0$ be some function of $N$ and take $N\\to\\infty$ such that $\\gamma$ stays at a constant value or something like this? As far as I'm aware, usually in NTKs, you can simply take $\\frac{d\\boldsymbol{\\theta}}{dt}=\\nabla_{\\boldsymbol{\\theta}}\\mathcal{L}$, possibly with a constant in front of the gradient. There shouldn't be a factor of $N$ that will take this to $\\infty$. \n\nI'm sorry if this is due to my ignorance, but is it obvious that (9) can be derived from (3) by taking fixed $\\beta$ and taking $t\\to\\infty$? I looked at the two references you gave for this, but one of them is a book without reference to a specific result, and the other is a paper in which I still couldn't immediately find a result that implies (9).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bZGap57OV5", "reviewer": "Reviewer_MHbU", "summary": "The paper shows that a GP describes neural networks trained in the rich infinite-width regime with a data-dependent kernel. Two different settings for the training dynamics are considered: one with a noise term and the other with no noise, where in both settings, there is a weight decay term. These correspond, respectively, to adaptive versions of NNGP and NT kernels. Explicit expressions for the kernel predictors and prescriptions to calculate them numerically are provided, using saddle point equations and dynamical mean field theory. It is shown that the adaptive kernels achieve lower test loss on benchmark datasets compared with kernels derived from the lazy regime.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Relationship to Fischer et al\n\nWe 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. \n\nOur 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. \n\n1. 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**).\n2. 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. \n\nIn 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. \n\nIn 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.\n\n3. 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. \n\nIn 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.\" \n\nThat 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).\n\nWe will cover this in greater detail in App B.\n\n### What is new in GD setting?\n\nThe 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.\n\n### More details about Expts\n\nAdditional experimental details are provided in the new version of Appendix F.\n\n### Questions\n1. 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. \n2. 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.\n3. In general, for gradient flow, one needs to know the full history of the NTK over training.\n4. $\\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. \n5. 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. \n\nInstead, GF + weight decay for non-homogeneous activation functions no longer satisfies a representer theorem for the $\\text{aNTK}$ kernel as in Equation (10). \n\n6. For CIFAR we resized the images to $28\\times 28$ pixels and then we convert them to grayscale.\n\n7. The simplest case is $P=1$ and $L=1$ with $\\phi(h) = \\tanh (h)$, \n$$p(h)  \\propto \\exp \\Big(-\\frac{1}{2}C^{-1} h^2 - \\frac{1}{2}\\hat{\\Phi}\\tanh(h)^2 \\Big)$$\nSince $\\hat{\\Phi} < 0$, $p(h)$ can develop a Mexican-hat shape.\n\n8. We suspect that Langevin noise acts as an entropic regularizer that reduces overfitting in small data regimes.", "claims_and_evidence": "As far as I can see, all claims are clearly stated and well-supported by theoretical or empirical evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense for the problem or application at hand. For the numerics, providing more details could be useful.", "theoretical_claims": "I checked the theoretical claims and found no major issues with their correctness.", "experimental_designs_or_analyses": "I found the experiments to be generally sound and valid. A few comments on the figures:\n\n1. Figure 3a - I don't understand the difference between NNGP and NNGPK here. \n\n2. Figure 2c - are the cols going from $\\ell=1$ to $\\ell=L$ ?\n\n3. Figure 3c, 2nd line - what is GF(wd)? is this the ground truth? \n\n4. Figure 5 - I did not find a mention of what data was used (CIFAR10 ?)", "relation_to_broader_scientific_literature": "Here are my main concerns with this paper: \n\nThere are two main parts to the contributions of this paper: the NNGP setting and the NTK setting. As far as I can see, there is limited novelty in both settings and insufficient acknowledgement of earlier related papers.\n\nIn the NNGP setting the theory leads to almost the same equations as in the missing reference Fischer et al. 2024 (see below), where the differences are only due to slightly different conventions and the choice of focusing on the regime where $P\\sim O_{N}\\left(1\\right)$ [an explicit resetting of the noise $\\kappa\\to\\kappa/\\gamma_{0}^{N}$ to account for the $\\mu$P scaling and ignoring terms of order $P$ such as the determinant term]. Can you point to any new result in your paper that does not already appear in this previous paper? Furthermore, I think that it would be fairer to give more emphasis in acknowledging other earlier works that are very much related e.g. the cited papers Seroussi et al. 2023 and Rubin et al. 2024a. \n\nAs for the NTK part, I am also not sure what is novel here relative to the earlier DMFT paper, other than adding the weight decay term and taking $t\\to\\infty$. Can you please clarify? \n\n• Fischer et al. 2024 - Critical feature learning in deep neural networks", "essential_references_not_discussed": "Here are some missing refs: \n\n1. The paper by Fischer et al. 2024 is missing, see “Relation To Broader Scientific Literature”. \n\n2. Lines 79-81: “interpreting the dynamics of gradient flow with added white noise as sampling the weights configuration from a Bayesian posterior”, the following papers are worth citing: Welling and Teh 2011, Naveh et al. 2021, Mingard 2021. \n\n3. Under: Related works - Neural networks as kernel machines, it's worth mentioning and differentiating your work from Avidan et al. Your work is well separated from that paper but might seem superficially similar for some readers. \n\n4. Lines 131-132: “Some works pursue perturbative approximations to the posterior in powers of 1/width”, the following papers are worth citing: Antognini 2019, Naveh et al. 2021. \n\n• Fischer et al. 2024 - Critical feature learning in deep neural networks\n\n• Avidan et al. - Connecting ntk and nngp: A unified theoretical framework for neural network learning dynamics in the kernel regime\n\n• Welling and Teh 2011 - Bayesian learning via stochastic gradient Langevin dynamics\n\n• Naveh et al. 2021 - Predicting the outputs of finite deep neural networks trained with noisy gradients\n\n• Mingard 2021 - Is SGD a Bayesian sampler? Well, almost\n\n• Antognini 2019 - Finite size corrections for neural network Gaussian processes", "other_strengths_and_weaknesses": "Weaknesses: \n\n1. The main weakness I see is the lack of novelty that I elaborate on under “Relation To Broader Scientific Literature”. \n*Edit*: I have raised my score to 4, provided that the authors made a substantial revision to acknowledge these prior works and differentiate their contributions. \n\n2. More details on the experiments could help. \n\n\nStrengths: \n\n1. The paper is well structured, clearly written, and relatively easy to follow, given the heavy math involved. \n\n2. It gives some nice intuitions.", "other_comments_or_suggestions": "Typos and other small issues: \n\n1. Figure 1a panel caption (right under the figure) reads $\\lambda>0$ and I believe should be $\\lambda=0$. \n\n2. Figure 1 main caption - \n\n(a) reads: “When $\\lambda=0$ the NN predictor of gradient flow *with* weight decay is not a kernel predictor” - should read “*without* weight decay”. \n\n(b) “... is not a kernel predictor” - why not? \n\n3. Line 181: know -> known. \n\n4. Line 221: the inline eq for $K^{\\mathrm{aNTK}}$ seems to be missing $\\mathrm{lim}_{t\\to\\infty}$ and $\\mu,\\nu$. \n\n5. Line 277: the inline eq for $s_{\\mu}$ seems to be missing $\\mu,\\nu$. \n\n6. Eq 17 lines 2,3 - $c\\to c_{L}$ ?", "questions_for_authors": "1. This paper studies feature learning and adaptive kernels in the $N\\to\\infty$ limit arising due to $\\gamma_{0}>0$ but with finite $P$. The paper mentions other studies that consider the same problem but in the lazy regime with $\\gamma_{0}=0$ and large $P$ (e.g. Li and Sompolinsky 2021, Fischer et al. 2024, Seroussi et al. 2023, Rubin et al. 2024 etc), but it would be interesting to comment on how the present theory behaves in this regime and if there is a simple relation between the scaling of $\\gamma$ and $P$ that would yield similar results. For example, in the deep linear case, perhaps one could directly compare to the aforementioned papers. I'm curious to get your reaction to this. \n\n2. What is new in the NTK part other than taking $t\\to\\infty$ and adding weight decay? Do you still need to numerically compute the time integrals? \n\n3. Why is it obvious that gradient flow with $\\lambda=0$ is not a kernel predictor? \n\n4. AFAIK, the lazy scaling is still the one used by practitioners - would you advocate for the adoption of the rich scaling in real world settings? or is it still mainly a theoretical tool? \n\n5. Eq. 11 - for the NTK theory $\\kappa$ plays the role of $\\beta^{-1}$ in the NNGP theory: \n\n(a) in theory, what happens for non-homogenous DNNs, e.g. with biases and/or non-homogenous activation functions (e.g. sigmoid, tanh)? \n\n(b) in the experiments, did you try activation functions other than ReLU? \n\n6. What is the dimension $D$ of the inputs in your experiments? did you do some dim reduction for the CIFAR10 data? \n\n7. The non-Gaussianity in your experiments always seems to be unimodal, can you say if there are other settings where we would expect a multi-modal distribution? \n\n8. From figure 3a it seems that aNNGPK is slightly better than aNTK - do you think that is the typical case? can you speculate why?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ucUfEL1ni2", "reviewer": "Reviewer_eXoJ", "summary": "This paper expands kernel theory to study networks in the feature learning limit. Coupled systems of equations are found for two different \"adaptive\" kernels, the aNNGP and aNTK using physics-based methods and based on different scaling assumptions on weight variances and learning dynamics. Helpful interpretations for the equations are given as well as complete solutions in the linear network case. The kernels are derived at the fixed-point of training dynamics. Thus the theory doesn't describe the evolution of training, just the endpoint. The theoretical results are well-supported with numerical experiments, and a major result of the paper are the numerical methods, since in general these kernels are hard to solve in the nonlinear case.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Analysis on the Cost of Solving Mean Field Eqns \n\nWe 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.\n\nRegarding 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$.\n\n 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.\n\nAn 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\n\n$$\\mathcal{Z}_\\ell  = \\left<  \\exp\\left( - \\frac{1}{2} \\phi(\\mathbf{h})^\\top \\hat{\\mathbf{\\Phi}}^\\ell \\phi(\\mathbf{h}) \\right) \\right>$$\n\n$$\\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) $$\n\nThis 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)$.\n\nFor 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.\n\nDespite 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. \n\nTo 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$. \n\n### Derivations could be clearer\n\nWe 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.\n\n### Experimental Details\n\nWe apologize for missing out the details concerning our numerical methods. We added a new section in Appendix F to help clarifying the settings. \n\n### More Numerical Details\n\nWe 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. \n\n#### Time complexity compared to known kernel \n\nThe 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.", "claims_and_evidence": "The main claims that are made are theoretical and concern the match between theory and experiment. I will treat these in the theory section below.", "methods_and_evaluation_criteria": "The main \"methods\" in this paper, beyond theory, concern the numerical solution to the theory equations. My understanding of the past results in this area is that solving the coupled mean-field theory equations is very costly. In fact, I have heard it said that solving the theory equations is more costly than just performing end-to-end training of a deep network. For this reason, I wish the authors had offered more:\n1. Explanation of the numerical methods, in sufficient detail to understand the implementation level.\n2. An analysis of the cost of these numerical methods or at least some description of the cost relative to network training.\nOverall, I found the numerical results to offer good support for the theory.", "theoretical_claims": "I'll admit I'm not an expert in the specific methods used by the authors, but I did read the derivations. I didn't notice any real errors but potentially caught some typos that I'll enumerate below. The methods presented are non-rigorous, so no proofs need checking. I did find some of the derivations difficult to follow, in particular the steps where the saddle point equations are derived which were left out. I'd encourage the authors to include some of those for readers who might want a hand filling in the details.", "experimental_designs_or_analyses": "The experiments are standard image classification tasks and seem standard. I don't find there to be enough detail to fully evaluate the design of the numerical methods.", "relation_to_broader_scientific_literature": "I believe that the authors are well-aware of related work in this area and appear to have addressed most of the important related work.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "I find the paper clearly written, for the most part, with good intuition provided for the forms of the equations. I like that the theory is able to predict overfitting and the generalization error for realistic networks (CNNs not just MLPs) on real data. The only real weaknesses I see are with respect to clarity.", "other_comments_or_suggestions": "* Line 72 right col: \"intermadiate\" typo\n* Figure 1: The y-axis \"Test Prediction $\\mu$P-AK\" labels in this plot don't seem consistent with the rest of the paper. If the results in (b) and (c) are comparing NN output with aNTK and aNNGPK, I don't know what you're using for comparison in (a).\n* In section 3.1, the authors move from having a single weight decay parameter $\\lambda$ to layer-dependent ones $\\lambda_\\ell$. These need to be defined.\n* Line 234 right col: the function $s(a\\theta) = a^\\kappa(\\theta)$, is this supposed to be the pre-readout of the network, in which case you'll have dependence of $\\kappa$ on the number of layers, or the homogeneity of a single layer nonlinearity $\\phi$? This jumped out at me.\n* Line 308, \"Implement the non-Markovian dynamics\" this is a big mystery to me. The appendix also doesn't really describe this, at least it doesn't describe a practical way to solve equation 62.\n* Figure 3: \"Langevin\" isn't described, is this the same mystery method compared in Fig 1(a)?\n* Supp, line 695, eqns 24, 25: The factor of 1/2 in front of the first sum and the $\\gamma_0^2$ in front of the second sum in (26) are not consistent with the form of the exponent in (25), if $NS$ is supposed to be the exponent. This appears to be a typo.\n* Supp, line 1323, distribution $\\mathcal{N}(0,\\Phi^{\\ell-1})$ seems like it should be $\\mathcal{N}(0,\\Phi^{\\ell-1}/\\lambda_{\\ell-1})$ to be consistent with the main text.\n* Supp, line 1333, I found this sentence confusing in terms of which functions are being optimized for which variables. To find the dual variables $\\hat{\\Phi}^\\ell$ you have to optimize $Z_\\ell$; to find the non-dual ones it seems you have to optimize $Z_{\\ell+1}$. This should be spelled out in a bit more detail, in terms of what order you do the optimization and for which variables.", "questions_for_authors": "Can you please expand the detail of the numerical methods explanation and/or release the code that solves the theory? This would be helpful for reproducibility.\n\nDoes the numerical approach offer significant advantages over previous methods, and how does it compare to training a network end-end? What kind of scaling does it have relative to, say, a kernel method with known kernel? This is relevant to the introduction's argument that \"directly training infinitely-wid feature learning networks\" could be a promising future approach.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BMUkdzS2JD", "reviewer": "Reviewer_Za7T", "summary": "The manuscript develops an approach for predicting deep non-linear neural networks outputs in mean-field and muP scaling. Central to their approach is the notion that kernels within each layer (Gram matrices) which, strictly speaking, are stochastic objects, concentrate leading to an effective description of the trained neural networks as a series of layer-dependent kernels (Adaptive Kernels) which obey deterministic non-linear matrix equations. They provide explicit expressions for these equations for linear settings (where they also solve the equations) and weakly-mean-field/perturbative settings (for the Bayesian case). In general, this leads to an effective GP like predictions though with those \"Adaptive kernels\" rather than the original data agnostic ones. The authors establish this for Bayesian inference and an NTK-like training away from the lazy regime (i.e. gradient flow). They also provide impressive numerical demonstrations which compare the solutions of their equations to the numerics and discuss statistics of activations.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Summary Clarification\n\nRegarding 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). \n\nTo 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. \n\n### Relations to Other Works\n\nWe 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.\n\nRegarding 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). \n\n\nThe 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-ﬁeld 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.*\n\nAgain, our theoretical results are substantially different:\n1. 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)$.  \n2. 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. \n\nRegarding 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**.\n\nAs an example, consider linear network with $L=1$ hidden layer. Our equation is\n\\begin{align}\n    \\Phi^{-1} = C^{-1} - \\gamma_0^2 \\Phi^{-1} yy^\\top \\Phi^{-1}\n\\end{align}\nThis differs from the equation 128 of Yang et al which states\n\\begin{align}\n    C^{-1} = \\Phi^{-1} yy^\\top \\Phi^{-1}\n\\end{align}\nwhich shows that our theoretical predictions are **different in general**, especially for small $\\gamma_0$.\n\nWe 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\n\\begin{align}\n    C^{-1} = \\bar{\\Phi}^{-1} yy^\\top \\bar{\\Phi}^{-1} + \\mathcal{O}(\\gamma_0^{-1})\n\\end{align}\n\nHowever our equation holds for a variety of richness levels, interpolating between lazy and rich regimes. \n\n\n### What is new for GF/GD?\n\nPlease see response to reviewer MHbU on this topic.\n\n\n### Positioning of Paper in Bayesian NN Literature\n\nWe 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.\n\n### Numerics Details\n\nThanks 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.", "claims_and_evidence": "The manuscript is generally well written and, while short of being a proof, shows solid arguments combined with numerical experiments.", "methods_and_evaluation_criteria": "The methods are appropriate.", "theoretical_claims": "I have checked their claims and agreed with their findings.", "experimental_designs_or_analyses": "I did not find any issues with the actual experiments.", "relation_to_broader_scientific_literature": "The main contribution of this work, in my mind, is numerical, especially in analyzing much larger datasets than previous works and deeper networks. \n\nThe authors, however, place too much focus on Bayesian results which already appeared, in minor variations, in various previous settings, often without noticing or informing the reader. Specifically, their equations for the kernels in the Bayesian appeared in a previous publication (https://proceedings.mlr.press/v235/fischer24a.html) as well as the perturbative analysis and the solution for a linear network. Furthermore, in that work it was established that these equations, upon taking the variational Gaussian approximation within each layer, also coincide with those of (https://www.nature.com/articles/s41467-023-36361-y). In addition, focusing on linear networks, the same equations were obtained by https://proceedings.mlr.press/v202/yang23k/yang23k.pdf. Earlier works by the latter authors also discussed the interpretation of trained networks in the rich regime via kernel flexibility/adaptation.\n\nSimilarly, their results regarding the NTK appear to largely overlap with Ref. https://arxiv.org/abs/2205.09653 and the *qualitative* addition made in this paper compared to that previous work are not clearly fleshed out. This seems to involve the introduction of weight decay which facilitates their numerical solutions, however, if so, this is appears as a somewhat technical increment (especially given that solvers are not posed in the center of the work).", "essential_references_not_discussed": "See above.", "other_strengths_and_weaknesses": "Stregths:\n\n1. The generalization of kernel adaptation to the ``equilibrium\" of gradient flow is a potentially interesting addition to the existing literature, however, only if this can be well separated from Ref. https://arxiv.org/abs/2205.09653.  \n\n2. The numerics are quite impressive compared to the existing ones on kernel adaptation. \n\n3. The paper is also well written and, apart from some minor points, the appendices are quite detailed and accessible for those with previous knowledge in statistical mechanics. \n\nIn terms of weakness, \n\n1. The positioning of this paper within the existing literature in the Bayesian setting should be improved. Similarly, since their NTK result involves running the full time dynamics, the novelty of their findings when compared to https://arxiv.org/abs/2205.09653 is not fleshed out. \n\n2. The need for numerics in writing down the kernel equation (i.e. App. F.) as well as the need for numerics to solve these equations is an obvious shortcoming, though admittedly, it is hard to see how one can avoid the latter for realistic datasets. More assurance on the convergence/needed-sample-size for the sampling procedure in App. F.  can be helpful. For instance, how does the number of samples scale with the number of data-points and input-dimension? How does this behave on the strength of feature learning?", "other_comments_or_suggestions": "I'd be willing to raise my score if the above presentation issues are fixed.", "questions_for_authors": "See above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "None."}}]

Author rebuttal:

Meta-review: This paper shows that infinite-width neural networks in the feature-learning regimes (e.g., via the mu-parametrization) also admit a kernel-regression/Gaussian-process view. This is useful since it enables further analysis of infinite-width neural networks. The authors show this theoretically, and provide numerical recipes to compute the derived data-dependent kernels.

While the reviewers pointed out that there are some crucial missing references, they agree that this work is novel and should be accepted. Feedback about the presentation of the paper also came up prominently during the review/rebuttal period. However, the majority of the reviewers think that this can be fixed easily.

In any case, I implore the authors to incorporate all reviewers' suggestions, especially re. references and presentation.
Decision: Accept (poster)