We thank the reviewer for their precise summary, for pointing out the strengths and weaknesses of our work, and for the insightful and helpful comments we received from their review. These have considerably helped us to clarify the message. We address below each and every comment.

Weaknesses:

• The physical picture should be made clearer, to help non-physicist readers. For example, it would be nice to have a figure illustrating what you mean by "droplet" in terms of physics, and in terms of grokking.

Response: Thank you for your suggestion (which is now included in the text). Note that following some similar comments from the first and fourth referees, the manuscript went through a major presentation overhaul. In particular, we delegated many of the model-specific derivation details to the appendix, provided a hopefully self-contained introduction to phase transitions, summarized the relevant parts of the kernel-adaptation work, and last but not least, separated our results from the derivation and presented them in a more concise way. We further provide much more compelling numerical evidence.

In addition, following your suggestion we added Figure 1. which shows the phase diagram of the model illustrating the equivalence between the feature learning phases we find and gas water, and vapor phases. We sincerely hope that the referee will give the manuscript another fresh look.

• The link to NN training is unclear. How to translate the theory to NN training results?

Response: Our analysis provides predictions on the output of a network trained using full-batch gradient descent at a vanishing learning rate with white noise added to their gradients (see section 2.3 in the revised text). Our results are in very good agreement with numerical experiments on actual neural networks (see revised numerical Fig. 2).

• The fact that the analysis only applies to two-layer networks is not satisfactory but understandable.

Response: The kernel adaptation formalism we use has been demonstrated to work well on deeper neural networks as well. Hence, we believe our approach, which is an extension of kernel-adaptation, can be extended to say three-layer networks at the price of tracking a few more parameters. However this would make the analytical solution significantly more complicated, and would obscure the general picture.

• Missing some literature review on grokking, e.g., "Omnigrok: Grokking Beyond Algorithmic Data" and "Explaining grokking through circuit efficiency".

Response: We agree and these references were added to the introduction.

Questions:

• In 3.1, "is kept fixed" $\sigma^2 / n$. Do you mean $\sigma^2 \sim n$ in scaling? Wouldn't the result of Cohen et al. rather hold for $\sigma^2 = O(1)$

Response: We agree with the referee that this is unclear, for a detailed explanation of the scaling here see App. A.1, in particular the $\alpha$ scaling. The results of this scaling also appear in the article itself. This does not contradict the conclusion in the referenced article as there to the EK approximation improves when $n,\sigma^2$ are increased simultaneously. More specifically, that work relies on Taylor expanding an interaction term of the form $n\int d\mu_x e^{-\sum_{p}(f_p(x)-y(x))^2/\sigma^2}$ to first order. This becomes even more justified as $\sigma^2$ grows larger.

• In 3.1.2, "thus the assumption that higher-order Hermite polynomials are irrelevant becomes inadequate." : I do not find which assumption this line points to... Could you elaborate?

