From Kernels to Features: A Multi-Scale Adaptive Theory of Feature Learning

We thank the reviewer for their overall positive feedback and are confident to address the raised points below.

Application to real-world data sets

The 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). 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).

Relation to van Meegen & Sompolinsky 2024

We 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) 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:

• 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)."

• 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) ..."

• 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)."

• p. 4, l. 166 "Following along the lines of (Segadlo et al. 2022a, Fischer et al., 2024b, vMS, 2024)..."

• p. 6, l. 324 "...and adaptive theories (Naveh & Ringel, 2021a; Seroussi et al., 2023; Fischer et al., 2024b; Rubin et al., 2024, vMS, 2024)"

Similarity of mean-predictors in the rescaling and the adaptive theory

We 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). 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).

Reviewer's specific questions

• 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.

• 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).

In a revision we will clarify these points in the main text.

We thank the reviewer for their feedback.

Claims and Evidence

We will add experiments on MNIST in a revision (see Fig. 1 in Supplement). 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).

Choice of title

We 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.

Improvement of accuracy of citations

We 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:

• 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)'.

• 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).'

• We will include Hanin & Zlokapa (2024) in the list of perturbative approaches in l. 149.

• 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.

• 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.

• 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.'

• 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.

Computational complexity and numerical stability

Solving 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.

A 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.

Further points

We 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}$.

We thank the referee for their positive evaluation and for the precise summary. In a revision, we will address the three mentioned weaknesses:

1. 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:

   (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).

   (b) We will include evidence of sample complexity reduction in the presence of feature learning in non-linear networks (see Fig. 3 in Supplement). 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.

   (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).

2. 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). 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).

   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). We believe that this unification is useful to the community to better compare the different variants of theories and understand their limitations.

3. 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.