From Kernels to Features: A Multi-Scale Adaptive Theory of Feature Learning
Authors: Noa Rubin
Kirsten Fischer
Javed Lindner
Inbar Seroussi
Zohar Ringel
Michael Krämer
Moritz Helias
Venue: ICML 2025 poster (2025)
Keywords: feature learning
deep learning theory
statistical field theory
lazy & rich learning
Primary area: theory->deep_learning
TL;DR: This work presents a multi-scale adaptive theory to understand feature learning in neural networks, bridging kernel rescaling and kernel adaptation across different scaling regimes and providing insights into directional feature learning effects.

Abstract: Feature learning in neural networks is crucial for their expressive power and inductive biases, motivating various theoretical approaches. Some approaches describe network behavior after training through a change in kernel scale from initialization, resulting in a generalization power comparable to a Gaussian process. Conversely, in other approaches training results in the adaptation of the kernel to the data, involving directional changes to the kernel. The relationship and respective strengths of these two views have so far remained unresolved. This work presents a theoretical framework of multi-scale adaptive feature learning bridging these two views. Using methods from statistical mechanics, we derive analytical expressions for network output statistics which are valid across scaling regimes and in the continuum between them. A systematic expansion of the network's probability distribution reveals that mean-field scaling requires only a saddle-point approximation, while standard scaling necessitates additional correction terms. Remarkably, we find across regimes that kernel adaptation can be reduced to an effective kernel rescaling when predicting the mean network output in the special case of a linear network. However, for linear and non-linear networks, the multi-scale adaptive approach captures directional feature learning effects, providing richer insights than what could be recovered from a rescaling of the kernel alone.