Response: To clarify this, let us first consider the infinite width limit along with the large $n,\sigma^2$ limit. In this one needs to perform a spectral decomposition of the target function on the eigenbasis of the kernel. Doing so one finds that Hermit polynomials are (eigenvectors have eigenvalues which scale as $d^{poly-power}$, (up to negligible $d^{-poly-power}$ corrections). Hence as long as $n/\sigma^2$ scale as $d$ only linear functions are learnable (meaning that our $c$ is actually fixed to $\epsilon$). As our saddle point argument shows, before the phase transition the network is a slightly modified NNGP. However, this slight modification is not enough to overcome the $d^{poly-power}$ scaling of the eigenvalues or alter the fact the $H_{1/3}$ are eigenfunction. associated with the third Hermite polynomial. Hence the ansatz we use. The same holds just after the phase transition where however, the scaling with $d$ of the eigenvalue associated with $H_{3}$ changes like that of $H_1$. Going deeper into the droplet phase, the Gaussian mixture components become more probable, $H_{1/3}$ seize to be eigenfunctions of the kernel induced by the mixture state. Specifically, components like $Erf(w_* \cdot x)$ enter the eigenfunctions which we do not track. Hence even though the target is still a combination of the two Hermite polynomials the network's output would start having an $Erf(w_* \cdot x)$ like component. Tracking these latter effects quantitatively is left for future work.

• Am I right to understand the theory quantitatively predicts the experiments only in the GFL phase, and only holds qualitatively in terms of phenomenology to explain the transition to GMFL-I and II phases?

Response: The referee is right to question this issue. For the modular algebra case, our results are exact and hold throughout the three phases. However, for the teacher-student setting, new discrepancy modes, proportional to $Erf(w_* \cdot x)$, start growing just after the transition to GMFL-I. Thus we believe our results are exactly up to GFFL-I and just after the phase transition.

Between submitting this work, and now, we conducted a more careful treatment of the scaling limit, which is now reflected in a much clearer match between theory and experiment.

• (Minor) In Fig. 2, solid lines correspond to the theory? It should be explicitly stated for better readability.

Response: We replaced Fig. 2 by Fig. 2 (c) of the revised version, There are less lines and we hope that this is now clearer.

• (Minor) It would be more compelling to complement Fig. 3 with experimental simulations, similarly to Fig.1 middle and right.

Response: We agree that this would be beneficial. However, in this setup we derived exact mathematical results and so did not see the need to further validate them. Experimental results can be added for camera ready version if necessary. Notwithstanding, we rely on the paper by Gromov in which the same qualitative behavior was shown experimentally.

We thank thank the reviewer for their honest opinion, which has helped us restructure the work in a way we believe to be a significant improvement, as well as for their insightful questions. We think that following the reviewer's advice our message is significantly clearer. We address below every comment.

Weaknesses:

• Overall, I find the presentation and clarity to be rather poor, making it hard to identify the claims and judge their soundness, and am in favour of rejection. If the authors clarify some points and promise to greatly overhaul the presentation I am happy to increase my score.

Response: We thank the reviewer for the time invested in our manuscript. We apologize for their poor reading experience. The manuscript now went through a major presentation overhaul. In particular, we delegated many of model-specific derivation details to the appendix, provided a hopefully self-contained introduction to phase transitions, summarized the relevant parts of the kernel-adaptation work, and last but not least, separated our results from the derivation and presented them in a more proof-like style. We further provide much more compelling numerical evidence.

General remarks:

• The main presentation issue is to me the excessive focus on the technical details of the derivation, from 2.4 to 3.1.2, at the expense of necessary details. For example, the discussion of the Langevin dynamics in 2.3 would strongly benefit from having the corresponding equation clearly written for definiteness, especially because several conventions exist, and also for the sake of readability.

Response: As aforementioned, we agree with the referee and made changes accordingly. In particular, the Langevin dynamics appear and are explained in further detail in the revised section 2.3.

• While I admit some technical details are essential in the understanding of the result and the flow of the paper, this is not the case of all of them (e.g. the discussion at the beginning of 3.1.1 could be shortened).

Response: We agree with the referee. These details can now be found in Appendix A.2,A.3.

• As a consequence, the assumptions on the scaling of $d,n,N$ are also dispersed through the text, making it difficult to follow which assumptions are needed in the end for the results to hold.

Response: We again completely agree with the referee, and thank them for this comment. The derivation of the scaling limit now appears in App. A.1, whereas the statement of what scaling limit we take now appears at the beginning of our results section.

We thank the referee for this relevant reference. We now cite it along our discussion of possible sample complexity aspects of our results.

Regarding the accuracy of the theory beyond the phase transition, we believe our approach captures the asymptotic behavior correctly before and exactly at the phase transition. In particular, we find that is also holds well at least as far as $S(w)$ and $b$ are concerned, just after the transition. We believe this is because correction to our discrepancy ansatz ($t(x) = b H_1(x) + c H_3(x)$) have a parametrically small effect. For instance taking $\sigma^2$ to be 0.015 smaller than our estimate $\sigma^2$ at the critical value (making $u$ larger than the critical value $u_c$), and allowing an additional $d H_5(x)$ correction to $t(x)$ we found that $b,c,d = 0.2437, -0.2859, -0.002$ whereas $b,c = 0.2433,-0.2855$ without allowing for this correction. Consequently, we do not see this as a fundamental limitation but just a technical issue where one would have to identify more properly the relevant kernel eigenfunctions which span the target just after the transition. We thank the referee for fleshing out this point and clarify this further in the text.

Questions:

• I am a bit confused about the scaling limits. What is the exact order of limits taken?...Which (if any) of these are actually adopted when deriving the action?

Response: Following the referee's comment, we clarified this point in the paper. In particular, we stress four conditions that our scaling limit should satisfy and find a two-dimensional family of scaling limits, parameterized by $\alpha,\beta$, which obey these conditions and yield the same performance as $\alpha,\beta$. The parameter $\beta$ embodies the large $N$ and large mean-field-scaling parameter limit together with the large input dimension limit. The parameter $\alpha$ embodies the large $n$ and GPR noise ($\sigma^2$) limit, making the equivalence approximation exact. Notably, apart from an overall scaling by $\sqrt{\beta}$ both $\alpha$ and $\beta$ do not change the action (see also Table 1 in the paper for the precise scaling for each model).

• In Appendix A.1 the authors describe how they adopt a mean field parameterization, yet the resulting action’s likelihood term vanishes as $N\to \infty$. Is it clear why finite-width networks would have less feature learning in the mean-field parameterization? I was under the impression that this should be constant with $N$.

Response: Indeed, this is due to our notations being slightly confusing. The non-linear term in the action has indeed a negative power of $N$, but also a positive power of $\sigma_a^2$ and two negative powers of $\sigma^4$ both of which scale as the inverse mean-field scale. Had we taken just $\chi=N^{-1}$, while keeping $d$ and $n$ fixed, this would indeed give a constant action which is independent of $N$. However, in our scaling limit, we take a more moderate increase of $\chi$ (roughly as $\sqrt{N}$) and use the increase of $n$ to ensure the non-linear coefficient remains constant. This is all now clarified in detail in App. A.1, Eq. (2) in the revised appendix.

• Could one study mean squared error so that the likelihood term was intensive in dataset size, corresponding to a loss... I am wondering if many of the factors of $n$ which appear in the computation are in fact artifacts of the comparison of the raw scale of the likelihood to the prior.

Response: The scaling we use is also related to the connection between the Langevin dynamics equilibrium distribution and the posterior distribution of Bayesian Neural Networks. Scaling the loss with a factor of $1/n$ implies in this setting that the variance of the noise scales with, $n$ i.e. $\beta^{-1} =2\sigma^{2}\sim n$ as well as the weight decay. Doing so one would find the exact same posterior and hence the same phenomenology.

You are thus right that one could avoid this training-wise-unnatural scaling of the loss, at the price of being less natural from the point of view of Bayesian Neural Networks.

• I understand that the density of $w$ undergoes transitions as the number of critical points in the action changes, requiring use of multiple saddle points when computing observable averages... Could the authors comment on these different approaches?

Response: Change in the number of saddles of $S[w]$ is a prerequisite for a phase transition, however without some large-scale parameter scaling up $S[w]$ it would induce a crossover rather than a phase transition. In this crossover, the probability weight around one saddle gradually shifts to the other saddle instead of doing so abruptly. Thus while one can avoid taking a saddle point, the accuracy of the saddle point is in fact the reason for the phase transition. In practice, we calculated integral with these quantities both numerically and with the saddle point approximation and found negligible differences in the parameter regimes used in the experiments.

• Could there be a connection between the phase transition reported here and the transition in...?

Response: This is an interesting connection with DMFT and might imply potential dynamical aspects of our transition. Indeed, we can similarly view the mean-field scaling parameter ($\chi \sim \gamma_0$) as driving the phase transition. However, for this to be a phase transition, all other parameters must scale up. Both in the Fig. 5c of this work and in the text we did not find an emphasis on phase transitions but rather on transitions/crossovers. Perhaps the referee can clarify if they think an actual sharp transition is expected. Is this issue with convergence discussed there? Is it reflected in bifurcation points in the ODE for the alignment in Saxe's setting? We'd be thankful for some further clarifications so that we may cite this properly.

We thank the reviewer for their precise summary, for pointing out the strengths and weaknesses of our work, and for the insightful and helpful comments we received from their review. We also appreciate the reviewer taking the time to ask relevant and thoughtful questions. The comments have helped us clarify our message and improve our work. We address below each and every comment.

Weaknesses:

• While the phenomena described and the resulting theoretical picture of 3 phases is quite impressive, I am not sure that this transition constitutes grokking... It also relies on some prior work from Seroussi et al 2022 and a short summary of this approach in the Appendix could be helpful. Lastly, some of the exposition is a bit challenging to follow (see questions below). That said, if these issues/questions are addressed I would be happy to raise my score.}

Response: We thank the referee for their insightful comments and are pleased to hear they find our work touches on fundamental issues. While we understand the viewpoint that Grokking is mainly about time, we wish to point out that (1) It is often accompanied by Grokking effects as a function of data-set size (See new Refs. in the introduction) (2) In general, behavior of DNNs as a function of time versus equilibrated neural networks as a function of dataset size is often similar, as reflected for instance in the use of One-Step SGD (You 2014 :"Exploring one pass learning for deep neural network training with averaged stochastic gradient descent"). We thus hope to keep this more timely framing, although we agree that our work addresses even more fundamental questions in learning representations.

The revised manuscript now includes a more detailed description of the relevant parts of the adaptive kernel approach we are using (Sec. 2.4.1.). If the referee finds it useful, we can further expand on this in the appendix.

We're happy to hear that the referee found our response useful.

While we need to consider the referee's example further, we can say some related things.

First, indeed the failure of a Taylor series to converge does not immediately mean a phase transition. For instance in quantum field theory, one often performs a perturbation theory keeping only 1-particle-irreducible diagrams leading to a mass renormalization (shifting of poles) in the Green's function $G(k,w)$. The resulting geometric series coming from perturbation theory is of the form $G=G_0 + G_0 V G_0 + G_0 V G_0 V G_0$ (where V is the self-energy) does not converge when $k,w$ are close to being on-shell and $V G_0$ is large. However, this does not imply a phase transition but rather just a limitation in this representation of the perturb Green's function. In practice one analytically continues the series. One way of doing so would be to re-sum this series leading to $1/ (G_0^{-1} - V)$ and treat this as the function. Another would be to find the Dyson equation obeyed by $G$ and solve it to again obtain this $1/(G_0^{-1} - V)$ form which has better analytical-continuation properties.