From Kernels to Features: A Multi-Scale Adaptive Theory of Feature Learning
Noa Rubin Kirsten Fischer 23
Javed Lindner "245 Inbar Seroussi Zohar Ringel Michael Krämer S
Moritz Helias 24
Abstract
1. Introduction
Feature learning in neural networks crucial for their expressive power and inductive biases, moti- vating various theoretical approaches. Some ap- proaches describe network behavior after train- ing through change in kernel scale from initial- ization, resulting in a generalization power com- parable to a Gaussian process. Conversely, in other approaches training results in the adapta- tion of the kernel to the data, involving directional changes to the kernel. The relationship and re- spective strengths of these two views have so far remained unde This work presents a theo- retical framework of multi-scale adaptive feature learning bridging these two views. Using methods from statistical mechanics, we derive analytical expressions for network output statisties which are valid across scaling regimes and in the contin- uum between them. A systematic expansion of the network's probability distribution reveals that mean-field scaling requires only a saddle-point approximation, while standard scaling necessi- tates additional correction terms. Remarkably, we find across regimes that kernel adaptation can be reduced to an effective kernel rescaling when predicting the mean network output in the spe- cial case of a linear network. However, for linear and non-linear networks, the multi-scale adaptive approach captures directional feature learning ef- fects, providing richer insights than what could be recovered from a rescaling of the kernel alone.
By mistake, this document lacks David Dahmen, Institute for Advanced Simulation (IAS-6), Computational and Sys- tems Neuroscience, Jülich Research Centre, Jülich, Ger- many, an additional author; he has been an author on this project from its beginning. To properly reference this work, please see our arXiv preprint (Rubin et al., 2025).
A central phenomenon that is essential for explaining the power of neural networks (NNs) is feature learning (FL), where networks learn meaningful high-dimensional repre- sentations of the data (Bengio et al., 2013). FL. plays an increasingly important role our ability to understand and rationalize the behavior of large language models (LLMs). Sparse autoencoders can extract so called monosemantic features from LLMs that are given by superposition of layer activations (Bricken et al., 2023); these features allow interpreting and even altering model behavior (Templeton et al., 2024). Beyond interpretability, FL is essential for efficient generalization with finite data, as it enhances infor- mative directions in the learned representations, reducing the sample complexity of learning functions of these direc- tions (Abbe et al., 2021; Paccolat et al., 2021; Dandi et al., 2024). Despite its significance, many open questions re- main regarding the mechanisms underlying the emergence of these feature directions.
A well-characterized case in NN theory is the limit of infinite-wiee and finite sample size, where networks behave as Gaussian processes (GPs) (MacKay, 2003), characterized by the neural network Gaussian process (NNGP) kernel (Neal, 1996; Williams, 1998; Matthews et al., 2018; Lee et al., 2018). However, the NNGP does not capture FL, which emerges at finite network width, the proportional limit, where both network width and sample size tend in- finity proportionally (Li & Sompolinsky, 2021), or in certain scaling regimes (Yang al., 2024). Hence the NNGP fails to capture the networks' nuanced internal representations that arise from feature learning (van Meegen & Sompolin- sky, 2024). Multiple theoretical approaches have emerged to describe this phenomenon, yet there is no consensus on how to characterize FL. A common approach is to study the change of the network kernel, though the existing frame- works differ in their predictions for this change.
Equal contribution 'The Racah Institute of Physics. The He- brew University of Jerusalem, Jerusalem, Israel Institute for Advanced Simulation (IAS-6), Computational and Systems Neuro- science, Jülich Research Centre, Jülich, Germany 'RWTH Aachen University, Aachen, Germany 'Department of Physics, RWTH Aachen University, Aachen, Germany 'Institute for Theoreti- cal Particle Physies and Cosmology, RWTH Aachen University, Aachen, Germany "Department of Applied Mathematics, School of Mathematical Sciences, Tel-Aviv University, Tel-Aviv, Isracl. Correspondence to: Noa Rubin <no..rubin @ mail. huji.a..u........
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Multi-Scale Adaptive Theory of Feature Learning
One prominent class of theories, which are commonly re- ferred rescaling theories (Li & Sompolinsky, 2021; Pacelli et al., 2023; Baglioni et al., 2024), predicts that the average network ouiput and variance can be described by a rescaled NNGP kernel. Initially developed for linear networks in the standard scaling regime', this framework surprisingly yields impressively accurate predictions even in mean-fichl scaling". Despite the strong FL in this regime, the average network outputs can be obtained from output kernel that is rescaled NNGP kernel.
(a)
(1) rt
(ii) rt
Trained Network
losp
mintniszntiom timm
minimizatio
hisbolin.
*
However, FL is often considered a structural phenomenon, as in the case of Gabor filters (Gabor, 1946; Rai & Rivas, 2020) that emerge in the latent layers of convolutional neural networks (Luan et al., 2018). Thus, the expectation is that the effect of FL on the output should be directional as well. The rescaling result raises fundamental questions about how learned features are represented in network outputs and can be captured theoretically.
Rescaling
Adaptive
Multi-Scale
(ii))
Adaptive Theory
directional feature learning
In contrast, adaptive theories of FL (Roberts et al., 2022: Seroussi et al., 2023; Bordelon & Pehlevan, 2023; Fischer et al., 2024b; van Meegen & Sompolinsky, 2024) consider learned features, predicting that the kernel undergoes a struc- tural change and incorporates features explicitly. Conse- quentially, these theories are able to predict phenomena in networks that stem from FL such as a reduction in sam- ple complexity - the required amount of samples to learn a given task- relative to that of GP (Naveh & Ringel, 2021b) as well as grokking (Rubin et al., 2024). However, adaptive theories are significantly more complex computa- tionally than rescaling theories, while yielding comparable predictions for quantities such as the network loss. A funda- mental open question remains: How can two such different deseriptions of FL. be valid at the same time?
[b)
0.7
NNGP
Experiment
0.6
Saddle-Point
Train
Squared Errors
Fluct. Corr.
Test
0.5
Rescaling
0.3
Fig.2
0.2
Fig. 3
0.1-
0.0
10
10-1
10"
10
Fluctuations Scale 1/x
Figure 1. (a) The multi-scale adaptive theory bridges between rescaling and adaptive theories of feature learning. Starting from the distribution of network outputs for trained networks, the choice of order parameter decides whether a rescaling (red) adaptive (blue) theory is obtained. The choice of order parameter recasts feature learning into either (i) now-dimensional minimization or (ii) high-dimensional minimization problem. An approxima- tion of the multi-scale adaptive theory in certain limits yields the result of the rescaling approach, but in addition describes (iii) directional aspects of feature learning. (b) Training (solid line) and test errors (dashed line) across scaling regimes for different approaches. While standard scaling (green shaded area) requires a one-loop approximation with fluctuation corrections (Fluct. Comr.), suddle-point or tree-level approximation (Saddle-Poiie Pointt is suffi- cient in mean-field scaling (orange shaded arca). We show results for the kernel rescaling theory by (Li & Sompolinsky, 2021) as reference (Rescaling). We here show results for linearly separa- ble task; for results on MNIST see Fig. in App. D. Parameters: Puss 10
In this work, we address this pivotal question by systemati- cally connecting different FL theories as well discussing their differences. Using methods from statistical physics, we recast the theoretical description of the posterior distribution of network outputs into a minimization problem with respect to quantity which we call the "order parameter". We find that different theories result from different choices of order parameters, in particular with regard to their dimensionality (see Fig. la). Our main contributions are:
We derive a multi-scale adaptive theory that is valid across the full range of scaling regimes, from mean- field to standard scaling, which allows us to system- atically include finite-width corrections (see Fig. 1b). This generalizes previous adaptive approaches, which were restricted to specific scaling regimes, and holds for arbitrary tasks as well as non-linear networks.
? = 1, Punin 80, N 100, D 200, Ko
ge = I = 0.5, Ap = 0.1.
'A scaling where readout weight variance scales 1 1/width. A scaling where readout weight variance scales as 1/ /wiitth".
2

---

Multi-Scale Adaptive Theory of Feature Learning
By analyzing the simplest non-trivial model, we rec- oncile adaptive and rescaling theories: we show that for the mean network output the multi-scale adaptive theory can be approximated in certain limits to yield an effective rescaling of the kernel. This explains why certain FL phenomena do not appear in rescaling theo- ries.
Bayesian framework: Zavatone- Veth & Pehlevan (2021b) study properties of the network prior, whereas focus the network posterior. Hanin & Zlokapa (2023) obtain rigorous non-asymptotic description of deep linear net- works in terms of Meijer-G functions. Zavatone- Veth et al. (2022) study the same setting but consider explicit models on the input data in the limit of infinite pattern dimension. Zavatone- Veth & Pehlevan (2021a) investigate deep linear networks in different proportional limits, recovering the results from Li & Sompolinsky in an adaptive approach. Cui et al. (2023) study non-linear networks, exploiting the Nishimori conditions that hold for Bayes-optimal inference, where student and teacher have the same architecture and the student uses the teacher's 's weight distribution as a prior; the latter is assumed Gaussian i.i.d., which allows them to use the Gaussian equivalence principle (Goldt et al., 2020) to obtain closed-form solutions.
Our theory reveals how the two approaches differ. In linear networks, rescaling and adaptive theories yield equivalent predictions for the mean network output, but not for the output covariance. In mean-fiell field scaling, the covariance exhibits clear adaptation to task-relevant directions, accurately captured by our adaptive theory and not by a rescaling theory. non-linear networks, the disparity between the approaches emerges already at the level of the mean predictions: our adaptive the- ory correctly predicts change in sample complexity class relative to the NNGP, a phenomenon that is not captured by rescaling theories.
Our work is distinct from perturbative approaches such as (Antognini, 2019; Naveh et al., 2021; Cohen et al., Roberts et al., 2022; Hanin & Zlokapa, 2023; Hanin, 2024) for the Bayesian setting or (Dyer & Gur-Ari, 2020; Huang & Yau, 2020; Aitken & Gur- Ari, 2020; Roberts al., 2022; Bordelon & Pehlevan, 2023; Buzaglo et al., 2024) for gradient-basee training that use the strength of non-Gaussian cumulants of the outputs as an ex- pansion parameter. In contrast, we perform an expansion in terms of fluctuations around the mean outputs, which is able to capture phenomena that escape perturbative treatments such as phase transitions: this technique corresponds to an infinite resummation of perturbative terms. Our approach is similar to van Meegen & Sompolinsky (2024) with the difference that they scale weight variances as I/N2, so that readout weights concentrate (see App. B for a comparison of the approaches).
2021; Halverson et al., 2021;
Overall, our findings suggest that a comprehensive under- standing of FL requires moving beyond kernel rescaling towards high-dimensional kernel adaptation.
2. Related works
The limit of infinite network width and finite amount of training data has been studied extensively, yielding among others the NNGP kernel (Neal, 1995; Williams, 1998; Lee et al., 2018; Matthews et al., 2018; Avidan et al., 2024). This theory relates network behavior at initialization to train- ing dynamics (Poole et al., 2016; Pennington et al., 2017: Schoenholz et al., 2017: Xiao et al., 2018). However, the NNGP cannot explain the often superior performance of finite-width networks (Li et al., 2015; Chizat et al., 2019; Lee et al., 2020; Aitchison, 2020; Refinetti et al., 2021), requiring either the inclusion of finite-width effects or dif- ferent infinite-widit limits such as μP scaling (Yang et al., 2022; Vyas et al., 2023).
Another line of work focuses on the dynamics of FL: Saxe et al. (2014) derive exact learning dynamics for deep linear networks, while Bordelon & Pehlevan (2023) use dynamical mean-field theory to describe network behavior in the early stages of training of gradient descent training in different scaling regimes; in contrast we consider networks at equi- librium. Yang & Hu (2020) consider the effect of network training dynamics and learning rate scales in networks. An experimental investigation of kernels in feature learning in gradient descent settings was performed by Canatar & Pehlevan (2022). Day et al. (2024) study the effect of weight initialization on generalization and training speed. A differ- ent viewpoint considers spectral properties of FL (Simon et al., 2023; Yang et al., 2024) as well as investigating the ef- feets of learned representations directly (Petrini et al., 2023). Maillard et al. (2024) derive polynomial scaling limits of the required amount of training data.
Describing FL in neural networks in Bayesian framework has led to concurrent views: kernel rescaling (Li & Som- polinsky, 2021; 2022; Pacelli et al., 2023; Bassetti et al., 2024; Baglioni et al., 2024) and kernel adaptation (Aitchi- son, 2020; Naveh & Ringel, 2021a; Seroussi et al., 2023; Fischer et al., 2024b; Rubin et al., 2024; van Meegen & Som- polinsky, 2024). These differ in the choice of order param- eters considered and in consequence also in the explained phenomena. A complementary perspective is provided by Yang et al. (2023), who propose unified theoretical frame- work showing that FI. extends classical kernel methods by enabling the learning of data-dependent feature maps in the infinite-width regime.
Various works study other aspects of networks in the
3

---

Multi-Scale Adaptive Theory of Feature Learning
3. Single hidden-layer network
work prior is given by
We consider the following network architecture
P+1
W(ifo.if.) in exp if of
Σ
ha = Viaa fa w ((ha), ya fa + Ea+
(1)
wi:hos
(3)
where { is Gaussian regularization noise & iid kO) We consider P tuples of training data D = {(I.. ((LaaBaahee with Ta E _2d and you E IR as well as an unseen test point (x..y.) y.) denoted by *. Here, Φ denotes a non-linear acti- vation function and Sa E R is the scalar network output. We study the Bayesian setting with Gaussian priors on the readin weights V R NND Vú ~N(0, g./ /D) and the readout weights w E RN
where the average (..,)wy,by is over the prior distribution on the network parameters and the hidden pre-activations
overj ((0,, C(ax) with (( = 9-/ D XXT E ((++1)((++) The detailed derivation can be found in App. A. The statistics of the conjugate fields fo are directly linked to the statisties of the network predictors fo via the output discrepancies A y - fDD on the training data
hag ~
as W, N(O0.ww/NT)
(A) -ixxoo'
(4)
We differentiate between two cases: (a) standard scaling for 7 = 1 and (b) mean-field scaling for 2. Accord- ingly. we scale the regularization noise as A = Kg_V' -7 so that it does not dominate the network output in mean- field scaling. For concise notation, we use the shorthands fo (fr) X = ( and y SSaSP in the following. Further, summations over repeated indices are implied VaLzz Σ VHL The code for theory and experiments can be found in 10.5281/zenode 15480988....5 15480988.
To obtain the statistics of the conjugate variables (JD, f.) and thus also of the network outputs (fD, we define a conditional cumulant- generating function W(h; = In (exp(i-f. + ik: fD)) s..fo which yields
j.l) -In s /dfexp lik fo + S.DD.. (5) In S(fD.i.lv) -iy io N0 N' NNT-1660
4. Multi-scale adaptive feature learning theory
+ W(iD,,so).
We compute the network posterior the test point (Ir.:
Here, we introduced source terms (k, (k,j.) with k E R", j. E R to obtain the statistics of (fv, as derivatives. On the training points, we have
by conditioning on the training data D and derive set of self-consistency equations for the average discrepancies (A) between labels y and mcan posterior network outputs (fo) on the training data. This description on the level of the discrepancies yields a high-resolution picture of the network behavior: allows us to explain kernel rescaling results in the proportional limit as well as predict directional aspects of FL.
(f)) =
(fbiT)) - W. (6)
with (FofbD) the covariance of fp. On the test point, we get
(f+) = -g Wls..=o. (f.) = 03, W (7)
4.1. Predictor statisties of the neural network
However, the namulant-generating function W(kjj-l) in (5) in general does not have an analytical solution. Instead, perform systematic expansion in terms of fluctuations of the network output using the I Legendre transform of the cumulant- -eenraatgg function W of the network posterior
We are interested in the Bayesian network posterior for the network output on training points fo and test point fa, which corresponds to training the network with Langevin stochastic gradient descent (LSGD) (Welling & Teh, 2011; M et al., 2017; Naveh et al., 2021) until convergence (see App. C.3 for details). We denote the joint vector of out- puts as ʃ = (/D, f.) E P-1 Following along the lines of Segadlo et al. (2022a), we write the joint distribution as
I r(f = extra ik (v).
(8)
where take the extremum with respect to k. This trans- form is function of the mean conjugate field f = (fo) (in the following, we drop the index D for readability), defined self-consistenly by the stationary condition given by
p(fD, = -(ullD,
x /afro dio fds. dī. exp [x [-
2 oyr(f..))) 1 o.
(9)
with f = f.) the conjugate fields to f= (fp, f=). The cwwulant-genae generatnng function W if.) of the net-
The Legendre transform r(F. J. j.ly) is thus a natural way of constructing a minimization problem that yields the quantity

---

Multi-Scale Adaptive Theory of Feature Learning
Training discrepancies
(b) Test discrepancies
1.0
1.0
NNGP
NNGP
22 2208 : ...
(Anoryy
0.5
Tree-Level
0.5
Tree-Level
0.0
0.0
-0.5
-1.0
-1.0
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
(Amp)
(Amp)
(f-) on Ising task in mean-field scaling. We show theoretical values for both NNGP and tree-level against empirical results, where the gray line marks identity. contrast the NNGP, the tree-level approximation accurately matches the empirical values. While we use á id here, the non-lincar case erf yiclds similar results (sce Fig. 9 App. D). Parameters: y = 2, Panin = 80, N = 100, = 200, Ko = 1, Pau = 10 go = 9w = 0.5, Ap 0.1.
Figure 2. (a) Training discrepancies (A) = (/r) and (b) test discrepancies (A.)
we are interested in. It recasts the problem of computing the statistics of the posterior, which is the stationary solu- tion of the stochastic minimization problem described by the LSGD training, into an effective deterministic optimiza- tion problem of I' with regard to the mean discrepancies f: intuitively, we may therefore think of I' as an effective loss function that explicitly depends only on the mean dis- crepancies (A) x f. but implicitly takes fluctuations of A into account. Moreover, it allows computing corrections to the mean network outputs in systematic manner, building on broad foundation of methods from statistical physies (Zinn-Justin, 1996; Helias & Dahmen, 2020).
4.3. Saddle-point approximation in mean-field scaling
In mean-field scaling, the exponent S of the cumulant- generating function in (5) scales linearly with the network width N, while the fluctuations of the network output scale (( 1/N and become negligible. Thus, we can perform saddle-point approximation for the integral in (5) and obtain the tree-level approximation of the Legendre transform (Helias & Dahmen, 2020) by replacing by its mean value + yielding
((f))+l)) = (f. = -(fff-)))
(12)
We derive this result more rigorously in App. A.2 using large deviation principle (Touchette, 2009). From the stationary condition in (9), obtain self-consistency equation for given by
Using the relationship between first-order parametric deriva- tives of the Legendre transform I and the cumulant- generating function W (Kkkj--b)) ly), we obtain
(f) -0 [ =o
(10)
i= =i -ioi(ko--'1 (F) cL
y.
(13)
Next, we consider systematic approximations of I for differ- ent scaling regimes, from which we determine the network output statistics.
On (J) = gw Sw N-1
[
Iw
(1E) ffr N2 DD
(14)
where Cler) E PxP refers to the training data submatrix of ((2z), In the remainder of this section, f refers to the solution of (13). We obtain the discrepancies on the training points as
4.2. Exact network prior for linear networks
The above expressions are exact but require knowledge of the cumulant- -oeellantgeeeratingg function W of the network prior in (3). For general non-linear activation functions o, the cumulant-gen geneurtnngg function W can only be approximated, e.g. using a cumulant expansion (see App. A.1.2 for details).
(A)TL = No (NOl + Cr (f) cheee ) y. (15)
For linear activations @(h) h, however, we derive an exact expression
For the test point, we get
(S.)m.
(16)
N W(ffo. if.) = Ind I+
f(xz) N
Iw
(11)
-[on cb
Cr (f) C y
(ax)
Since the goal of this work is to connect existing theories while also studying their differences, we consider the linear setting in the following as the simplest setting possible. De- spite this choice, the theoretical framework presented here also applies to the non-line non-linear case as discussed in App. A.
where C.boa) = Dow/Dac I I.tisosp E B recover- ing results by Seroussi et al. (2023). In Fig. 2, we compare theoretical values for training and test discrepancies against empirical measurements for linear networks trained on a
5

---

Multi-Scale Adaptive Theory of Feature Learning
(a) Training discrepancies
(b) Test discrepancies
1.0
1.0
NNGP
NNGP
(Anory)
Tree-Level
0.5
Tree-Level
0.0
One-Loop
0.0
One-Loop
LS
LS
-0.5
-0.5
-1.0
-1.0
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
(Aimp)
(Aum))
(Amory) - (Annea))
0.1
A Amnenn
0.1
1:
0.0
0.0
Tree-Level
Tree-Level
-0.1
One-Loop
-0.1
One-Loop
-0.2
One-Loop simpl.
-0.2
One-Loop simpl.
LS
-0.3
-0.3
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
-1.00
-0.75
-0.50
-0.2
0.00
0.25
0.50
0.75
1.00
(nN)))
(ANNEe)
Figure 3. (a) Training discrepancies (A) - y (fb) and (b) test discrepancies (A.) = y. (f-) on an Ising task standard scaling. Upper row; theoretical values for different theories against empirical results, gray line marks the identity. Lower row: difference of theoretical values to the NNGP as baseline against NNGP predictions, indicating small differences between the different approaches. Results of the kernel rescaling approach by Li & Sompolinsky (2021) are shown as reference (LS), Parameters: = 1, Pimin = 80,
N = 100, 200, Ko = 0.4, Pree 10", g- = 0.5, gv = 0.2, Ap= = 0.1.
linearly separable Ising task (see App. C.3 for details). Com- paring to the NNGP as baseline, we find that, while the NNGP fails to match network outputs, the multi-scale adap- tive theory accurately predicts the values observed in trained networks. For similar results on non-linear networks, see Fig. 9 in App. D.
where A(J) NOT + (n) clez) and g(m) refers to the n-th derivative of the exponent S with respect to f evaluated f (see App. A.2 for details). In the remainder of this section, ff refers to the self-consistent solution of (19), which is not necessarily the same as the one of (13) in the previous section. This yields for the training discrepancies (A) = iso. as in (4) and for the test point from (10)
4.4. Fluctuation corrections in standard scaling
(f~)i-aup
(20)
In standard scaling, output fluctuations are not scaled down by the network width N and instead become non-negligible. To obtain the leading-order fluctuation corrections, we ex- pand the exponent S of the gumulant-generaan (5) around its saddle-point, / to second order as
= =No C( (wx)) Cn (T) + 1 (1)))
-1 (()
) Sa
function in
In the next section, we will see how these expressions re- duce to a kernel rescaling theory in the proportional limit N xx P- DO, which in linear networks we refer to as one- loop simplified in Fig. 3, where we compare theoretical pre- dictions empirical measurements on the Ising task. We show results for the multi-scale adaptive theory presented here as well as the rescaling theory by Li & Sompolinsky (2021), which was derived for the standard scaling regime. Due to the weak FL in standard scaling, all theories match the network behavior relatively well. However, by taking the NNGP as a reference, the differences between the theories become discernable: The tree-level solution shows devia- tions from the other solutions, predicting overly small test errors compared to the one-loop solution and compared to empirics. Furthermore, predictions of the one-loop solution agree to those of the rescaling theory by Li & Sompolinsky (2021).
s(fo.i-l)) = + (t (Fo- (17)
where (t2) denotes the Hessian of S(f, j-ly) with respect to f at the aaddle-point f. Caleulating the Gaussian integral in (5), we obtain the one-loop approximation of the Legendre transform (Helias & Dahmen, 2020)
log det(- (18)
-
The self-consistency equation for J from the stationary con- dition in (9) is then given by
7. = [A( 0B
1 [(()
S(3)
iyy-
(19)
2
c8R
The one-loop solution takes into account leading-order fluc-
j =0=
6

---

Multi-Scale Adaptive Theory of Feature Learning
tuation corrections. The latter vanish in mean-field scaling, so one expects the one-loop approximation to converge to the tree-level result in this scaling regime. We show this ex- plicitly in Fig. 1,, where we demonstrate how the different theories transition between the two scaling regimes by seal- ing Ky + Ko/X X and IN + X with 0.1/N < 1/x < 10 determining the scale of fluctuations. As expected, train and test errors decrease for increasing FL in the mean-field regime. Due non-negligible fluctuations, the tree-level and one-loop solutions differ in standard scaling. When further increasing the fluctuations scale, even the one-loop solution does not accurately predict empirical measurements anymore since this regime requires fluctuation corrections beyond first order. In principle, the multi-scale adaptive approach allows computing these higher-order correction terms (Helias & Dahmen, 2020). When decreasing the fluctuations towards the mean-field scaling regime, the one- loop solution converges to the tree-level solution. Notably, the here presented multi-scale adaptive approach accurately predicts train and test errors across both scaling regimes, including the intermediate regime.
equations obtained by Li & Sompolinsky (2021) can be obtained from (2) by marginalizing over the hidden pre- activations h in (3) and performing a change of variables so that the posterior is function of single scalar order parameter Q == | P. A saddle-point approximation with respect to this variable yields a self-consistency equation for Q and consequently expressions for the predictor statisties on test and training points, such as the mean and fluctua- tions. As the order parameter is scalar here, it is limited to describing scalar changes to the kernel.
Conversely, the choice of the high-dimensional order param- eter in the multi-scale adaptive approach, which in mean- field scaling reproduces equations from the approach in (Seroussi et al., 2023), results in structural changes to the kernel. Notably, the choice of a high-dimensional order parameter results in the need to correct for fluctuations that arise in standard scaling, requiring us to go beyond the saddle-poin point approximation by using fluctuation corrections.
Surprisingly, as we have shown in the previous section, for a linear network on the level of the mean predictors, the multi- scale adaptive approach converges to that of the rescaling one, even though they have qualitatively different kernels. This observation motivates (ii), showing that for linear network in the proportional limit N x P - co, regardless of the initial choice of order parameter, the mean network output can be obtained from kernel regression (Rasmussen & Williams, 2006) with a rescaled NNGP kernel. For non- linear networks, however, differences already arise on the level of the mean predictors, as shown in the next section.
5. Kernel rescaling theory as an approximation of the multi-scale adaptive theory
Existing rescaling theories (Li & Sompolinsky, 2021; 2022: Pacelli et al., 2023; Bassetti et al., 2024; Baglioni et al., 2024) and adaptive theories (Naveh & Ringel, 2021a; Seroussi et al., 2023; Fischer et al., 2024b; Rubin et al., 2024; van Meegen & Sompolinsky, 2024) make both qual- itatively and quantitatively different predictions regarding network behavior. On the one hand, rescaling approaches predict that the mean network output is equivalent to that obtained by a rescaled NNGP kernel. On the other hand, adaptive approaches such as the multi-scale adaptive theory presented here, as well other existing approaches, predict that the kernel adapts to the data in richer manner, show- ing changes in specific directions that are determined by the training data's statistics. While these approaches are quite different, we here expose the close relation between them in two respects: (i) We show that the adaptive and the rescal- ing approach can both be derived from the same starting point; the expression for the joint distribution of the network outputs (2). (ii) We show that for linear networks the adap- tive approach in the proportional limit N x P + can be approximated by kernel rescaling for the mean outputs. For output fluctuations and for certain non- non-linear networks, however, such a reduction does not hold.
In the kernel rescaling case, the predictor for the mean output is obtained by replacing the NNGP kemel KNNGP = 9w - C((p)) with rescaled kernel
Krescaling = Q/ KNNGP.
(21)
For the multi-scale adaptive approach presented here, the output statistics in mean-field scaling are obtained from
K Kadupivee TI.
I +
Iw (2x) NY DD FT..
KNNGP. (22)
The appearing matrix product allows a non-trivial change of the NNGP kernel in certain meaningful directions, e.g. a teacher direction or dominant eigendirections of the input kernel that align with the target, thereby yielding additional insights. However, we derive an equivalent equation for the mean predictor by simplifying (11) using the matrix- deterninant-eemma, which yields the mean output from a rescaled NNGP kernel given by
Technically, the differences between the two viewpoints stem from different choices of the order parameter used in the approximation of the posterior, utilizing either a saddle- point approximation or including fluctuation corrections. Specifically, with point (i), we show in App. A.4 that the
K Krewaling, TL. Qn (f) N1->) KNNGP,
(23)
where QT. (T) = 9n, N1- /(1 (T î) and j sat- isfies (13). So even though the adaptive approach in mean- field scaling considers a directional change to the kernel,
N
7

---

Multi-Scale Adaptive Theory of Feature Learning
in terms of the mean output this is equivalent to a rescaled kernel.
Directional Output Distrihution
Adaplive
In standard scaling, one cannot immediately express the mean output in terms of a rescaled kernel. However, in the proportional limit x P - oo, certain fluctuation correc- tion terms become negligible, reducing the expressions to a rescaling form again (see App. A.5). The rescaling factor is given by
4
NNGP
d-p d-
Rescaling Experiment
Directional Output Distribution
é g
(f+
(24)
2
d P
Qi Tr [A- (A- clozl
=Qm. N
where A(f) =x0o++ + Qn (n) and satisfies (141).
10-2
10-1
100
We thus find that known theoretical approaches are all de- rived from the same original posterior distribution by con- sidering different order parameters, while their resulting pre- dictions for the mean network output behave like a rescaled NNGP. However the rescaling equivalence of mean predic- tors holds only for linear networks, as well as non-linear networks approximated using cumulant expansion of the con-linearity as described in App. A.1, since one only needs to substitute c(pe) + C,00 However this equivalence does not necessarily hold for more fine-graine tineigraiedd methods of approximating the non-linearity, such as the approach described in App. A.2.3. Applying such approximations within the adaptive approach allows predicting various phe- nomena that emerge in non-linear networks such as phase transitions (Rubin al., 2024) and changes to sample com- plexity App. A.2.3, that escape description by a rescaled kernel.
Fluctuation Scale 1/x
Figure 4. Relative directional feature learning in teacher student setting as a function the fluctuation scale 1/x. Both NNGP and rescaling theory fail to capture directional feature learning, while the multi-scale adaptive theory accurately predicts network behavior. Insets show the output distribution in different direc- tions; a detailed version can be found in Fig. 12 in the Appendix. Parameters: Pirain 80, N 200, D 50, xo 2, gv 0.01,
9w = 2.
Then, (()) Tr ((ffT)) indicates that the feature direc- tion á dominates the covariance, implying that this feature has been perfectly learned, whereas & ((3) < Tr (()ST)) is an indication of weak directional FL. As derived in App. A.2, we obtain for the covariance of the network out- puts the training data
2Q N
= -T- x2 (A
F) -1
(26)
6. Directional feature learning emerges in adaptive description
where we observe a structural change in the covariance ma- trix in form of the term F = cbe (A)A)) TC(e)) which is not present in rescaling of the NNGP, whose covariance is (see (129) in App. A.4 for details)
The power of NNs stems from their ability to detect high- dimensional features in the data, implying that in the transi- tion from the lazy to the rich regime this would be reflected in the network output statistics in a non-trivial manner. It is well established that the network weights adapt during train- ing in an anisotropic manner, detecting relevant directions present in the training data (Seroussi et al., 2023; Fischer et al., 2024b); yet surprisingly, for the mean output of linear network this adaptation seems to be equivalent to isotropic rescaling of the NNGP kernel.
"fff"rescalinn <I - x² (KI Q (
(27)
As evident from the expressions for the covariance, the directional FL measure & differs significantly between the two approaches, which is illustrated most easily for a kernel ((arr) x I: the isotropy of the rescaling theory then results in the same value of & (0) independent of the direction of c. whereas the structural change of the covariance in the adaptive theory by the rank-one term (A)( T in (26) may yield larger values of Φ ( for features & || (A).
Teacher-student setting In this section, we demonstrate that the directional aspect of FL is nonetheless present in output fluctuations, which is only captured by the adaptive approach. Given normalized feature direction d. we de- fine directional FL measure & (Φ)) that indicates to which degree this feature is represented by the network as
We show the directional aspect of FL in a teacher-student setting, where the teacher is given by y = Xw. with X ~ N I) and the student is a linear network as in Section 3 with @(h) = h. In this setting, the teacher
3 = T NO.
(25)
8

---

Multi-Scale Adaptive Theory of Feature Learning
defines a feature direction U. X / [Xa/] and for comparison, we consider another possible feature direction #1 Xww/ / wil |, orthogonal to the former in the sense that W1 T w.. The latter can be thought of as the direction of a randomly selected teacher that differs in the weights of the hidden layer from that of the actual target teacher. In Fig. 4, we show the relative directional FL measure & (0.)/((tL) between the target teacher and a random, orthogonal teacher direction. While the rescaling theory does not differentiate between these directions, the adap- tive theory accurately predicts amplification of the teacher direction when entering the mean-field regime.
learn the non-linear component of the teacher, in contrast to the NNGP and kernel rescaling approaches which pre- dict the non-linearity would not be learned (see Fig. 5 in App. A.2.3). These results reflect a change to the sample complexity class, and establish principled connection be- tween feature learning and improved network performance.
7. Discussion
In this work we present unified theoretical framework to understand feature learning (FI.) in the Bayesian set- ting across scaling regimes, from lazy to rich learning, This framework describes both effects of data adaptation in trained networks, i.e. directional changes of the net- work's s output statistics in response to statistical dependen- cies present in the training data, as well as output rescaling phenomena that were described in previous works (Li & Sompolinsky, 2021; Pacelli et al., 2023). Our theory thus creates links between existing and so far unconnected previ- ous theories. In the rich regime, the presented multi-scale adaptive theory clearly exposes directional aspects of FL, thus going beyond rescaling theories. By considering the simplest non-trivial case of linear network, we finally reconcile the apparent contradiction between directional adaptation and rescaling by recovering the latter ap- proximation of the former on the level of the mean output.
Amplification of kernel structure The expression (26) likewise exposes a cooperative effect caused by feature learning as a result of shaping the kernel jointly by the input and the output statistics. To illustrate this effect, we consider the Ising task (see (C.3)), where the input ker- nel c = 9o (I e yy") contains a rank-one term com- posed of the binary classification labels, unlike the teacher- student input kernel where Chpp I. Although the factor
1+ + 4p(1 - p) is potentially small, as is the case for weakly distinguishable patterns (0 p 1-p-11 11122 & 1), the matrix factor F appearing in (26) can significantly amplify this signal. Consider the case that the target is not learned well, for example at large ridge 16, so one has (A) = y. The matrix factor
Furthermore, the multi-scale adaptive theory presented here applies to both standard and mean-field scaling and the en- tirety of the scaling spectrum. The latter is possible since the presented theoretical framework allows systematically computing fluctuation corrections depending on the scaling regime. While the tree-level solution and its equivalence to rescaling applies only to 7 = 2, the one-loop theory applies to the full regime y E [1, 2] and thus also the equiva- lence to rescaling (see Section 5). Note that this equivalence is restricted to the mean predictor, while the adaptive the- ory captures additional aspects, like directional FL in the variance terms (see Section 6). Moreover, this equivalence breaks down in the case of non-linear networks, where the presence of FL. results in changes to the sample complexity class. In addition, our theory does not make any assump- tions on the data set; we show results for an Ising task, a teacher student task and MNIST.
F F vy T (1 (1+cP)2
(28)
then shows that the weak structure X present in the input kernel has been amplified by factor P (see Fig. 11 in App. D). For an unstructured kernel ((=x) x 1 (E = 0) the rank-one component still persists F = -6 yy but is not amplified by P here. The cooperative effect, in contrast, shows that feature learning may utilize coherent statistics between input and output.
Non-linear networks Directional effects in linear net- works emerge only in the covariance of the outputs. For non-linear networks, however, these effects appear in the mean outputs, and have qualitative effects network per- formance. The adaptive approach described in App. A.2.3 enables structural changes in the kernel which enhance spe- cific features of the data. This enhancement, in turn, al- lows the network to learn non-linear functions of those features with significantly fewer data points than would typ- ically be expected. In particular, introducing even small non-linearity can lead to a fundamental difference in the behavior of approaches that rely on a rescaled or adapted kernel. InApp. A.2.3 we demonstrate this for a teacher- student setting similar to the one above but where we take both teacher and student to be non-linear. We find that the adaptive approach accurately predicts that the student will
Outlook In this work, we study shallow networks, but we expect directional FL to be crucial for network perfor- mance in deeper networks as well. Beyond this, it will be valuable to extend the theoretical framework to other net- work architectures such as convolutional networks, residual networks, and transformers, using the respective network priors Saarria--Aoonse et al., 2019; Hron et al., 2020; Fis- cher et al., 2024a). To study the effect of noise in input data on FL (Lindner et al., 2023), we would like to include fluctuations of the input kernel in the theoretical framework.

---

Multi-Scale Adaptive Theory of Feature Learning
Acknowledgements
Avidan, Y., Li, Q., and Sompolinsky, H. Connecting ntk and nngp: A unified theoretical framework for wide neu- ral network learning dynamics, 2024. URL https: / //arxiv. wwg. hbtp://b///a 2309. 04522.
This work was partly funded by the Deutsche Forschungs- gemeinschaft (DFG, German Research Foundation)
368482240/GK24416 and the Isracli Science Foundation 374/23. The authors gratefully acknowledge the computing time granted by the JARA Vergabegremium and provided on the JARA Partition part of the supercomputer JURECA at Forschungszentrumm Jülich (computation grant JINB33). MK would like to thank the Institute for Advanced Simu- lation (IAS-6) at Juelich Research Center and its directors Markus Diesmann and Sonja Grün for their hospitality dur- ing regular visits. We thank the reviewers for providing valuable feedback.
Baglioni, P., Pacelli, R., Aiudi, R., Di Renzo, F,, Vez- zani, A., Burioni, R., and Rotondo, P. Predictive power of a bayesian effective action for fully connected one hidden layer neural networks in the proportional limit. 133:027301, 2024. doi: 10.1 10.Le03/hhssRewLett. 133.027301. https:::: //ink/inn./////////////////// aps. /tt///////////// 10. 1...3//hyyeeeeet...mmmm 133. 027301. 2,
Bassetti, F., Gherardi, M., Ingrosso, A., Pastore, M., and Rotondo, Feature learning in finite-width bayesian deep linear networks with multiple outputs and convo- lutional layers, 2024. URL https: htt////xxvv.rrr///// abs 2bbs/2404 03260.032..... 2,5
Impact Statement
This paper works towards understanding feature learning, thus aiming to advance explainability of networks. While the latter surely has societal impacts, these will be much further down the line.
Bengio, Y., Courville, A., and Vincent, P. Represen- tation learning: A review and new perspectives. 35 18):1798-1828, aug 2013. doi: 10. 10.1109/tpmi..2011 50. URL https://doi..nn 10tg//0.1199tpamm 2013.50.
References
URL https: 10. 10. http://enndo..188899.8 15480898.
Bordelon, B. and Pehlevan, C. Self-consistent dynamical field theory of kernel evolution in wide neural networks
2023(1 11) 14009, nov 2023, doi: 10. 10.1088//742-5468/ ad01bo. URL https: : dxx doi www..oogg....00000 ad01bo.
Abbe, E., Boix-Adsera, E., Brennan, M. S., Bresler, G., and Nagaraj, D. The staircase property: How hierarchical structure can guide deep learning. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang P. and Vaughan, W. (eds.), Advances in Neural Information Processing Systems, vol- ume 34, pp. 26989-27002. Curran Associates, Inc., 2021. URL oc ccc/ppet_ffllleesppprr/////////e htwp://ooccceiiin....niiss..eee. neurips. 22001188673322222116676668888883
1742 5468/ 3466//d0100. 1,2
Bricken, T., Templeton, A., Batson, J., Chen, B., Jermyn, A., and team, A. Towards monosemanticity: Decom- posing language models with dictionary learning, 2023, URL https://transfon wwwpeeerr-rrrttteeee pub/ 2023/mono http://oosmmmiii---ftteeeeeeeeee
html.
Buzaglo, G., Harel, I., Nacson, M. S., Brutzkus, A., Srebro, N., and Soudry, D. How uniform random weights induce non-uniform bias: Typical interpolating neural networks generalize with narrow teachers. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Pro- ceedings of the 41st International Conference on Ma- chine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 5035-5081. PMLR, 21-27 Jul 2024. BRLhttps: httpas:://ppccppepepeeeee mlr. wtt.pisss.........s./.........ee v235 /235/bozagglo/// 2233/bbzaglo21::. html. 2
pdf. 1
Aitchison, L. Why bigger is not always better: on finite and infinite neural networks. In Proceedings of the 37th Inter- national Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 156- 164. PMLR, 2020. URL httph::tt//.o.........ttt htt.thts////ccceeeeee /wwcceddoeeeeeeeeeoooooocc ngs. mlr.press v http///////aaccisnnneeaeeeee html. 2
Aitken, K. and Gur-Ari, G. On the asymptoties of wide net-
works with polynomial activations, 2020. URL https:
//arxiv. / abs / 2006 1668.22 2
Canatar, A. and Pehlevan, C. A kernel analysis of feature learning in deep neural networks. In 2022 58th Annual Allerion Conference Commu- nication, Control, and Computing (Allerton), pp. 1-8, 2022. doi: 10.1109//Allerton44993.. 2022. 9929375. URL /tttttppp////a.l.or//1.....00 2022. 9929375.
Antognini, J. M. Finite size corrections for neural network gaussian processes. pp. 10030 10030 [cs. 2019, URL httts://arxiv. http///axxvv.v// owg./ abs 1908 10030. 2
Ariosto, S., Pacelli, R., Pastore, M., Ginelli, F., Gherardi, M., and Rotondo, Statistical mechanics of deep learning beyond the infinite-widd limit. 2023. A.4, A.4.2
Allerton4 4933777n49997.000 9937
10

---

Multi-Scale Adaptive Theory of Feature Learning
Chizat, L., Oyallon, E., and Bach, F.
On lazy
Halverson, Maiti, A., and Stoner, K. Neural networks and quantum field theory. Machine Learning: Sci- ence and Technology, 2(3):035002, apr 2021. doi: 10. (088/2632-2115//becaa) URL https: /ttpiooorggroooooo 10. -212232222222/22 abeca3. 2
training in differentiable programming. volume 32, 2019. URL https:/ http://eteeeeeeenee/eeeeeee id=rkgXDVSILB. 2
Cohen, O., Malka, O., and Ringel, Z. Learn- ing curves for overparametrized deep neural net- works: A field theory perspective. 3:023034, 2021. doi: 10.116///hysseeReeeem...00000000 URL https://link. org/ BlyaRevResearch.3.O...
Hanin, B. Random fully connected neural networks perturbatively solvable hicrarchies. Journal of Machine
Learning Research, 25(267): 2024. URL http: 2322223--3342000 html.
Hanin, and Zlokapa, A. Bayesian interpolation with deep linear networks. 120123) 2023. doi: 10. 1073/pnas.23013455122 URL https: /www: pnas. org/ abs////s/////// /10 1073 pnas. 200..2301130 2301345120. 2
Cui, H., Krzakala, F., and Zdeborova, L. Bayes-optimal learning of deep random networks of extensive-width. In Proceedings of the 40th International Conference on Ma-
chine Learning, volume 202, pp. 6468 -6211. PMLR, 23-
29 Jul 2023. URL https: /ptppppeeonnns....// mlr. press /2202/222222//////222////2 cui233. html.. 2
Helias, M. and Dahmen, D. Statistical Field Theory for Neural Newworks. Springer International Publishing. 2020. doi: 10.
4.1, 4.3, 4.4, 4.4, A,
Dandi, Y., Krzakala, F., Loureiro, B., Pesce, L., and Stephan, L. How two-layer neural networks learn, (giant) step time. Journal of Machine Learning Research, 25(349):1-65, 2024. URL http: http://jmlr. org/ peggpapers/// htt////////22///// html.
A.2
Hron, J., Bahri, Y., Sohl-Dickstein, J., and Novak, R. In- finite attention: NNGP and NTK for deep attention net- works. In III. H. D. and Singh, A. (eds.), Proceed- ings of the 37th International Conference on Machine Learning, volume 119, pp. 4376-4386. PMLR, 13-18 Jul 2020. UtLhttps: pttptts::/pproeeeeee /ttppppeccnngs..c ml press/ v119 v119/hron20a. html. 7
Day, H., Kahn, Y., and Roberts, D. A. Feature learning and generalization in deep networks with orthogonal weights, 2024. URL nttps://arxiv. org/ abs/ 3310.07765.
Huang, J. and Yau, H.TT. Dynamics of deep neu- ral networks and neural tangent hierarchy. In Pro- ceedings of the 37th International Conference on Ma- chine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 4542 4551. PMLR, 13-18 Jul 2020. URL https: htttttss:/ppocceednngge. mlr.presss v119 1119/huaae hts9/hhang201..0 html.......mmmmm
Dyer, E. and Gur-Ari, G. Asymptoties of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020. URL https: /// openreview.ner net/ formm?oormmmm/mm d= 1F 2
Fischer, K., Dahmen, D., and Helias, M. Field theory for optimal signal propagation in resnets, 2024a. URL
https: /
paops://arxiive //taxxixvvvv/////////eee/eaaaav/ org/ abs 2305. 07715.
Krogh, A. and Hertz, A simple weight decay can improve generalization. volume Morgan-Kaufmann, 1991. URL ittps::/prroceddiggs.eeus cc cc/papers [iaaper__Lless/ppee///9///// Beef feeffcoof55300ee411offfffffffff pdf. C.3
Fischer, K., Lindner, J., Dahmen, D., Ringel, Z., Krämer, M., and Helias, M. Critical feature learning in deep neural networks, 2024b. 1, 2, 5, A
Gabor, D. Theory of communication. Journal of the Insti- fution of Electrical Engineers Part I: General, 94:58- 58, 1946. URL https::t////// semt::: icscholar. org/ ID: 61327032. 1
Lee, J., Sohl-Dickstein, J., Pennington, J., Novak, R. Schoenholz, S., and Bahri, Y. Deep neural networks gaussian processes. In International Conference on Learning Representations, Vancouver, 2018. Open- Review.net. URL https:/ / openreview.ne f Forum?1d=B1EA-M-02 1,2
Garriga-Aonso, A., Rasmussen, C. E., and Aitchison, L. Deep convolutional networks as shallow gaussian pro- cesses. In International Conference on Learning Rep- resentations, 2019. URL https: http::://///eeeeeeee... net/forum Biw.unndddddkllsieideeea 7
Lee, J., Schoenholz, S., Pennington, J., Adlam, B., Xiao, I.., Novak, R., and Sohl-Dickstein, J. Finite versus infinite neural networks: an empirical study, volume 33, pp. 15156-15172. Curran Associates, Inc., 2020. URL https: http:///poceennss..//// neurips. wwtpcccepppp.....///////// ffli//iie////////
Goldt, S., Reeves, G., Mézard, M., Krzakala, F,, and Zdeborová, L. The Gaussian cquivalence of gener- ative models for learning with two-layer neural net- works. 2020. URL https:/ /ttp://aaivvv.//ii org/ abs/ 2006.147090

---

Multi-Scale Adaptive Theory of Feature Learning
ad a006f59924ffa0077355022222222
in wide deep neural networks. In International Confer- ence Learning Representations, 2018. URL https:
pdf. 2
/twpenneeiew..ee/eeeeeeeeeeeeeee wwtp::/weee/ttooeeee/o/oe/ nGgWC- 1,
Li, Q. and Sompolinsky, H.
Statistical Mechanics
2
of Deep Linear Neural Networks:
The Back-
propagating Kernel Renormalization.
11(3):
Naveh, G. and Ringel, Z. A self consistent theory of gaus- sian processes captures feature learning effects in finite CNNs. virtual, 2021a. NeurIPS 2021. URL https:// openreview net L 2, 5
031059, 2021. doi: 10.1I0//hhysReex....11053
URL https: htw.///orrosssaoeooeo aps orgg//xx abstract /10. hwt...11///hyseeee... 11. 11t.........1.1.11 031059. 1, 1, 5, 7, A.4, A.4, A.4.2, A.5,
Li, Q. and Sompolinsky, networks. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 14789-34801 Curran Associates, Inc., 2022. 2,5
Globally gated deep linear
Naveh, G. and Ringel, Z. A self consistent theory of gaussian processes captures feature learning effects in finite cnns. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Infor- mation Processing Systems, 2021b. URL https: // openreview net forumm 1d=vBYwWBVVccS..
Li, Y., Yosinski, J., Clune, J., Lipson, H., and Hopcroft, Convergent learning: Do different neural networks learn the same representations? In Storcheus, D., Ros- tamizadeh, A., and Kumar, S. (eds.), Proceedings of the 1st International Workshop on Feature Extraction: Modern Questions and Challenges at NIPS 2015, vol- ume 44 of Proceedings of Machine Learning Research, pp. 196-212, Montreal, Canada, 11 Dec 2015. PMLR. www.mlr.... parr.presssvvv
Naveh, G., Ben David, O., Sompolinsky, H., and Ringel, Z. Predicting the outputs of finite deep neural networks trained with noisy gradients. 104: 064301, Dec 2021. doi: 10.1103/PhsseevE.104 064301. URL https::// aps org prgg/do/100 10t...434........11 064301. 2, 4.1, C.3
104.
Neal, R. M. Bayesian learning for neural networks.
URL httpthtss///ppoeeeeeee lil5convergent html. 2
Springer, 1995. 2
Neal, R. M. Bayesian Learning for Neural Net- works. Springer New York, 1996. doi: 10.1007/ 978- 978-1-4612-0045-0. URL https: http://doi..rgggg 1007/978- 1777881-46122 4612-4622220 0011461111100000 0745-0.
Lindner, J., Dahmen, D., Krämer, M., and Helias, M. A theory of data variability in neural network bayesian in- ference. arXiv preprint arXiv:2307 16695, 2023. URL pettss//arxxie //t////iivrrggab///////////
owg/ htt////a/////// 16695. 7
Paccolat, J., Petrini, L.., Geiger, M., Tyloo, K., and Wyart, M. Geometric compression of invariant manifolds in neural networks. Journal of Statistical Mechanics: Theory and Experiment, 202 242((4)044001,, apr 2021. doi: 10.1088/ 1742-5468/abf1 https:/ 1088/1742 5 54 /aaffff....
Luan, S., Zhang. B., Zhou, S., Chen, C., Han, J., Yang. W., and Liu, J. Gabor convolutional networks. In 2018 IEEE Winter Conference on Applications of Computer Vision (WACV), pp. 1254-1262. 2018. doi: 10.1109/WACV. 2018.00142. 1
Pacelli, R., Ariosto, S., Pastore, M., Ginelli, F., Gherardi, M., and Rotondo, P. statistical mechanics frame- work for bayesian deep neural networks beyond the infinite-wide width limit. Nat. Mach. Intell., 5(12):1497- 1507, December 2023. ISSN 2522-5839. doi: 10.1038/ $42256-023-007767-6 URL https://oo::: http://too.oo/o/////////o /10/ 1038/ 023 0ttp://42226002-00000000 007 6. 1, 5, A.4
M, S., t, Hoffman, M. D., and Blei, D. M. Stochastic gra- dient descent as approximate bayesian inference. Jour- nal of Machine Learning Research, 18(134):1-35, 2017. URL http:/ http://mmrr.oooo wwwprg/appps/////p//ppp html. 4.1
MacKay, J. Information theory, inference and learning algorithms. Cambridge university press, 2003. 1
Pennington, J., Schoenholz, S., and Ganguli, S. Resur- recting the sigmoid deep learning through dynamical isometry: theory and practice. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Gameett, R. Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.,
Maillard, A., Troiani, E., Martin, S., Krzakala, F., and Zdeborova, L. Bayes-optimal learning of an extensive- width neural network from quadratically many samples. In The eighth Annual Conference on Neural In- formation Processing Systems, 2024. URL https:// openreview. net. forum ( dwwddd=88nnnn.. 1d=
2017. URL nttps::/proceedinss. neurips...... cc/paper_ http:ppeeffllleeppppr/////////// BBF0ccaBTT3300034444222206000000 pdf. 2
Matthews, A. G. d. G., Hron, J., Rowland, M., Turner, R. F.,, and Ghabramani, Z. Gaussian process behaviour
12

---

Multi-Scale Adaptive Theory of Feature Learning
Petrini, L., Cagnetta, F. Vanden-Eijnden, E., and Wyart, M. Learning sparse features lead overfitting in neural networks*. 2023(11): 14003, nov 2023. doi: 10. 10.1088//744-5546//ddd0 URL https:://// //dx doi. doi. org/10. 1080/088000 -4 ad/0db1
Saxe, A., Mcclelland, J., and Ganguli, S. Exact solutions the nonlinear dynamics of learning in deep linear neu- ral networks. In International Conference on Learning Represenatations, 2014. 2
Schoenholz, S. S., Gilmer, J., Ganguli, S., and Sohl- Dickstein, J. Deep information propagation. Sth International Conference on Learning Representa- tions, ICLR 2017 Conference Track Proceedings, 2017. URL https: / wwwpreeeeeeee../////// net/ forum? id-HiWIUNsgg. 2
Poole, B., Lahiri, s., Raghu, M., Sohl-Dickstein, J., and Ganguli, S. Exponential expressivity in deep neural networks through transient chaos. In Advances in Neural Information Processing Sys- tems 29. 2016. URL hitps: httpa::/prooeddngg: /proceedinggs neurips oc/ htt.ipsscc//aaerr/0/1////ee 44510031196626655000533f11222123 pdf. 2
Segadlo, K., Epping, B., van Meegen, A., Dahmen, D., Krämer, M., and Helias, M. Unified field theoretical approach to deep and recurrent neuronal networks. 2022 (10):103401, 2022a. URL https: : /dxx 10.
Rai, M. and Rivas, P. A review convolutional neural networks and gabor filters in object recognition. 2020 International Conference on Computational Science and Computational Intelligence (CSCI). pp- 1560-1567, 2020. 2020. doi: 10.1109/CSCI5 COACCCISS811000000 202.000822.0..9
00..6///cc5.... 4.1
Segadlo, K., Epping, B., van Meegen, A., Dahmen, D., Krämer, M., and Helias, M. Unified field theoretical approach to deep and recurrent neuronal networks. 2022b. accepted. A
Rasmussen, and Williams, C. Gaussian Processes for Machine Learning. Adaptive Computation and Machine Learning. MIT Press, Cambridge, MA, USA, January 2006. 5
Seroussi, I., Naveh, G., and Ringel, Z. Separation of scales and thermodynamic description of feature learning in some cnns. 14(1):908, 2023. URL ht https: /doi. org/10. 10....108///00 00138//41467--0 5, 6, A.2.3
1, 2, 4.3,
Refinetti, M., Goldt, S., Krzakala, F., and Zdeborova, L. Classifying high-dimensional gaussian mixtures: Where kernel methods fail and neural networks suc- ceed. In Meila. M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learn- ing, volume 139, pp. 8936-8947, PMLR, 18-24 Jul 2021. htt.httss://proott 1.39/refinetti222 html. 2
Simon, B., Dickens, M., Karkada, D., and Deweese, M. The eigenlearning framework: A conservation law perspective kernel ridge regression and wide neu- ral networks. Transactions Machine Learning Re- search, 2023. ISSN 2835-8855 URL https:// openreview net 2d-ouunii-- 2
Templeton, A., Conerly, T., and team, A. Scal- ing monosemanticity: Extracting interpretable features from claude sonnet, 2024. URL. http:::/trransfoomee-.do 2024/scaling-nnnnoemmmttiiiddd
Risken, H. The Fokker-Planck Equation. Springer Verlag Berlin Heidelberg, 1996. doi: 10.1007/
778-3-642-61544-.11 URL tatttps:/ /ttttt::://ooo/ htt./di.oorrggg// 1007/978- 3- 6422--266 61544 10344--33 3_4. C.3
pub/
Touchette, H. The large deviation approach to statistical
Roberts, D. A., Yaida, S., and Hanin, B. The Principles of Deep Learning Theory. Cambridge University Press, URL https: May 2022. doi: 10. /tt.//99111010100 1,2 // //doi. .00oi.o0:00 wwtpo/gg./...ooo/ 1017
mechanics. 478(1):1-69, 2009. ISSN 0370-1573. URL https://wwes http:///www.sieeeediieee. con/ science/ science/articlln sccinnee/rriiaee///////777700000
4.3, A.2, A.4.1
Rubin, N., Seroussi, I., and Ringel, Z. Grokking as first or- der phase transition in two layer networks. In The Twelfth International Conference on Learning Representations, 2024. URL https:tt///.o : wpw..eveeeeeeeeeee/e/pp/e. net forum? id=3ROGsTX31R, 1, 5,
van Meegen, A. and Sompolinsky, H. Coding schemes in neural networks learning classification tasks. pp. 2406. 16689, 2024. 1, 2, B, B,
Vyas, N., Atanasov, A., Bordelon, B., Morwani, D., Sainathan, S., and Pehlevan, C. Feature-leaming net- works are consistent across widths at realistic scales. In Thirty-seventh Conference on Neural Information Pro- cessing Systems, 2023. URL https: /ppnreviee.. net:/oorumidd-i.ne
Rubin, N., Fischer. K., Lindner, J., Dahmen, D., Seroussi, I., Ringel, Z., KrÅvEmer, M., and Helias, M. From kernels to features: multi-scale adaptive theory of feature learning. arXiv preprint arXiv:2502.032 /0, 2025.
13

---

Multi-Scale Adaptive Theory of Feature Learning
Welling, M. and Teh, Y. W. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings the 28th International Conference International Conference Machine Learning, pp. 681-688, 2011. 4.1
Zavatone- Veth, A., Tong, W. L., and Pehlevan,
C. Contrasting random and learned features in deep bayesian linear regression. 105:064118, Jun 2022. doi: 10.100PPhyykeeeE.155.61111
URL. Abstract/10.11000 htt.::100///yyaee............ /
Williams, C. K. Computation with infinite neural net- works. 10(5):1203-1216, 1998. doi: 10.1162)
2
0899766983000014122 URL https:ttppt.o :ttppp::/ddo/ ////ci.oogg/1..../ 1162/08997669800112228
1,2
Zinn- Linn-Justin, J. Quantum,iee theory and critical phenomena.
Clarendon Press, Oxford, 1996. 4.1
Xiao, L., Bahri, Y., Sohl-Dickstein, J., Schoenholz, S., and field
Pennington, J. Dynamical isometry and theory of CNNs: How train 10,000-layer vamilla convo- lutional neural networks. In Dy. J. and Krause, A. (eds.), Proceedings of the 35th International Conference on Ma- chine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 5393-5402. PMLR, 10-15 Jul 2018. URL https::: httppps://ppreeeeeeeee /ptppeeeeeeggsssseeeeeppp.....se mwr. press/ v80/x 180/xiaa00 8a html. 2
Yang. A. X., Robeyns, M., Milsom, E., Anson, B., Schoots, N., and Aitchison, L. A theory of represen- tation learning gives a deep generalisation of kernel methods. In Krause, A., Brunskill, E., Cho, K., En- gelhardt, B., Sabato, S., and Scarlett, J. (eds.), Pro- ceedings of the 40th International Conference chine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 39380-39415. PMLR, 23-29 Jul press/ 2023. httsttps: / wwlppeodddinggoee mlr. htt.ggsss.....s////...... v202/yang236 html. 2
Ma-
Yang, G. and Hu, J. Feature Learning in Infinite -ihth Neural Networks. 2020. URL httpttttp//// hellhttps///aaae /ttp://xxxxvvooooee
abs/2011 0bos/201.14522, 14522.
Yang, G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, J. Tensor
Programs V: Tuning Large Neural Networks via Zero- Shot Hyperparameter Transfer. 2022. doi: 10.48550/ arXiv.2203.0344
Yang, G., Simon, J. B., and Bernstein, J. A spectral condition for feature learning, 2024. URL https: //arxiv.org/ org./ abs/ 2210. 17813. 1,2
Zavatone- Veth, J. A. and Pehlevan, C. Depth induces scale-averaging in overparameterized linear bayesian neural networks. In 2021 55th Asilomar Confer- ence on Signals, Systems, and Computers, pp. 600- 607. IEEE, 2021a. doi: 10.109991eeec0053335220.... 9723137. URL http: //t: doo./xx.oooooooo/ooooo doi www.dddi..o.....ooooooo www.o.....00.. 1109/ IEEECONF 53345. 2021. 9723137.
Zavatone- Veth, J. A. and Pehlevan, C. Exact marginal prior distributions of finite bayesian neural networks. virtual, 2021b. NeurlPS 2021. URL https: : //pp///eee......///// net/ forum? d=MxE7 2
14

---

Multi-Scale Adaptive Theory of Feature Learning
A. General approach to train and test statisties
- (fa) with (fa) denoting the mean network output, and the are interested the training discrepancies (Ao) - yo - mean network output (f.) for test point 2, after conditioning on the training data D {(r.: 1 For clarity, in the appendix we make all index notations explicit instead of using as in the main text, and denote summations over
training data points with Greek letters. We refer to the mean network outputs predictors. The joint prior distribution for (J, f) can be computed as in (Segadlo et al., 2022b; Fischer et al., 2024b) and is given by
P+1 E + wiff.if)).
vf).f. N(0l5 f
af. df. exp (
(29)
P+1
N
w(bb, fo) = f+)= In
exp
Σ
ja
Σ
Wy (((bss)
(30)
mishai
d=1
j=1
where we use the shorthands
I_x (2T), the P + 1 index corresponds to the test point, and i is the
imaginary unit. The i.i.d. distribution of the readin weights Va implies that has Lid. over) NN(,, ((az)) with the covariance matrix of the hidden-layer representation given by
{ clame) p clax) Ja=1
((ex) { (la
(((x))
(31)
{
P Jam1
To keep notation concise, summations over repeated indices on where Clax) - 9v/ X ,nd Clax) gv/ D To the right are implied in the following.
We may obtain training discrepancies (Aa) and the test predictor (f.) from the joint cumulant-generating function W for the test point defined as
w(--l)) =I in f df. / exp(i.f.)
(32)
Taking its derivatives w.r.t. to either training labels Ya or the source term j. yields the posterior of the desired quantities
-- 1))
aW(j-
(Aa) ~-
3.
(33)
By
OW(i-y)
(f+) Oj+
(34)
6.=0
because the outer derivative of the logarithm produces the normalization by the model evidence (marginal likelihood) 1/p(y) = 1/ f df f df= fff.f"
Likewise, the variances follow as
82
(A.A,) NI - 2
oya 15-=0-
(35)
a2W(ily)
(f?) a(j-)²
(36)
1==0
By inserting (29) into (32) and performing the integration over f. we can rewrite W as
W(j-l) = In / d exp ( iy No 2
N' + w(fo..
(37)
Comparing (33), (35), and (37), we note that y acts as a linear source term for fp. from which we see that the physical meaning of the field fo is related to the discrepancy between target and network output
(aa) = No. \
(38)
AA)) 1o. + N 22-2 (..f..
(39)
15

---

Multi-Scale Adaptive Theory of Feature Learning
For computational convenience, now introduce source term k:
W(k. = =In / df exp (ikafa iyafa Ko N +
(40)
allowing us to compute moments of f by differentiating by l instead of and subsequently setting k: = 0. We define the latter part of the exponent of W as the action
SU j.lv) = -iyasa N + W(ff.
(41)
Depending on the scaling y, the network outputs fully concentrate on their mean values or require corrections due to non-negligible fluctuations. To treat both cases jointly and systematically, we introduce the so-called effective action (Helias & Dahmen, 2020)
[(f. j.l) = extr, ik WKk. j-1y),
(42)
where we explicitly keep the dependence on the souree term j+ for the test point in order to compute parametric derivatives to obtain test point statistics. This corresponds to the Legendre transform of the sumulant-generatn function W; in the case that j+kj)()) has a scaling form, a large deviation principle can be applied and the effective action corresponds to the rate function (Touchette, 2009).
The argument/ / is implicitly defined by the stationary point (sometimes referred to as the equation of state)
ar (f.
ihi = 0,
(43)
0fa
as we set the source term & to 0 by definition. Using the definition of I in (42), the extremum condition yields self-consistency equation for
ifii) OW(k.
(44)
Ok
In the following we determine approximations of the Legendre transform I(ffj+ly)) to different orders statistical fluctuations, corresponding to different scaling regimes. From the definition of the effective action follows as well that we obtain the mean output that test point from
aW(k,
(r
(f+) =
(45)
0j.
0j-
k.j. m0
A.1. Cumulant-generrating function W of the network prior
We recall that haj is i.id Gaussian with covariance given by (31). Likewise, the weights w; are i.i.d., so that we may factorize the expectation to get overall factor of N in (30)
P+1 faw ((hwa)
w(fo.f.) = N In exp
Σ a=1
(46)
w.h.
reducing the expectation over and / to scalars with regard to the former neuron index j.
1.1.1. LINEAR ACTIVATION FUNCTION
For linear activation function ((h) = h, we compute first the Gaussian integral in (46) over to and then over the hidden-layer representations h, yielding
W(f.. f.)= N In
(Fhw
exp
whha (((x)) FfT]
N
det
[
Der
(47)
Nv
16

---

Multi-Scale Adaptive Theory of Feature Learning
Taking the integrals in the opposite order yields
{ (2 ((ex) f fu²)
W(fD.f = N In
exp
N
9w [
In
(48)
2
N
The expressions identical also by the matrix-determinat lemma.
A.1.2. NON-LINEAR ACTIVATION FUNCTION
For non-linear activation function (h), the integral in (46) in general cannot be solved in closed form. However, one may perform cumulant-expansion of the cumulant- -genrratngg function (46) in terms of the first two cumulants of ø as
W(f.. f+)= = N In
exp f (T)())
(49)
whha
= N In
exp
1 wm? (w f m
fT(ioo0)))
2
where we introduced the short hands for the cumulants of á as
ma = (ó (
(50)
((dd) = (6(ha) (he (ba.0))))(00,0om)))
Taking the expectation over w of (49) yields
W(FD. f4) =
N
In
Iw
fo ((oo)
(51)
2
NY
+ 2
N
[
m
2
N7
the
For moint-symmetric activation functions o(-h) -(h). such as erf or tanh activation, the mean m 0 vanishes and comparing (48) the only replacement that appears is ((ax) C(00)
A.2. Tree-level approximation
To compute the output statistics, one technically requires the exact effective action (42). However, in general it does not have an analytical solution and we instead determine systematic expansion. well-established method from both statistical physics and quantum field theory is the loopwise expansion (Helias & Dahmen, 2020), expands the effective action I'(f) in terms fluctuations around its mean value f. The lowest-order term of the loopwise expansion is called the tree-level approximation, which hence corresponds standard mean-field approximation: one replaces fo by its mean fin the action itself
(m.((.ilv)) --S(f,i-lv)
(52)
=ayafa
Ko "2"
N NN ja- -
(53)
The average value of fa is given by the equation of state (43) of the effective action
(f.
0.
(54)
0fa
1.=0
In mean-field scaling (7 and for N + this result becomes exact using the Gäärter--lllis theoreom: the output cumulant-generating function W in (5) has scaling form
iyaF + Ko fafa =
(55)
2N
17

---

Multi-Scale Adaptive Theory of Feature Learning
with (k) iyaka + sa kak W(ik, j.). Thus, we can approximate the probability distribution of network outputs as (Touchette, 2009)
-pyy R (f.j.l.)))))
(56)
Due to the strong suppression of fluctuations in mean-field scaling with N oc, the tree-level approximation is sufficient to describe the network behavior and in particular
lim - P(y/C(E)))/N (f..l...1)) 10000
(57)
However, in the case of larger output fluctuations as in standard scaling (? = 1), we need to take into account the output fluctuations systematically by including higher-order corrections to the tree-level result. We derive the leading-order correction the following section A.3.
A.2.1. LINEAR ACTIVATION FUNCTION
From the equation of state (54) we obtain self-consistency equation for fa
y. Iw ((axa) FfT (n) N [1
==
-i -NoNiNNN--I Cn c(an))
-1
(58)
(59)
N
Using the relation between the statistics of the discrepancies A and fo (38), we obtain for the training discrepancies
(Aa) = -x (W(kk,,l))
aka 10=00
(60)
=No (kOl + Cr. ((xx)
y3-
(61)
ag
For the test point, we get
(f-) (f.Jll)))))
(iF
Oj. 1i==0 0j.
11==0
(62)
9w Clar)
[I+
Iw
((ax) 871
-1
ifg.
(63)
N1-7
N
ag
By substituting the self-consistency equation for 8, we obtain
(f-)m. = -nn 8a (luw)
[ + CT. (î)
-11
U1
(64)
aa
For the covariance of the network predictors, we use that Do = ya - fa implies
((AnAs)) ((A. - (Aa -(aa)))
= ((fa (fa))(f (fo)) = ((fafa)).
(65)
so that we may obtain the covariance from (39)
82 (
((fos))) = No ~
Kig
N
2
iy
5. 0-
(66)
Using the involutive property of the Legendre transform, we can express the Hessian of W by the inverse of the Hessian of using 2y(2 = (() (112))-1,00 yielding
(Pr(AW)
(Jasz)) NoN
2--2
li-=0
(67)
0faî
18

---

Multi-Scale Adaptive Theory of Feature Learning
In free level approximation with FTL -S, then have
(oJas) = No + N222-22 (
-1
1i-00
(68)
Computing the Hessian of the action S, we get
((-S(f.1)) = KoOaB + Qul(f) Cas -o--oulle Qu(() Casfs Qu Coafe
(69)
0j.0f/
N
50=0
=A(F) as - N
2 on [ciic] aB
(70)
where the shorthand A() kall + (f) C. Using the relation between f and the mean deviations from (38)
(Aa) inofa, may rewrite this expression
02 (P--(((...) 0i.0f.
= A(F) a3 + (1) N=²
(obe
(DHA)
(ww) C
CD
(71)
5-
o8
which yields the expression from the main text (26) for the covariance of the network output on training data with N1-7 and hence
((ff)) = KI n² A(F +
20
F
(72)
Nk
Chor) with F F-C (AYA) ( (Ex) DD
A.2.2. NON-LINEAR ACTIVATION FUNCTION
Here we consider point-symmannn symeetrcm non-linear activation function, so that the mean m = O: the extension to the case with m # 0 straightforward. Due to the similarity of the expressions (48) and (51), the final expressions here have the same structure as in the case of the linear activation function, Eqs. (60) (62), but with the replacement ((ax) + (( throughout. Because these two cases throughout lead to identical expressions except for this replacement, in the following we denote ((2z) or ((+0) simply by C.
A.2.3. NON-LINEAR ACTIVATION FUNCTION BEYOND CUMULANT EXPANSION
In this section we consider a more fine-grained approach compared to the cumulant expansion in the previous section. Following results from previous sections, the self-consistency equation for f in a non-linear network is given by
i= - (K Fo/N
(73)
with
de ) f d' ho (h) exp (-
K= g=/NNN fdPhoxp (dPhexp(-tnTCC (-ANTClea))---
(74)
du ZN ZN= (T)()))))}))))
In the case of linear networks, the integrals in (74) are tractable, and are replaced in previous sections by their explicit solution, as can be seen in (47). In the case of non-linear networks, the value of K must be approximated, either via a cumulant-expansion as the previous section, via a variational Gaussian approximation (VGA) as (Seroussi et al., 2023), or via higher-order distributional approximations such as the variational Gaussian mixture approximation as in (Rubin et al., 2024), as well potentially richer approximations. Here we consider the VGA, but we emphasize that our framework is in no way limited Gaussian assumption.
The VGA defines matrix Σ which is chosen such that it minimizes the KI. divergence between the distribution of h as it appears in (74) and Gaussian distribution with covariance Σ. The matrix Σ is then determined by the following equation
-ー 15 = [c(am)] - -yT (Kx + no/ No/NT)
0Kr σE,
(Kx + ko/NI) y
(75)
19

---

Multi-Scale Adaptive Theory of Feature Learning
(a)
[b)
0.9
NNGP
NNGP
0.8
Rescaling
Rescaling
31. Leaambbiltyy
0.7-
Experiment
Adaptive
H6 Learnabiity
0.3
Experiment
Adeptive
0.8
0.2
0.1
30
P 120 240 30
120 240
Figure 5. g-learnability of target components in two-layer erf network, trained the target given in (77), where H; learnability is defined according to (78), with g(X) H;(Xw, As can seen in panel (b), the adaptive approach as derived using VGA predicts that the network will begin learn higher-order components of y P (D), due the manifestation of directionally dependent feature learning. On the other hand, kernel methods such as the NNGP and the rescaling approach predict that the network output will be lincar, and the cubic component the target would require P (D*) training samples. Parameters: Pies 4.000, N 1.000, 30,
Ru 2, yu 1, 1,== -0.1.
where : = := (1)))
(which is tractable for activations such as Erf, ReLU). Thus the equation for k~N((,,))
j can be written as
fVG = (K' + Ko/NT)
(76)
note that for a linear network, the VGA approximation exact, as the distribution of indeed Gaussian. The application of the VGA approximation allows for structural changes the kernel which do not emerge from a simple cumulant expansion. The structural encoding of certain directions in the kernel could allow the network to learn complex functions of these directions with significantly less data points than would be naively expected. Thus, even small non- linearity could lead to fundamentally different predictions by the adaptive approach compared to rescaling approaches. To demonstrate this, we consider a teacher-student setting, where the teacher is given by
y(z) = Hi (w) + eH?(wwa)).
(77)
with H1.3 being the first- and third-order Hermite polynomials, and æ (O,, We quantify the ability of a network learn the different target components (linear cubic) with parameter call learnability, defined that gllearnability is given by
(Xtea) 1 g (Xuss)
g g-eernaabiltyy
(78)
(Xres)) g (Xues)
where f. y, 9 are applied row-wise, and g is any function. Having learnability for given component implies that the network has successfully learned this component. In Fig. show the theoretical predictions as well experimental observations for the learnability of both target components (H1,a [o][z) x)). The experimentally measured learnability of the non-linear component is in good agreement with the adaptive predictions (obtained by VGA), which significantly outperforms the predictions of the kernel approaches (both rescaling and NNGP).
A.3. One-Loop corrections in standard scaling
While in mean-field scaling (r 2) the comulant-generatiin function has scaling form and the network outputs concentrate, we need to account for their fluctuations in standard scaling (> 1). In the following, we thus set = 1. To leading order, also called one-loop approximation, we have
Loooop (f. f.jz b) -.....)) 2
1 log det (((())))))
(79)
20

---

Multi-Scale Adaptive Theory of Feature Learning
The self-consistency equation for f then becomes
L (f. v)
as(
( -1 80 (-)-(()))))) 1 0.
(80)
0fs
50=0
0fe
j'=0 2
o
Ba 0f.0f00f.
5-=0
Given the form (48) of the cumulant- -ennerttggnn function W, taking the derivative of S with respect to f in (54) yields different expression than in the previous section
as(J,--)
=iy
KOSO + QTL C
(81)
0f,
j--0
(
9w
(82)
1+ N4 TCf
where now Qn. is a scalar. Note that in this form the tree-level equation for (62) can be written as
i= -i (KO (+) y,
(83)
thereby obtaining expression in which the imput kernel C' is only rescaled by scalar, which we call kernel rescaling expression. For the second and third derivatives, we obtain
0 (-S((f.)))) 0fa0fo
2 N
=nodas + +Qn Cas
Qn Casfs Qn (f) Cgafe
(84)
5-
A(F)
2 N
(n [ciFd] as"
(85)
0 (---(().))))
2
2 (i [CadCoefe + Cas Cacfa
(86)
0f.0500f0
N
5-
8 ( Can for Ip Is
(87)
N2
2
2
(h (Cas [c]] [css
Cas [c]] 3
+ Cas [c]]]
(88)
N
+ N2
Q (f)
[c]] c [c 5"
(89)
Here, we use the shorthand A(F) = KoI + Qn (f) C. Overall, we obtain
S(2) -1 (
- (A()))) -1
iye +
(90)
8c
2 a no 0f.0f00f. 1=
Similarly to the previous section, the training discrepancies are given by
(An) = inofa.
(91)
For the test predictor, we have
(f)) = OT
(92)
Oj
1==0
aw(F.j))
-1 0' (--S((
(((
(93)
0j*
0f.0500j.
5-=0
1==0
= (?) CuaFo - Σ
1
(2)) (-(2))
07 ((-.(...))))
(94)
30
0f30..
a3
=0=
21

---

Multi-Scale Adaptive Theory of Feature Learning
The appearing derivatives of the action are structurally similar but replace the training point by by the test point Ix yielding
((-S((.1))
Oh (Cas [ci]. + Car [co] Cg. [c]]c]
2
N
(95)
0fa0fz0j.
4==0
[ct] [o]] [c]].
(96)
When solving these equations, we backtransform to the imaginary variables f + if, which changes multiple signs and absorbs the appearing imaginary units.
A.4. Kernel rescaling approach
We here derive the results by Li & Sompolinsky (2021) and Ariosto et al. (2023) in our multi-scale adaptive theory: Using that Maj ~ (0, ((#x)) i.i.d. over the neuron index j. we can rewrite the cumulant-generating function W (30) for the case of a linear activation function (hh) = conditioned on readout weights w as
1 (mm) W (( (w) = In (exp(-[aww,taa)) haj = Cas aB Sa |w]|
(97)
where we drop the test point here to keep notation concise. The result for the test point will follow naturally later. Likewise, performing cumulant expansion up to second order in for point-symmetric activation function in (49), we obtain
W W(Dlww) In (exp((fu,,((hs))) has
1 fa c(ee) fa ||0
(98)
where C(++) is defined as in (50). As in the adaptive approach, we here again write for short to refer to C(ar)) in the case of linear activation function and to ((6)) the case of the non-linear point symmetric activation.
We observe that the readout weights only appear in the form of the squared norm ||| The distribution of the network output is hence
1 ply, fIC) = N KO) / dfo (exp - ifafa - 2 CanSs [1]l())) wi 15 "+)
(99)
Since both, the prior measure of the weights w ~ is only in the form of [|w]|,, we may introduce this quantity as an auxiliary variable, which we name Q := |/|| = Σ wn
x exp(N || and the explicit appearance of w.
and which corresponds to the Euclidean norm of the readout weight vector w. Note that, given ||w|| the integral over fo simply yields S'jwja ~ ~N(0, [=]]²C), so
ply. fIC) = - | (QNN(10.,CC)()))
(100)
Here the distribution of the squared norm is
p(Q)= = (81-Q + ||00|11) mi A0 zy
(101)
100
dQ
2 ( cxp (Øl - Q =lwxl31))
(102)
wz in NO0, 2F)
dQ
2tí exp ( 1-QQ+wal
(103)
where W(²)= = In (exp (Ô [|]])3))
is the cumulant-generating function of Q Using that the wy are i.i.d, we
w. 2n Ai00 )
get
W(Q) = In ( (exp (Q|(w(13))
(104)
N In [
29 m
Q
(105)
22

---

Multi-Scale Adaptive Theory of Feature Learning
where we performed the one-dimensional Gaussian integral over wc Up here, all steps are exact.
(100) shows that the auxiliary variable Q being scalar may only carry fluctuations of the overall scaling of the kernel and hence all descriptions and approximations in terms of Q can only change the scale of the kernel, which is consistent with the results in (Li & Sompolinsky, 2021; Pacelli et al., 2023).
1.4.1. APPROXIMATION OF NETWORK PRIOR FOR WIDE NETWORKS
One expects that Q concentrates for large N according to the central limit theorem since Q ||w]|/ Σ of with i.i.d. Wy (0, g//N) The cumulant- gennrraing function W can be written as a scaling form AN (k)) := N ((NN) k)) - In [1 2g and its limit N oo then exists trivially, so that we may approximate p(Q) with the Gärtner-Ellis
theorem (Touchette, 2009)
In p(Q) sup
QQQ+ + W(Q)
(106)
Q
N
[1
9w
Q
N
(107)
29w
2
N Q
Q -Q(
1
In
(108)
2 9w
Jw
Intuitively, by the scaled gumulant-generatie function of the form N N In 2g Φ the mean of is of order (Q) = (1) and all higher-order cumulants of Q are being suppressed by at least ON So on exponential scales, one may parametrize the probability by the mean, namely one obtains the distribution of Q from the rate function as e eQQ))
To obtain (106), the supremum condition has been used - Q + 9w 2g o solved for 25w Q 9u and
N
Q
Rw and inserted into the first line of (106) to obtain the second line. The rate function, being the Legendre 2gwr Q
.V. [
transform of W, obeys the equation of state
d r(Q) = -
N
[
9v Q
(109)
dQ
2gw
So the final expression for the joint probability of y and f. the network prior, is
p(y, = (Nlls / (Q)
(110)
=f iQ e (.Q)
where the action s)Q is
S(QIf.v) = 11y f1 P
In Ko
(111)
2k0
2
1 TT (QC)
-1 f
N In det (QC) - I(()) + const.
1
A.4 MAXIMUM A POSTERIORI ESTIMATE FOR Q
To obtain the posterior distribution for Q we marginalize (110) over the network outputs f. which yields
p(u|C)= / d$ p(v,
(112)
/ dQ exp (S(Q
which yields the action
1
S(Qw) - (QC + =ol) y In let (QC++ T(Q),
(113)
23

---

Multi-Scale Adaptive Theory of Feature Learning
and which reproduces Eq. All in Li & Sompolinsky (2021) after inserting the rate function (106) and using C ( for the linear network. It likewise reproduces Eq. (33) in (Ariosto al., 2023) when inserting C = ((éó) for the non-linear activation function.
When computing the maximum posteriori value QLs, it only depends on the numerator of
p(Quv) plylQ) p(Q) p(y)
(114)
since the form of (112) is p(y) = p()) comprises the maximum posteriori (MAP) QLs in p( =S(Q y) + const has the same stationary point as p(vlQ) pQQ).
(QQ). Thus, computing the Q-integral in saddle point approximation
The length Q = |!| in their theory is obtained by the maximum of (113), which is given by
0 - as aQ
1 2" T +ol) c + "1
y
(115)
N 1
-IrC + +xol)
2 Dw
This self-consistency equation yields the tree-level approximation for ||wi| 2 = QLs.
A.4.3. PREDICTOR STATISTICS
To obtain predictions beyond the length of the readout ||l]] we start from 112. We obtain statistics the training discrepancies A Va fa from
0
MAP d sup
Inp(y(C)
S(Q1y)
(116)
Bya
dy a Oy
as aQ
S(Qnsly)
(117)
aQ Oy Q=Q"1
where the derivative by Q vanishes because Qts has been determined by the supremum condition as the stationary point of the action. The partial derivative by yo only acts -y (QC =ooT) 9/2 in the expression for (113)
(Aa) = Ko (QtsO + kol) -1 y.
(118)
In consequence, the test predictor is identical to the NNGP predictor with different regularizer f
(f-Lss [C..] (C + k0/ ko/Qus!l) - 48.
(119)
To compute the variance of the predictor, we generalize(99) such that instead of the variance iog in N RoTl), we insert general covariance matrix K into the Gaussian measure N(y f, K) and perform an integration over f
/ / N(
P II
Ply)
C(a))) ((aa))
8 [fa
Σ wa o(hai))) w: N0, 9F), hus iid i Ni00,
(120)
The presence of the general matrix K allows us to measure the statistics of the discrepancies A, = = - Za, because writing the Gaussian (s).f, K) exp ( [ f) k- (y- f) + 1/2 In det (K (K-1)) explicitly we observe that derivatives by
[K-]] a yield
i
1
In lnp(y) Inp(y|K,C) K,C)
= - ((v Daly - Na) + Ko dag.
(121)
0[K]]
K=Nol
24

---

Multi-Scale Adaptive Theory of Feature Learning
With the manipulations that led (112) then has
/ faQ exp (S(Q)K,2))
C)
(122)
where the action, corresponding to (113), is
v. K) - - (C ( ()+ y -ー det (C+
N
In -1-
Q yv
l=
(123)
2
So in the approximation replacing Q by its MAP LLS we get
d
In 1n p,, |[K]ap
sup 3(Q
(124)
K=xoll
d[K]
K=xot
0
S(Q1sly: K)
(125)
o
K=NoT
where the inner derivative by 0s/aQ drops out due to stationarity at QtS. which is given by the solution of (115). The latter partial derivative evaluates to
S(Qly, K) 0|K
K c K] K|_ uy [e+ K|
K+ K (c+ KY- K
(126)
K=Rgl
2
as c
2 No 2
[c koo] -1 uvv |c+ ko=]
+ (c kol) -
(127)
2
a3 Qus
where we used that /KKK// - Kya Kas, which follows by symmetry from 0 /0Kag - -Kad K So the second moment of the discrepancies with (121) is
AA) = Nodn + Nḥ [c kol] yy [c kol] - (c+ kol)- aad
ais
= (Aa))As) + KodaB - n (c + Ko ko
(128)
c=
where used (118) in the last step. Because A - and the target label do not fluctuate, the latter terms in (128) the variance
(Aa: Ag)
fsl)
(129)
Kgdag ni (c kol)
(130)
cQus C
c-clc -
(131)
c
which is the usual expression for the variance of the NNGP predictor a Gaussian process with the kernel c QLs C.
A.5. Connecting kernel rescaling and adaptive approach
While the kernel rescaling approach holds in the proportional limit x P - oc, the one-loop approximation holds also for large but finite P.N N >I As we have seen in Fig. the main text, they yield almost identical results certain settings. By considering the proportional limit, we may connect these two approaches: some correction terms vanish in this limit, leaving only a scalar term.
To this end, we look the scaling of each correction term with both I. and N. We have
-2(2) = -ool + QT. (1x))
N Un ((uz) 7J clas)
(132)
Kol QT. (xx) O(IN))
(133)
25

---

Multi-Scale Adaptive Theory of Feature Learning
since ((=x))) + +*)
x y = (11) and thus also C(ax) j 0(1). Here, we drop the dependence of QnL on f for brevity,
The fluctuation correction given by
(2) 0 (-S(,.))))
(134)
2
Ba
05. 0f.050018
5-=0 (clss)
(-8)2) )Ba
[clen)]
j], + Clare [clex) F] +
Cg By
(
(135)
aB
a9
4
Un
Σ aB
(1)(2)
[(lex)] j], [clamt) }] [(tex)
(136)
N
ga
Looking at the individual terms, we have for the first term in the second line
1 Qi. Tr [MOOT+ C(az)
( + O(1/N))) -((m))
clan) fo OP/NN)
(137)
where the factor P results from the appearing trace. Assuming the regularization noise Ko to be small compared to the kernel ((z)), we see for the other terms that they scale as
1
Q
I(HOI + Q()( + +1(1NN)) (Chey) clar) fo + cl c 16)
4 Qu 75 8o = O(1NN), (138)
N
4
Σ
(MOl + Qn.C(an) + +0(1/N)) [c(em) i| (((mn) 7] [C(E))]] , -P(/N9)
(139)
N2
o
In the proportional limit P x N oc, only the first term does not vanish and the self-consistency equation for f becomes
iy + No. + Qn. clam) j-
1 N th (f) Tr [Noll + QTL ( -( (x(x) j=
(140)
yielding
1 iJ ( (1-()(11 - (f) Tr[ Qn (f) (((x))y-t ((xcer)))ccz)) -1
(141)
The rescaling factor is thus given by
QI-Loop= = Qn. N on Tr[ + QT
(142)
where QT. = Qn. depends on the self-consiste solution in (141). The tree-level solution is the leading term here and receives correction due to the output fluctuations. We cannot directly compare the expression for this rescaling factor to the one in (Li & Sompolinsky, 2021), since the latter is given by the self-consistency equation (113) and the former by the self-connistency equation (141) for the training discrepancies inserted into (142). Nevertheless, Fig. in the main text shows that these two expressions yield the same value and thus the same predictions for the mean discrepancies numerically.
B. Other scaling regimes
We here expose the relation of our theory to the work by van Meegen & Sompolinsky (2024) where yet another scaling regime is considered. We start from the effective action in tree-level approximation (53) and the form for W given by (46)
Tn((F.i-)) iyaaa+ + Ko N- N--jaaaa -W(f.j-)
(143)
wifo. f.) N In (fdu exp
P+1
Σ a=1
fa
w
((ha)
P
w²
const.
2gw ha (
where, as in the original work, we scale the variance of the readout weights by an additional factor P (see ref. (van Meegen & Sompolinsky, 2024), their Section "I. Weight Posterior" second last paragraph). In the proportional limit P x N it
26

---

Multi-Scale Adaptive Theory of Feature Learning
entails to & N 2, which scales down the readout weights even more strongly than mean-field scaling. This scaling allows taking the integral over in saddle point approximation because the exponent EeZ saw ((haa) P u² /2 scales with P. With (57) one therefore has with I = (f.j+)/N
In
(144)
K0 1T1
iy +
2
In / du
exp
w²
w [it (((oo()))))
m-Nio.ccss) + const.
2gw
PocN iy'
T
+
K0 [
P u
extr w
In
exp
(w [it ((ll)) hn.N( const.
2gw
The extremum condition may have multiple degenerate values wy that appear with the relative log probability given by the entropy
w In P. = -P + In
Wy [il [il (oTTo)) h=N(0. + const.,
exp
(145)
2gw
which corresponds to Eq. (B10) in (van Meegen & Sompolinsky, 2024) and the constant determined such that Σ, = 1. The distribution of w hence approximates the posterior p(w) Σn Py 5(w - w which used in (144) to obtain the equation of state (54) for ff by taking partial derivative of by la, which yields
Yo Σ Pyy W.,, [Qa]r, + kail"
(146)
which corresponds to Eq. (B17) in (van Meegen & Sompolinsky, 2024), where the expectation value [...]] [...]r,,,. 10 is given by
jd ... exp W T [(h) -
[TTC(mm)))))
(147)
h exp
w
)hTC(a)))))
The extremum condition appearing in (144) leads to stationary values of the weights:
Iw T [o(h)] P
w
(148)
which corresponds to eq. B16 of (van Meegen & Sompolinsky. 2024).
In summary, the difference to our work is that in the scaling considered in the work by van Meegen & Sompolinsky (2024), w x L/(N /P), the readout weights w tend to concentrate to non-fluctuating values w while in mean-field scaling 1/N, which we treat in the main part of our work, only the network outputs concentrate. Compared to the standard scaling 1/N both scalings only require tree-level approximation to describe the mean predictor whereas the fluctuations in the standard scaling require the inclusion of one-loop corrections.
C. Details of experiments
C.1. Self-consistency equations for numerics
In App. A, we derive train and test statistics in framework involving imaginary variables f. To solve the resulting self-consistency equations, we need account for their imaginary nature and substitute in all of the results above f if.
27

---

Multi-Scale Adaptive Theory of Feature Learning
changing various signs in the process. The final expressions read as follows: In tree-level approximation, we have
i == - [60] + Qn (fa) y,
(149)
(T)
9m
(150)
1 No
(T(())))
(Ahm = No0NN11 No (NOT= + Qnl(f) clear) y.
(151)
(f-)n. Qn. (f) Clan! (KOII + (n) ((2x) 2/.
(152)
In one-loop approximation, we have for the train discrepancies
E (
f = [A(F)]]
de Ve
+
2 a9
30
of ff.ff
(153)
A(f) NOll ( ((ex),
(154)
8
=a + 2 N
2
ou (f) ((ax)) (((()
(155)
0f.0f 50=0
(-S)
( (cl.ar) (clam) j] + c((x)
(((x)
+ +Clea)
[clax] }la))
(156)
0f.0f.05 f=
as
8 N2
9 Pi (J) [c(an) [clam) ((aa)
+
88 6¹
(Aa) = - Kap
Jar
(157)
and for the test predictors
89(-)) Ba
((')) = (f) claa) Ja E
(-5(()
S(2)
(158)
0fa0fo0j. 1-
aB
(-S-)) (f) (clez) (C(aa))
Clux) (( 5l
(159)
0f 5-=0
8
o oi [C(am) F] [c(aa) j
[((=))
+ N2
Finally, in the proportional limit P N + 0o this reduces to
j = (NOT+ + Q1 (1) (f)()z))
(160)
QI-Loop (f) (n) 1
N Qh Tr Tr(( + Qn (f)
(161)
(A) No (nTT
N1-77(011 + Q (f) ((oz))
y,
(162)
[Q1-Loop (f) claa)) (ROll + Q1-Loop (f) cloxa) -1 a v3.
1-Loooo rescaling
(163)
C.2. Numerical stability and computational complexity
To improve the numerical stability when solving these self-consistency equations, we use the following scheme: standard scaling, we use the train predictors given by the NNGP as starting values for solving the tree-level equations. Then, we use the tree-level solution as the initial value for the one-loop equations since the latter contain additional fluctuation corrections to the tree-level result, Further, we anneal solutions from the standard to the mean-field regime since the solutions are more unstable in the mean-field regime (see pseudo code in 1).
The computational complexity for solving these equations is O(P3) with a different fre-factor for tree-level and one-loop. Note that neither the input dimension D nor the network width N but only the number of training samples P affect the computational complexity.
28

---

Multi-Scale Adaptive Theory of Feature Learning
Algorithm 1 Annealing of solutions across scaling regimes
Input: data X, labels Y. scales {xi
Compute NNGP train predictors PNNGP from data X and labels Y. Set initial value to NNGP predictor fa JNNGP.
for in {xi): do
Set 9h0 + 9w/
Solve self-consistency solution for tree-level approximation iLL with initial value JANOP Solve self-consistency solution for one-loop approximation fo Fo- with initial value Fan. end for
C.3. Network tasks and training
Ising task We use linearly separable Ising task: Each pattern V. in the Ising task is D-dimensional and Tai {#1). If the pattern belongs to class -1, cach Tai realizes Zea +1 with probability of p1 0.5 Ap and the value Zai -1 with 0.5 + Ap. The value for each pattern element Tai is drawn independently. If the pattern belongs to class +1, the probabilities for aais = 1 and Yas = -1 are inverted. The task complexity decreases with larger Ap. We use Ap 0.1
throughout, corresponding to an oracle accuracy on the classification task of Poracle = 99, 78%.
Teacher-student task this setting, the target is given by a ya = wa Ta where Ta E IR D is standard normally distributed Foi (0, T).. The teacher direction w. RDD is chosen to é in the standard basis.
Network training We train networks using Langevin stochastic gradient descent (LSGD) as detailed in (Naveh et al., 2021) that the trained networks are effectively sampled from the posterior distribution (29). Here evolving network parameters O such weights V, with the stochastic differential equation
0,0(t) -pe(t) Ve((e(t):v) V2T((t).
(164)
(c))))(())) = 61,ō(1-1)) s),
with the squared error loss L(O;y) = Ei= (e) )= the network output for sample a = l...., P. leads to sampling from the equilibrium distribution for θ for large times t which reads
< S.C
unit variance Gaussian white noise, and (O) denoting
lim p (9(t)) ~ exp (- -
(165)
Using the Fokker-Planck equation (Risken, 1996) one can derive this density for θ Further, this implies a distribution on the network output
p(( & / de exp (-
(-
P
||012
T llf y)l²)
(166)
2T
a eep ( 7 ||ʃ -
1
@N(vlf.T/2) (6 f(0)] f
In fact, p(/IX) (8 (8[s) ( "5
leads to the posterior in (5) if one identifies Kn T/2 with the regularization noise and T/p= 9/N with the variance of the parameter θ₁, Implementing the sampling in practice this corresponds to requiring different weight decay for each parameter, as weight variances can differ in the input and output layer.
The time discrete version of (1 (164) is implemented in our PyTorch code as
θ = 0,-1 - 17 (p0,-1 + VeL(O.-1:)) + /2TTGe
(167)
(CCC.) = 6ist
with standard normal Çi and finite time step which can also be interpreted as a learning rate. To accurately reflect the time evolution according to (1 64) the learning rate 7) needs be small enough.
29

---

Multi-Scale Adaptive Theory of Feature Learning
Hence the LSGD implement corresponds to full-batch gradient descent with the addition of i.i.d. distributed standard normal noise and weight decay regularization (Krogh & Hertz, 1991). The value for Ko corresponds tradeoff in the optimization between the weight priors and the likelihood in terms of the loss L. Choosing large Ko corresponds to large T = 2ng and hence a large noise in the LSGD and therefore putting more emphasis on the Gaussian parameter priors. Small regularization values Ko favor the training data terms of the loss in the exponent of (165).
To faithfully compare the numerical results with our theoretical results, the LSGD needs to sample from the equilibrium distribution. For this it needs to be ensured that the distribution is equilibrated by evolving the networks for 10.000 steps We ensure uncorrelated network samples by initializing different networks with different random seeds.
For the Ising task, we average over Netwoork 100 with different initial weights to obtain the training and test predictors. For the teacher-student task, we average over Networks 5.000 with different initial weights to obtain the covariance of the network output projected onto different directions.
D. Additional figures
D.I. MNIST
Since the presented approach does not make any assumption on the data, is applicable arbitrary data sets. We here show results for binary classification MNIST across scaling regimes.
MNIST
0.25
Squared Errors
0.15
0.10
NNGP
Exp.
Saddle-Point Fluct. Corr.
Train
0.05
Test
0.00
10-2
10-1
10c
10
Fluctuations Scale 1/x
Figure 6. Binary classification on MNIST: training (solid line) and test errors (dashed line) across scaling regimes for different approaches.
While standard scaling (green shaded area) requires one-loop approximation with fluctuation corrections (Fluct. Comr), a saddle-point
or tree-level approximation (Saddle-Poinn is sufficient in mean-field scaling (orange shaded area). Parameters: Prain 80, N 100.
= 784, Ko = 1, Pes 10". 9. = I=
30

---

Multi-Scale Adaptive Theory of Feature Learning
Training discrepancies
(b) Test discrepancies
1.0
1.0
NNGP
NNGP
(Anoryy
Tree-Level
Anbory
0.5
Tree-Level
0.0
-0.5
-1.0
-1.0
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
(mmp)
(Amp)
Figure 7. (a) Training discrepancies (A) (f)) and (b) test discrepancies (A.) (f.) for binary classification MNIST in mean-field scaling. We show theoretical values identity. In contrast to the NNGP, the tree-level approximation accurately matches the empirical values. Parameters: 2, Punin 80, 100, 784, Ko 1, Plex 10a, 90 = gw 2.
NNGP and tree-level against empirical results, where the gray line marks the
(a) Training discrepancies
(b) Test discrepancies
1.0
1.0
NNGP
NNGP
Amerey
0.5
Tree-Level
0.5
Tree-Level
One-Loop
One-Loop
LS
LS
-0.5
-0.5
-1.0
-1.0
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
(Amp))
(Demp)
0.1
( (nwar))
0.1
Tree-Level
Tree-Level
One-Loop
One-Loop
-0.2
One-Loop simpl.
-0.2
One-Loop simpl.
L.S
L.
-0.3
-0.3
-1.00
-0.75
-0.50
-025
0.00
0.25
0.50
0.75
1.00
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
(A (wwar)
(AxnaP)
Figure 8. (a) Training discrepancies (A) = y- (/r) and (b) test discrepancies (A.) = y- (f.) for binary classification on MNIST standard scaling. Upper row: thcoretical values for different thcories against empirical results; gray line marks the identity. Lower row: difference of thcoretical values to the NNGP as baseline against NNGP predictions, indicating small-scale differences between different approaches. Results of the kernel approach by Sompolinsky (2021) shown reference (LS). Parameters: y
Puiin = 80, N 100, D = 784, Kg 1, Pes
10 9w = 2.
31

---

Multi-Scale Adaptive Theory of Feature Learning
D.2. Ising task in mean-field scaling
[a) Training Discrepancies
(b) Test Discrepancies
Tree Level
Tree Level
0.75
NNGP
NNGP
0.50
0.25
Sheery
0.00
0.0
A
0.25
-0.5
-0.50
-1.0
0.75
-1.5
-1.00
-0.75
-0.50
-0.25
0.00
0.75
1.00
-1.5
-1.0
-0.5
0.0 (Aemp)
0.5
1.0
(emmp)
Figure (a) Training discrepancies (A) = = (fb) (b) test discrepancies (A.) = y- (f-) on Ising task mean-field scaling. We show theoretical values for both NNGP and tree-level against empirical results, where the gray line marks the identity. In contrast to NNGP, tree-level approximation accurately matches empirical values. We here use non-linear activation function ø = erf.
Parameters: 2, Pemin 80, N 100, D 200, Ko 1, Pen 10 go Iw 0.5, Ap 0.1.
(a) Training discrepancies
(b) Test discrepancies
1.0
NNGP
1.0
NNGP
Tree-Level
Tree-Level
2 Se... : 2.6.
a
One-Loop
0.5
One-Loop
One-Loop Rescal.
One-Loop Rescal.
LS
LS
-0.5
-0.5
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
(Amp)
(Ampp)
Figure 10. Scatter plots of (a) training discrepancies (A.) = y.. - (fa) and (b) test discrepancies (A.) = y. (f-) on an Ising task in feeanfiiedd scaling. We show theoretical values for NNGP and different feature learning theories against empirical results, where the gray line marks the identity. In contrast to the NNGP, the tree-level approximation accurately matches the empirical values. Further, the different feature learning theories lie top of another in mean-field scaling. Parameters: 7 Punn 80, N 100, 200,
Ko 0.4, Pies 10², = 0.5, 0.2, Ap = 0.1.
32

---

Multi-Scale Adaptive Theory of Feature Learning
D.3. Coherent amplification of low-rank kernel structures
(a)
(c)
20
0.0
20
Ap
20-
NNGP
Adaptive
15-
15
15-
Rescaling
(0)))
Experiment
10
10
10-2
10
10
10
10a
0.00
0.05 0.10
Fluctuations Scale 1/x
Fluctuations Scale 1/x
Ising probability Ap
Figure 11. Relative directional feature learning on the Ising task as function of the fluctuation scale 1/x for (a) an input kernel without structure and (b) an input kernel with block structure (input kernels shown insets). Both NNGP and rescaling theory fail to capture directional feature learning, while the multi-scale adaptive theory accurately predicts network behavior. (c) An increase in structure in the input kernel inereases with the Ising probability Ap and lcads significantly higher dircctional feature learning in the adaptive
theory than in both NNGP and rescaling, matching the experiments. Parameters: Pixi = 80, N 100, D = 200, Ko 2, 0.01, gn 0.5.
D.4. Teacher-student task
(a)
(b)
(c)
Adaptive
Exp
Exp
NNGP
Theory
Theory
Rescaling Experiment
d= 9
b-6
d= y
1.5
3-
6
1.0
2
4
(c)
0.0
10-1
10-1
10"
-0.2
-0.1
0.0
0.1
0.2
-0.6
-0.4
-0.2
0.0
0.2
0.4
0.6
Fluctuation Scale 1/x
(f (f))
(S (M))
Figure 12. (a) Directional feature learning in teacher-student setting as function of the fluctuation scale 1/x. Both NNGP and rescaling theory fail to capture directional feature learning, while the multi-scale adaptive theory accurately predicts network behavior. Output distribution in different directions (b) in mean field scaling (x N) and (c) in standard scaling (x 128). Parameters: Punin = 80,
200, D 50, Ku = 2, g. = 0.01, I = 2.
33


Reviews:
[{"review_id": "i8NYYu8VKI", "reviewer": "Reviewer_YLkj", "summary": "This paper introduces a unifying theoretical framework which connects the kernel rescaling approach with the adaptive kernel approach in the Bayesian setup. The authors demonstrate that both of these approaches can be derived from the same starting point (the network’s posterior) but differ in the choice of order parameters in an effective action formulation. Specifically, the paper explores two different initialization scalings:\nMean-Field Scaling: In this regime, network outputs concentrate, and feature learning is captured through a tree-level (saddle-point) approximation. While the kernel adapts directionally in principle, for a linear network’s mean predictions, this adaptation can be effectively approximated by a simple rescaling of the NNGP kernel.\nStandard Scaling: Here, fluctuations in network outputs become significant, necessitating a one-loop approximation. The authors demonstrate that in the proportional limit of N and P, even these corrections can be expressed in a rescaling-like form for the mean predictor, although the kernel undergoes non-isotropic modifications in key directions.\nExperiments using linear networks on synthetic tasks validate these theoretical insights. Notably, the study underscores how this new framework effectively captures directional feature learning effects in the covariance of network outputs, which conventional rescaling approaches fail to account for.", "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": "We thank the referee for their positive evaluation and for the precise summary. In a revision, we will address the three mentioned weaknesses:\n\n1. Our work aims at a feature learning theory from first principles. We hope that in the long term this theory can be leveraged to build a theoretical foundation for mechanistic interpretability. We will provide three additional results in that direction:\n\n    (a) As the theory makes no assumption on the distribution of the data, we will show in a revision that the theory makes accurate predictions for train and test predictor also on more practical datasets (MNIST, see [Fig. 1 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)).\n\n    (b) We will include evidence of sample complexity reduction in the presence of feature learning in non-linear networks (see [Fig. 3 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)). To this extent, we consider a teacher-student setting where the teacher is of the form $y(x)=H_1(w_*\\cdot x)+\\epsilon H_3(w_*\\cdot x)$ with $H_{1,3}$ being the first and third-order Hermite polynomials. In this setting the adaptive theory correctly predicts that the network will learn the non-linear components at $P\\sim\\mathcal{O}(D)$, whereas both the NNGP and the rescaling theory capture only the linear component. This result further demonstrates that the adaptive theory captures feature learning phenomena beyond rescaling.\n\n    (c) The theory explains the collaborative effect of input ($x$) and output ($y$) statistics in shaping the resulting kernel and thus the anisotropy of the predictor's variance. In particular, we show that low-rank structures in the input kernel $\\propto \\epsilon$ get amplified by the number of training patterns $P$, if the labels $y$ contain coherent information (see [Section 4 and Fig. 4 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)).\n\n2. The presented framework is indeed general and allows the inclusion of non-linearities. The difficulty lies in approximating the cumulant-generating function $W$ in Eq. (4) of the main text: In a revision, we present two approaches. (a) A cumulant expansion that for point-symmetric activation function leads to a simple replacement of $C^{(xx)}$ by a different matrix $C^{(\\phi\\phi)}=<\\phi(h) \\phi(h)>_h$, while general non-symmetric non-linearities yield additional terms (see [Section 2 for theory and Fig. 2 for empirics in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)). This allows us to obtain the rescaling theories by Li & Sompolinsky (2022) and by Pacellli et al. (2022) within our framework. (b) A variational Gaussian approximation on $h$ yields results that differ qualitatively from the rescaling approach, leading to a reduction of sample complexity as discussed in the previous point (see [Section 3 and Fig. 3 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)).\n\n    We will also add an appendix showing how to derive the seminal results by van Meegen and Sompolinsky (2024) on coding schemes within our framework; in brief, their results are valid in “super-mean-field” scaling for the readout weights $w\\propto1/(N\\sqrt{P})$, which allows them to make an additional saddle point approximation on the readout weights (see [Section 5 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)). We believe that this unification is useful to the community to better compare the different variants of theories and understand their limitations.\n\n3. We agree with the referee that the Bayesian approach is in general different from practical ways of training. However, while Bayesian sampling does not exactly correspond to SGD, under certain conditions and with appropriate parameter tuning, it can serve as a useful approximation of SGD. This approximate correspondence builds on the noise inherent in SGD to explore the parameter space in a manner that aligns with Bayesian principles.", "claims_and_evidence": "Their claims are convincing.", "methods_and_evaluation_criteria": "Methods: The use of a Bayesian field-theoretic framework is well-justified given the theoretical focus. Restricting experiments to linear networks makes the math tractable.\nEvaluation Criteria: Measuring both the mean and covariance of the posterior outputs on carefully chosen synthetic tasks is appropriate.", "theoretical_claims": "I did not review the supplementary material, where the detailed theoretical analysis is presented. The theoretical results and equations in the main paper look decent.", "experimental_designs_or_analyses": "he authors use synthetic tasks for which analytic predictions exist, train networks using Langevin stochastic gradient descent, and systematically compare theoretical vs. empirical mean and covariance statistics. While the scope is necessarily limited to linear networks and synthetic data (Ising tasks and simple teacher-student setup), it is appropriate for verifying the paper’s key claims about the theory of feature learning across scaling regimes.", "relation_to_broader_scientific_literature": "This paper bridges separate lines of work of kernel rescaling and kernel adaptation by showing they can be viewed with a field-theoretical approach but differ in the choice of order parameters. This unified viewpoint is a valuable contribution to the literature.\nSpecifically, regarding the kernel rescaling theories, the authors demonstrate why and when the kernel-rescaling perspective is valid (e.g., for the mean predictor in linear networks, especially in the proportional limit). They identify the dimensionality of the order parameter (just a scalar vs. high-dimensional) as the crux behind whether the kernel is rescaled or adapted (to specific directions).", "essential_references_not_discussed": "I am not aware of papers which should be cited/discussed.", "other_strengths_and_weaknesses": "Strength:\nThis paper creatively combines statistical physics methods with existing neural network kernel/rescaling theories to derive a unified perspective for kernel adaptation in Bayesian setup.\nTheir numerical experiments with toy tasks beautifully agree with the theoretically derived quantities.\n\nWeaknesses:\n1.Unclear Practical Benefit\nWhile the multi-scale adaptive theory elegantly unifies kernel rescaling and kernel adaptation, its practical utility remains unclear. One would hope for novel or stronger predictive insights. For example, phenomena previously unexplained that this theory can now elucidate. but the paper does not provide explicit examples of such breakthroughs.\n2. Limited to Single-Layer Linear Networks\nThe work focuses on single-hidden-layer linear networks, which is a highly constrained setting. It is not evident how to extend the framework to more realistic architectures (e.g., multi-layer or nonlinear) without incurring significant complexity. This limits applicability to modern deep-learning scenarios.\n3. Bayesian Setup vs. SGD Practice\nThe theory is developed under Bayesian posterior sampling assumptions, but most modern neural networks are trained via (variants of) stochastic gradient descent. It is not straightforward how the results might carry over to typical (non-Bayesian) training regimes, potentially reducing the theory’s direct impact on practical deep learning methods.", "other_comments_or_suggestions": "No suggestions.", "questions_for_authors": "No questions.", "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": "draG4Xf99M", "reviewer": "Reviewer_Dpc2", "summary": "This paper aims to provide an account of feature learning in linear Bayesian neural networks that bridges the gap between results emphasizing isotropic rescaling and anisotropic reshaping of kernels.\n\n## Update after rebuttal\n\nAfter the authors' rebuttal, I weakly recommend acceptance.", "questions": "", "limitations": "", "rating": 3, "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 the reviewer for their feedback.\n\n**Claims and Evidence**\n\nWe will add experiments on MNIST in a revision (see [Fig. 1 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)). We also extend the formalism to non-linear activations and include experiments (see Fig. 2 & 3 in Supplement; for details, also see reply to point 2 to reviewer YLkj and Sections 2 & 3 in Supplement).\n\n**Choice of title**\n\nWe regard the formalism that is valid across all scaling regimes as a main merit of our work. Adequate fluctuation corrections are one important technical aspect, but we fear this term is not an easy to comprehend title; instead we will clarify in abstract and introduction that we perform systematic expansions in the output fluctuations. We further add a derivation of the theory by van Meegen & Sompolinsky 2024 within our framework: It results by an additional saddle point approximation justified in their scaling $w\\propto 1/(N\\sqrt{P})$ (see Section 5 in Supplement). Thus our work covers the entire range of scalings $w\\propto 1/(N\\sqrt{P})...1/\\sqrt{N}$ and explains the relation between kernel rescaling and kernel adaptation.\n\n**Improvement of accuracy of citations**\n\nWe are glad to hear that the reviewer in general agrees to our presentation of prior work and are grateful for the detailed and helpful suggestions to improve this section further:\n\n• We agree that feature-learning either arises at infinite width or in infinite-width limits different from the NNGP; we will change l. 122 as 'However, the NNGP cannot explain the often superior performance of finite-width networks [...], requiring either the inclusion of finite-width effects or different infinite-width limits such as $\\mu$P scaling (Yang et. al, 2021, Vayas et. al, 2023)'.\n\n• We agree that Canatar & Pehlevan (2022) is better placed below l. 129 as: 'These differ in the choice of order parameters considered and also in the explained phenomena. An experimental investigation of kernels in feature learning in gradient descent settings was performed by Canatar & Pehlevan (2022).'\n\n• We will include Hanin & Zlokapa (2024) in the list of perturbative approaches in l. 149.\n\n• We will change the citation of Cui et. al. (2023) in l. 139ff. to \"Cui et al. (2023) study non-linear networks [...]. \". Concerning Maillard et. al. (2024) we feel that the current citation is appropriate since this work takes into account a different scaling in the amount of training data.\n\n• We already cite Aitchison (2020) in l. 121; we will further include it in the enumeration of adaptive approaches, e.g. in l. 127. The additional reference of Zavatone-Veth & Pehlevan (2021b) will be included as “Zavatone-Veth & Pehlevan (2021b) investigate deep linear networks in different proportional limits; recovering the results from Li & Sompolinsky in an adaptive approach.\n\n• We already cite Hanin & Zlokapa (2023) in l. 138 and will add there: 'Zavatone-Veth, Tong & Pehlevan (2022) study the same setting but consider explicit models on the input data in the limit of infinite pattern dimension.'\n\n• In the revised manuscript, we discuss the relation to van Meegen & Sompolinksy (2024) in the related works and show how to obtain their theory within our framework. For details please also see our reply to reviewer QcC4.\n\n**Computational complexity and numerical stability**\n\nSolving the tree-level self-consistency equations requires $\\mathcal{O}(P^3)$ operations and the one-loop self-consistency equations $\\mathcal{O}(P^3)$, with a difference by a pre-factor independent of $P$. Note that neither the input dimension $D$ nor the network width $N$ but only the number of training samples $P$ affect the computational complexity.\n\nA naive implementation of the self-consistency equations indeed faces problems. Our stable implementation of the one-loop equations first solves the tree-level equations with the NNGP as initial value and then uses the tree-level solution as the initial value for the one-loop equations. The tree-level and one-loop equations are more unstable in the mean-field regime; we resolve this by annealing solutions from the standard to the mean-field regime. We will include this in App. B of a revision.\n\n**Further points**\n\nWe will carefully incorporate the comments on text improvements in the revised manuscript; in particular in l. 362-364, clarifying that meaningful directions are either the teacher direction or dominant eigendirections of the input kernel that align with the target. We will also explain additional insights due to this result: encoding of significant directions allows for more effective learning by reducing the sample complexity. Please see response to reviewer YLkj for details. We will rename App. A to 'General approach to train and test statistics' as the common starting point for train and test statistics in linear and non-linear networks across all scaling regimes $w\\propto 1/(N\\sqrt{P})...1\\sqrt{N}$.", "claims_and_evidence": "Most of the authors' claims are well-supported, but there is room for improvement of clarity. I elaborate on this issue in various places below. \n\nOne issue I will note up front is that the authors should emphasize in the Abstract and Introduction the fact that they focus on linear networks.", "methods_and_evaluation_criteria": "The analytical methods used are standard in statistical field theory, and are appropriate here.", "theoretical_claims": "This is a physics paper using standard methods from statistical field theory, and the derivations appear correct (though I have not checked them line-by-line).", "experimental_designs_or_analyses": "The experiments largely seem well-executed, but the choice of datasets is a bit disappointing. The paper would be enhanced by some experiments on real data, even subsets of MNIST or CIFAR-10 as in Baglioni et al. PRL 2024, Zavatone-Veth, Canatar et al. NeurIPS 2021, or Seroussi et al Nat Comm 2023. I don't expect experiments at the scale of Izmailov et al 2021 (https://arxiv.org/abs/2104.14421), but the authors should be aware that such things are possible.", "relation_to_broader_scientific_literature": "I think the paper generally does a good job of situating itself relative to past works on the statistical physics of Bayesian neural networks, though the accuracy of its citations to past works could be improved (see below).", "essential_references_not_discussed": "The discussion of related works in Section 2 is not complete, and incorrectly describes some of the works which are cited. \n\n- The claim of the \"often superior performance of finite-width networks\" made in the opening paragraph of Section 2 is out of date *in the context of gradient descent trained networks* thanks to work on mean-field parameterizations where the infinite-width limit can be taken without losing feature learning. See for instance Yang et al. (NeurIPS 2021) https://arxiv.org/abs/2203.03466, or Vayas et al. (NeurIPS 2023) https://arxiv.org/abs/2305.18411. Here, larger networks generally perform better as increasing width suppresses spurious predictor variance due to initialization. \n\n- The citation to Canatar & Pehlevan (2022) on Lines 134-135 is somewhat misplaced, as they use SGD in their experiements, and do not take an explicitly Bayesian perspective. \n\n- Lines 136-140 do not make clear the close relationship between Zavatone-Veth & Pehlevan (2021) and Hanin & Zlokapa (2023): the former paper characterizes the prior in terms of Meijer G-functions, and the latter leverages that result to characterize the variance in the predictor at zero temperature. The authors should also cite Hanin & Zlokapa's recent perturbative extension of that work to nonlinear networks in https://arxiv.org/abs/2405.16630. \n\n- Cui et al (2023) arises as an attempt to rigorously justify the Gaussian approximation of Li & Sompolinsky (2021) in a mathematically tractable setting. Also, Maillard et al. (2024) is in a sense a follow-up to this work. And, a key feature here is that these works allow one to study nonlinear networks and tasks. \n\n- I think the authors should make a greater attempt to connect their results in the linear case to those of Aitchison (2020), whose work emphasizes the adaptive aspect of feature learning and is in the limit he considers exact. Also relevant in the linear case is Zavatone-Veth & Pehlevan (https://arxiv.org/abs/2111.11954) who derive exact expressions for the predictor mean and covariance, generalizing Li & Sompolinsky and Aitchison. \n\n- When discussing the properties of linear networks in the proportional limit, the authors should reference Hanin & Zlokapa (2023) as well as Zavatone-Veth, Tong, & Pehlevan PRE (2022; https://arxiv.org/abs/2203.00573), both of which furnish sharp asymptotics for the generalization error. The latter paper is also relevant because it compares that performance against sharp asymptotics for a random feature model of matched architecture.\n\n- The authors should cite and discuss van Meegen & Sompolinsky (https://arxiv.org/abs/2406.16689), which considers mean-field parameterization. Here there is a matrix-valued order parameter, and the authors find drastically different properties of the posterior depending on the nonlinearity.", "other_strengths_and_weaknesses": "- This is partially a matter of taste, but I think the clarity of the paper would be improved by giving it a title that more clearly states what you actually have done. For the same reason, I am not in favor of the choice to describe the results as a \"multi-scale adaptive theory\". It would be more illuminating to state in plain words what you have done: a standard expansion in fluctuations of the predictor. \n\n- The clarity of the paper could be improved. In particular, there are many overly long sentences, like that starting with \"Another line of work...\" on Line 161. \n\n- In the Introduction, the authors should more clearly emphasize that they focus on linear networks. I don't view the specialization to linear networks as a flaw, but it does limit the generality of the authors' \"Universal theory of train and test statistics\" (as they claim in Appendix A).", "other_comments_or_suggestions": "- It might be worth noting that the \"directional feature learning measure\" introduced in eq. (27) is closely related to measures of feature learning based on centered kernel alignment, but differs in that those measures usually consider hidden-layer statistics of the trained network rather than output statistics. \n\n- Following on my earlier comment about writing: the sentence \"The appearing matrix product allows a non-trivial change of the NNGP kernel in certain meaningful directions, yielding additional insights\" in Lines 362-364 is not meaningful, because it does not clarify what those insights are.", "questions_for_authors": "I'm curious about the computational complexity and numerical stability of evaluating the one-loop corrections, as these have been issues for previous perturbative approaches to feature learning (see discussion in Bordelon & Pehlevan 2022). Could you provide some comments along these lines, at least in the Supplement? These issue could raise particular challenges in nonlinear networks, where higher moments of activations cannot be computed analytically. This is a key advantage of the Bayes-optimal setting as studied by Cui et al: the nonlinear case is not so much more challenging.", "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": "y8hlMgrocU", "reviewer": "Reviewer_QcC4", "summary": "The paper addresses the problem of Bayesian learning with a wide two-layer linear network. It develops a formalism which bridges between prior works, and in particular covers in unified manner the rich (mean-field) and lazy (standard parameterization) regime. The analysis consists in rephrasing the problem as the computation of the legendre transform of a generating function, which can be approximated using expansions in the various considered regimes. The theory is corroborated by numerical simulations across the various scalings. Notably, the mean predictor is found to be the same as a rescaling of the NNGP kernel. However, the covariance of the predictor is anisotropic, and reveals feature learning.", "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": "We thank the reviewer for their overall positive feedback and are confident to address the raised points below.\n\n**Application to real-world data sets**\n\nThe theory indeed applies to arbitrary data, since it only requires the input kernel $C^{(xx)} = \\frac{g_v}{D} X X^{\\mathsf{T}}$ without specific assumptions on the task. While we focused on linearly separable tasks in the initial version of the manuscript, we add results for MNIST in the revised manuscript (see [Fig. 1 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)). Further, we extend the formalism to include non-linear activation functions in an approximate manner (see reply to reviewer YLkj and [Fig. 2 & 3 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)).\n\n**Relation to van Meegen & Sompolinsky 2024**\n\nWe thank the reviewer for pointing out that we accidentally missed to cite van Meegen & Sompolinsky (2024). We sincerely regret this oversight. The main distinction is their mean-field scaling of readout synapses $w\\sim 1/(N\\sqrt{P})$, which differs from ours $w\\sim 1/N$. In the proportional limit $P\\propto N$, they thus effectively consider a “super mean-field scaling” $w\\sim1/N^{3/2}$. As a result, not only the network output but also the posterior readout weights concentrate, justifying an additional saddle point approximation on $w$. The latter enables them to find spontaneous symmetry breaking, the mechanism behind the seminal finding of coding schemes they describe. We will explain this relation in a revision. To expose it on a formal level, we derive their theory in our framework (see [Section 5 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)) and will add it as an appendix in a revision. We share the reviewers assessment on the high relevance of van Meegen & Sompolinsky (2024) as a crucial work in the field and a highly relevant reference for our manuscript. We will include it prominently in the following places:\n\n• p. 1, l. 31ff: \"However, the NNGP does not capture FL, which emerges at finite network width as well as in the proportional limit, where both network width and sample size jointly tend to infinity (Li & Sompolinsky, 2021), or in certain scaling regimes (Yang et al., 2024). Hence the NNGP also fails to capture the networks' nuanced internal representations that arise from feature learning (*vMS, 2024*).\"\n\n• p. 2, l. 106: “In contrast, adaptive theories of FL (Roberts et al., 2022; Seroussi et al., 2023; Bordelon & Pehlevan, 2023; Fischer et al., 2024b; *vMS, 2024*) ...\"\n\n• p. 3, l. 121 \"(*vMS, 2024*) considers the proportional limit in a regime where weight variances scale as $1/N^3$ (see Appendix for a derivation of their theory in our framework and a comprehensive comparison of the approaches).\"\n\n• p. 4, l. 166 \"Following along the lines of (Segadlo et al. 2022a, Fischer et al., 2024b, *vMS, 2024*)...\"\n\n• p. 6, l. 324 \"...and adaptive theories (Naveh & Ringel, 2021a; Seroussi et al., 2023; Fischer et al., 2024b; Rubin et al., 2024, *vMS, 2024*)\"\n\n**Similarity of mean-predictors in the rescaling and the adaptive theory**\n\nWe agree that the mean predictors for adaptive and rescaling in Fig. 1 and Fig. 3 of the manuscript are close. This point is precisely the question we target in this manuscript: How can these apriori very different theories be reconciled? We therefore show explicitly in Section 5 that the rescaling theory can be derived from the adaptive theory for the mean predictor. We then show in Section 6 that beyond mean predictors, the adaptive theory captures directional aspects of feature learning that escape rescaling theories. In a revision we will strengthen this point by demonstrating that only the adaptive theory captures the emergence of structure in the kernels driven by coherent statistics between input $x$ and target $y$ (see [Fig. 4 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)). We will also show that for nonlinear networks the predictions of the different theories differ qualitatively (see reply to reviewer YLkj and [Fig. 3 in Supplement](https://drive.google.com/file/d/1i2bhJbJYuJIvyNj12eWLjyAhrRlmCf9I/view?usp=sharing)).\n\n**Reviewer's specific questions**\n\n• In Fig. 4, the curves are computed from Eq. (27) with Eq. (29) for the rescaling theory. The variance for the rescaling theory is calculated in Eq. (116)-(118) in App. A.4.\n\n• While the tree-level solution and its equivalence to rescaling applies only to $\\gamma=2$, the one-loop theory applies to the full regime $\\gamma \\in [1,2]$ and thus also the equivalence to rescaling as discussed in Section 5. Note that this equivalence is restricted to the mean predictor, while the adaptive theory capture additional aspects, like directional feature learning in the variance terms (cf. Section 6).\n\nIn a revision we will clarify these points in the main text.", "claims_and_evidence": "The theoretical claims are supported by numerical experiments in all figures. One possible source of concern is that the results of the proposed approach at one-loop level and those of the rescaling approach (Li and Sompolinsky, 2021) are very close in Fig.1, 3, making it hard to tell whether one significantly matches better numerical experiments. It would be informative to, if possible, provide a figure in a setting where the two approaches are farther from each other.", "methods_and_evaluation_criteria": "The paper is primarily theoretical in nature. All setups considered are stylized. As a minor suggestion, since to the best of my awareness the theory makes no distributional assumptions, would it be feasible to include experiments on simple real datasets e.g. MNIST? If not, what are the barriers?", "theoretical_claims": "I did not check the technical parts in detail.\n\nI have a few questions regarding the theoretical parts, which I list below.\n- regarding Fig. 4 : from which expression is the curve for the rescaling approach (pink) evaluated ? I have only found a computation of the predictor mean and variance in the Appendix, but it is possible I might have overlooked a part.\n- Does the equivalence with the rescaling approach in the proportional regime hold for all $\\gamma$ or just for $\\gamma=1$?", "experimental_designs_or_analyses": "I have not identified any particular issue with the experiments.", "relation_to_broader_scientific_literature": "The main contribution of the paper is to reach a formalism holding across various previously considered scaling limits, in particular the lazy (e.g. (Li and Sompolinsky, 2021, Pacelli et al., 2023) ) and rich limits of Bayesian learning. This allows authors to recover a number of previous expressions, e.g. (Seroussi et al., 2023). Appendix A.4 in particular connects their formalism with that of (Li and Sompolinsky, 2021). As such, I believe it could be of interest to researchers working on this topic. \n\nAs a comment, the paper focuses on shallow, linear networks. I believe this is an important precision which should be at least mentioned in the abstract and main contributions, as this limitation is not shared by all the related works.", "essential_references_not_discussed": "The authors crucially do not mention (Van Meegen and Sompolinsky, Coding schemes in neural networks learning classification tasks, 2024), despite sharing a large closeness in topic, and in some claims. This reference is not a concurrent work. For example, the fact that the mean predictor of linear networks in the rich regime still coincides with a rescaled NNGP kernel is already observed in this reference, yet is not mentioned in the manuscript under review.", "other_strengths_and_weaknesses": "I believe the results can be of interest, as they connect a number of prior works, but have limited confidence in my assessment. Due to the concerns I listed above, and more crucially the key missing reference, I am in favor of not accepting the current manuscript, but am happy to increase my score if the authors address and clarify my questions, and revise the manuscript accordingly.", "other_comments_or_suggestions": "I do not have further comments or suggestions.", "questions_for_authors": "I do not have further questions for the authors.", "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": {}}]

Author rebuttal:

Meta-review: All reviewers recommend acceptance. The paper develops a theory of feature learning in Bayesian linear networks that bridges kernel rescaling and kernel adaptation in different scaling regimes. Their approach uses methods grounded in statistical field theory, and provides consistent predictions in both standard and mean-field limits. Numerical experiments on synthetic tasks and MNIST validate the theory. While the framework currently targets linear Bayesian networks, the authors explain how approximate Bayesian sampling can inform SGD-based neural networks. They have also clarified relationships with prior work and extended the analysis to nonlinear activations.
Decision: Accept (poster)