Second, It would be fascinating to see these Landau-like phase transitions showing up in the DMFT and obtain their dynamical signatures. We wonder though if those can happen for deep linear networks. At least in equilibrium, we found that having a non-linear activation function is essential to get out of the GFL phase.

We again thank the referee for his/her questions.

We thank the reviewer for recognizing the novelty in our work and for the insightful and helpful comments we received from their review, as well as pointing out relevant further directions. These have considerably helped us to clarify the message. We also appreciate the attention to detail in the "Nits" section. We address below each question as well as every comment.

Weaknesses

• The connection to Grokking, as observed in prior work across various tasks and modalities, seems almost secondary in this paper. This work focuses on some toy models, solves them, and refers very little to the general phenomenon of delayed generalization.

Response: Our work is indeed about learning representations, a topic we found fitting for ICLR. We understand the referee's viewpoint that framing this as a Grokking phenomenon only ignores the dynamic aspects of Grokking. Still, as the revised introduction now states and cites, Grokking, in the delayed generalization sense, is accompanied by Grokking as a function of dataset sizes. More generally, similar behavior is often observed between fully trained networks as a function of dataset size and networks trained on an infinite amount of data for a finite time (See for instance You 2014 :"Exploring one pass learning for deep neural network training with averaged stochastic gradient descent"). For instance, by taking very low learning rates and gradually increasing the data-set size, this correspondence between dataset size and time becomes linear and exact.

With regard to the referee's comment that we just solved two toy models, we note that we put out a mapping between abrupt feature learning effects and the theory of phase transitions. Our starting point action (log posterior distribution) is completely dataset and targets generic, and thus potentially captures a much bigger family of phase transitions in two-layer non-linear networks. We can solve it analytically for common toy models, like the ones we choose. Notably, one of which is a very natural and widely used toy model for Grokking (modular algebra learned by an FCN is also presented in other works such as Gromov 2023: "Grokking modular arithmetic" and in Nanda 2023: "Progress measures for grokking via mechanistic interpretability" ) hence hopefully elevating doubts that we somehow tailored the toy model to the theory.

Last we note that to the best of our knowledge and according to some of the other referee reports, solving a non-linear neural network with two trainable layers in the feature learning regime, is a considerable push to the theoretical state of the art in the field.

• The general feeling I have from this paper is that the approach somewhat obscures the contributions and implications... But overall, I like the approach and would want to see more work in this direction, so I give a score of 6 (weak accept).

Response: We sympathize with the referee's wish to see more practical implications. The current manuscript has been revised to stress our claims and the conditions under which they apply more clearly. It now places more places more emphasis on the generality of our approach by structuring the results sections so that it is evident that the same procedure is preformed in both cases and can be preformed similarly in many other similar cases. Furthermore, we extended the background sections thus hopefully further helping readers understand the potential generality of our approach.

Turning to the second point regarding experiments, establishing a correspondence between feature learning in latent representation and water-vapor transitions, providing analytical pen & paper predictions on two canonical non-linear models, and verifying this experimentally is in our minds a sizable theoretical contribution. Moreover, as we now clarified in the text, our solution method (initial action and set of approximations we apply) is completely general and in particular shared between both models (see the derivation overview section 2.4). Furthermore, the droplet pictures, wherein some neurons are largely target agnostic whereas some carry important representations, can be used to rationalize why pruning works. Following this, we are currently conducting pruning experiments, showing some promising preliminary results.

Weaknesses- continued

• In my current understanding, the whole "adaptive kernel approach" ... The adaptive kernel theory in this paper cannot cover this case at all.

  Response: The precise outcome of training a neural network surely depends on all training details. However, it is rarely the case in practice where SGD at one batch size fails whereas at a large batch, it does spectacularly well. Underlying the celebrated NTK for example, is a seemingly impractical full-batch SGD at a vanishing learning rate. We assume full batch SGD at a vanishing learning rate and, heuristically, replace the SGD colored/correlated noise with white noise. This allows us to make contact with Bayesian Inference formalism. We comment that the latter shares much resemblance with SGD as argued in various places in the literature (for instance in Ard Louis, "is SGD a Bayesian sampler well almost"). We additionally note that the adaptive kernel approach can also be viewed from the Bayesian perspective (see Li & Somplolisky 2021, Welling, M. & Teh, Y. W 2011 and more), where Langevin noise can be viewed as noise on the target and in many cases it is reasonable to assume that it is small and uniform over samples.

  Separating proofs from arguments and exact agreement with quantitative agreement, we believe, based on experience, that the adaptive kernel approach is a good heuristic for realistic networks despite its assumptions. Indeed, our work is in qualitative agreement with various aspects of actual Grokking happening in SGD-trained networks (Gromov 2023), in particular, the types of filters that are learned and the fact that only a fraction "droplets" of neurons have the desired Fourier-transform filters.

Questions

• I would like to ask the authors to give more introduction of the technical tools they are using and rephrase their results in terms of theorems and lemmas

  Response: As aforementioned, the revised manuscript now contains a self-contained introduction to phase transitions and kernel adaptations approach (section 2.4), as well as more precise claims regarding our main results. We further made efforts to better state the conditions for our theory to hold and its implications.

We thank the reviewer for their precise summary, for pointing out the strengths and weaknesses of our work as well as the places in which our background information was lacking. These have considerably helped us to both clarify the message and appropriately introduce the relevant concepts. We address below each and every comment.

Weaknesses

• This paper is poorly written ... I don't believe the current version of the paper is ready for ML audience to read.

  Response: We thank the referee for the time invested in our work. We apologize for their poor reading experience here. We agree that while tools from statistical physics are making their way into the field, they are not at the stage where they can be taken for granted. Furthermore, our results in the previous version were too tangled up with derivation details.

  The revised manuscript thus went through a major presentation overhaul, and we believe it now answers all points brought up by the referee. In particular, it now has a more self-contained introduction to kernel-adaptation results and to the theory of phase transitions. Furthermore, our main claims and the limits under which they apply are now clearly stated.

  In essence, the idea behind kernel adaptation methods, is to extend the known result about Neural network at infinite width to the finite width regime. Previous results in this direction have shown that the output of the network at the end of training behaves like a Gaussian process. The kernel describing this Gaussian process as opposed to the Gaussian process at the infinite width limit do depend on the training data, both features and labels. Another useful aspect of this approach is that layers' probability distributions decouple, allowing one to formally treat representation learning as layer effect.

  The reason why we focus on the action as our main result is that the action is the log of the posterior distribution up to a constant factor. Hence it is just a formality which helps avoid tracking normalization factors and terms in the exponent. We extend more on this idea and introduce the concept in section 2.4 in the main text.

• If this paper focuses on theory, then it is better to organize the claims into theorems and lemmas...

  Response: We agree with the referee that the conditions under which we believe our theory works should be stated much more clearly. Accordingly, we revised the scaling limit appendix (App. A.1) and clarified this limit in the main text. We also now sate all our claims explicitly in the main text in a way we hope will make it easier for the referee to follow.

  Unfortunately, we do not aspire to prove all the claims in this work. Rather, we provide a physics-style treatment of this system using the same level of rigorousness and common practices found in most theoretical physics publications. For instance, if we find a small parameter, we expand in it and assume that such expansion convergences quickly. We however do not bound the Taylor residue, which is often very complicated to do. To elevate various doubts that accompany such a paradigm, we always perform numerical experiments of our approximations with numerics. In particular, the revised manuscript shows a very good match with experiments.

  Notwithstanding, some of our claims are exact. For instance, with regard to the modular algebra section, our setup there is sufficient symmetry in the system such that our results are indeed mathematically exact.

  We note that being non-rigors is also the strength of the Physists' approach, as it allows one to make progress and introduce claims motivated by physicist intuition without focusing on the technical details. These ideas can sometimes be beyond the current mathematical knowledge. Physists' claims can be then introduced for mathematicians as open problems and conjectures that are interesting to prove mathematically and sometimes take years to be proven. Such a process has happened many times in science in the past, spin glass theory is a very good example of that.