Grokking as a First Order Phase Transition in Two Layer Networks
Authors: Noa Rubin
Inbar Seroussi
Zohar Ringel
Venue: ICLR 2024 poster (2024)
Keywords: Grokking
deep neural networks
Gaussian Process
phase transitions
Primary area: general machine learning (i.e., none of the above)
TL;DR: Analytical predictions of feature learning and Grokking properties which demonstrate a mapping between Grokking and the theory of phase transitions

Abstract: A key property of deep neural networks (DNNs) is their ability to learn new features during training. This intriguing aspect of deep learning stands out most clearly in recently reported Grokking phenomena. While mainly reflected as a sudden increase in test accuracy, Grokking is also believed to be a beyond lazy-learning/Gaussian Process (GP) phenomenon involving feature learning. Here we apply a recent development in the theory of feature learning, the adaptive kernel approach, to two teacher-student models with cubic-polynomial and modular addition teachers. We provide analytical predictions on feature learning and Grokking properties of these models and demonstrate a mapping between Grokking and the theory of phase transitions. We show that after Grokking, the state of the DNN is analogous to the mixed phase following a first-order phase transition. In this mixed phase, the DNN generates useful internal representations of the teacher that are sharply distinct from those before the transition.

Published as a conference paper at ICLR 2024
GROKKING AS A FIRST ORDER PHASE TRANSITION IN TWo LAYER NETWORKS
Noa Rubin"
Inbar Seroussit
Zohar Ringel
ABSTRACT
A key property of deep neural networks (DNNs) is their ability to learn new fea- tures during training. This intriguing aspect of deep learning stands out most clearly in recently reported Grokking phenomena. While mainly reflected as a sudden increase in test accuracy, Grokking is also believed to be beyond lazy- learning/Gaussian Process (GP) phenomenon involving feature learning. Here we apply recent development in the theory of feature learning, the adaptive kernel approach, to two teacher-student models with cubic-polnnomial and modular addi- tion teachers. We provide analytical predictions on feature learning and Grokking properties of these models and demonstrate mapping between Grokking and the theory of phase transitions. We show that after Grokking, the state of the DNN is analogous to the mixed phase following first-order phase transition. In this mixed phase, the DNN generates useful internal representations of the teacher that are sharply distinct from those before the transition.
INTRODUCTION
Feature learning is a process wherein useful representations are inferred from the data rather than being engineered. The success of deep learning often attributed this process. This is reflected, in part, by the performance gap between actual deep neural networks and their infinite-wine width Gaussian Process (GP) counterparts (Williams) 1996. Novak et al. 20188) Neal 1996). It is also key transfer learning applications (Weiss et al.) (016)) and interpretability (Zeiler & Fergus) [014] Chakraborty al. 2017). Yet despite its importance, there is no consensus on how measure or let alone classify Teature learning effects.
Several recent results Li & Sompolinsky (202) Ariosto et al. (2022) began shedding light on this matter. One line of work (adaptive kernel approaches) treats the covariance matrices of activation within each layer (kernels) the key quantities undergoing feature learning. Feature learning would manifest as deviation of these kernels from those of a random network and their adaptation to the task hand. While providing a quantification of feature learning quite generic settings, the equa- tions governing these latent kernels are quite involved and may host variety of learning phenomena. One such phenomenon (Seroussi et al., 2023), capable of providing a strong performance boost, is that of Gaussian Feature Learning (GFL): A gradual process in which the covariance matrices of neuron pre- reeaactivations change during training so as to increase their fluctuations along label/target relevant directions. Remarkably, despite this smooth adaptation, the pre-activations' fluctuations, across width and training seeds, remain Gaussian. At the same time, the latent kernel itself develops notable spikes in the target direction, indicating feature learning.
Another phenomenon often associated with feature learning is Grokking. This abrupt phenomenon, first observed in large language models running on simple mathematical tasks, involves fast changes to the test accuracy following longer period of constant and poor performance (Power et al. 2022) Though usually described
time-dependent phenomenon, Grokking as function of other param- eters. it also occurs as a function of other parameters, such as sample size Power et al. (2022): Gro-
mov
(023)
Liu al. 2022a). More broadly, DNNs behaviour as a function of time and dataset size
HHebrew University, Racah Institute of Physics. Jerusalem. 9190401, Israel
1 Department of Applied Mathematics, School Mathematical Sciences, Tel Aviv University, Tel Aviv 69978, Israel
1 Hebrew University, Racah Institute of Physics, Jerusalem, 9190401, Isracl

---

Published as a conference paper at ICLR 2024
is often similar as reflected for instance the use of One Pass SGD You al. (2014). Several authors provided quantitative explanations in the context of specific toy models wherein one can handcraft or reverse engineer the solution obtained by the network (Gromov) [023] Nanda et al., 2023) in
suitably tailored perceptron models (Liu et al. 2022a Zunkovië & Ilievski 2022) where, however, representation learning is tricky to define. Given the aforementioned adaptive kernel approaches to
deep learning, as well as the universality of Grokking across different DNNs and hyperparameters, it is natural to look for a more unifying picture of Grokking using a formalism that applies to generic deep networks.
In this work, we study Grokking as an equilibrium (or Bayesian) phenomenon driven by sample size, noise, network width. Utilizing the aforementioned theoretical advancements, show that Grokking in large-scale models can be classified and predicted through the mean field theory of phase transitions in physics (Landau & Lifshitz 2013). Studying two different models, teacher- student with cubic teacher and modular algebra arithmetic, we show the internal state of the DNN before Grokking
well described by GFL. In contrast, during Grokking, it is analogous to the mixed phase in the theory of first-order phase transitions, and the statistics of pre-activations are described by mixture of Gaussian (GMFL). In this GMFL. state, the latent kernels associated with the DNNs develop entirely new features that alter their sample complexity compared to standard infinite-width GP limits. After Grokking the weights are all specialized to the teacher. Besides pro- viding framework to classify feature learning effects, our approach provides analytically tractable and quantitative accurate predictions for the above two models.
Our main results are as follows:
We establish a concrete mapping between the theory of first-order phase transitions, internal representations of DNNs, and Grokking for two non- non-linear DNN models each having two tunable layers.
We identify three phases related to Grokking, one which is smoothly connected to the GP limit and two distinct phases involving different GP mixtures.
For both our models, we simplify the task of learning high-dimensional representations solv- ing a non-linear equation involving either two (cubic teacher) or one (modular addition teacher) variables. Moreover, for the latter, we determine the location of the phase transition analytically.
We flesh Grokking- mechanism that can reduce the sample complexity compared to the GP limit.
Prior works. Phase transitions are ubiquitous in learning theory (e-g. Refs. Garddnr & Derrida 1988]; [Seung et al. (1992); Györgyi (1990)), often in the context of replica-symmetry breaking. Connections between Grokking and phase transition were suggested by Nanda et al. (2023) but as far as analytic predictions go, prior work mainly focused on one trainable layer Zunkovi & Ilievski (202); ; Arnaboldi et al. 2023) some suggesting those as effective theory of representation learn- ing Liu et al. (2022a). This can be further investigated by analyzing the loss landscape Liu et al.) 2022b). Varma et al. (2023) suggest that the generalizing solution learned by the algorithm is more efficient but slower to learn than the memorizing one, using this interpretation they define regimes of semi-grokking, and ungrokking. The formalism of Refs. Arnaboldi al. (2023); [aad & Solla (1995) can potentially be extended to online Grokking with two trainable layers, but would require reducing the large matrices involved. Phase transitions in representation learning have been studied in the context of bifurcation points in the information bottleneck approach (e.g. Tishby & Za- slavsky (2015), nonetheless, the connection deep learning remains qualitative. To the best of our knowledge, we provide a novel frrst-principals connection between grokking and phase transition in representation learning.
MODELS
2.1 NON-LINEAR TEACHER MODEL.
Our first setting consists of student Erf-network learning a single index non-linear teacher. The student trained on a training set of size n, 'D (xp, (In)) with MSE loss. In the following,
bold symbol represents a vector. The input vector is 2 Rdd with iid Gaussian entries of variance
2

---

Published as a conference paper at ICLR 2024
1. The target function y. is scalar linear function of with an orthogonal non-linear correction. Specifically, y is given by-
y(x)= = w w. te (ww" xx)² I² "²= a
H1 Hiww =))
Ha(w" m))
where Hi, H3 the first two odd Hermite polynomials, and w² E Rd the teacher weights. For simplicity we take here the norm of the teacher weights to be 1, but this has no qualitative effect on the theory as long
require | (1)). We consider fully connected non-linear student network with one hidden layer of width N given by
N
f (x) = a (w x)
where w, Rd for [] are the students weights. Evidence that this model Groks can be found in App. D).
2.2 GROKKING MODULAR ALGEBRA
Here we consider the setting of Ref. Gromov 2023 where the learning task is addition modulo P where prime. The network is trained on the following data
D [m y (Xrm) [m, Zp}
where Zam R2P is vector such that it is zero in all its coordinates except in the coordinates and P + m where is 1 (a "two-hot vector"). The target function y E IR"" is given by
(2(mmm) = (x mod I/P.
(4)
where 5 is the Kronecker delta and mod P denotes the modulo operation which returns the remainder from the division by P. For the student model, we consider two-layer deep neural network with square activation function, given by
N
fp (Xmm)
aps aps(W; (w: Fmnn )2
(5)
where we E 2P for E N] are the students weights. For brevity, we denote Up (Fmn) = Vhim and fp (Ium) fm
2.3 TRAINING THE MODELS
In both cases, we consider networks that are trained with MSE loss to equilibrium using Langevin dynamics via algorithms such Durmus & Moulines (2017); Neal al. (201I). The continuum- time dynamics the parameters are thus
@(1) = -o ( I.
2aE(1)
(6)
where (t)) is the vector of network parameters in time y is the strength of the weight decay, is the loss function, the noise is given by (E,(t), (t') 6. (t t') and is the magnitude of the noise. We set the weight decay of the output layer so that with no data af az both average to o /N under the equilibrium ensemble of fully trained networks. The input layer weights are required to have a covariance of of. d in the teacher-student model with o² O(1) and covariance of 1 in the modular algebra model. Note that the covariance of the hidden layer is given by ///y. The posterior induced by the above training protocol coincides with that of Bayesian inference with Gaussian prior on the weights defined by the above covariance and measurement noise σ² Naveh al. (202))
3

---

Published as a conference paper at ICLR 2024
2.4
DERIVATION OVERVIEW
2.4.1 BRIEF INTRODUCTION TO MEAN FIELD THEORY PHASE TRANSITIONS
Phase transitions Landau & Lifshit 2013 Tong 2011). such as the water-vapour transition, ubiquitous in physics. They are marked by a singular behavior of some average observables as a function of a control parameter (say average density are typically smooth, phase transitions are inherently large-scale, or thermodynamic, phenomena loosely analogous how the sum of many continuous functions may lead to non-continuous one.
function of volume). As the laws of physics
In a typical setting, the probability p(r), of finding the system at state x can be approximately marginalized track a single random variable called an order parameter (Φ). As the latter is typ- ically a sum of many variables (i.c. macroscopic), it tends to concentrate yielding a probability of the type log(p($)) x (d) where d is a macroscopic scale (e.g. number of particles) and S(Φ) some well behaved function that does not scale with d. Given this structure, the statistics of Φ can be analyzed using saddle point methods, specifically by Taylor expanding to second order around minima of 5.
Phase transitions occur when two or more global minima of ((2) appear. First-order phase transi- tions occur when these minima are distinct before the phase transition and only cross in Š value at the transition. Notably, the effect of such crossing is drastic since, at large d, the observed behaviour (c.g. the average Q) would undergo a discontinuous change. Depending on the setup, such sharp change may be inconsistent it would immediately change the constraints felt by the system. For instance, in the water-vapour transition, as one lowers volume the pressure on the vapour mounts. At some point, this makes a high-density minima of O favourable as the low-density minima, signifying the appearance of water droplets. However, turning all vapor to droplets would create a drop in pressure making it unfavourable form droplets. Instead, what is then observed a mixture phase where as a function of the control parameter, a fraction of droplets forms so as to maintain two exactly degenerate minima of Š. Lowering the volume further, a point is reached where all the vapour has turned into droplets and one is in the liquid phase.
In our analysis below, & will be a property of the weights in each neuron of the input layer, say their overlap with some given teacher weights. A high input dimension will be analogous to the large-scale limit, and the density loosely corresponds to the discrepancy in predictions. The phase transitions are marked by a new minima of (()) which captures some feature of the teacher network useful in reducing the discrepancy. What we refer to as droplets corresponds to some input neurons attaining & values corresponding to the teacher feature while others fluctuate around the teacher agnostic minima. However the spatial notion associated with droplets not relevant in this case as in the mean field theory of phase transitions.
2.4.2 ADAPTIVE KERNEL APPROACH AND ITS EXTENSION GAUSSIAN MIXTURES
Our main focus here is the posterior distribution of weights in the input layer (v:))) and posterior averaged predictions of the network (f))). Such posteriors are generally intractable for deep and/or non-linear networks. Hence, we turn to the approximation of Ref. [Seroussi et al. (2023) where mean-field decoupling between the read-out layer and the input layer is performed. This exact in the limit of N + and vanishing of (i.e. mean-field field scaling). Following this, the posterior decouples to a product of two probabilities, a Gaussian probability for the read-out layer outputs and a generally non-Gaussian probability for the input layer weights. These two probabilities coupled via two non-fluctuating quantities. The average kernel induced by the input layer and the discrepancy in predictions. As shown [Seroussi et al. (2023). specifically for a two-layer FCN, the resulting probability further decouples into iid probabilities over each neuron p(w:). Below, we thus omit the neuron index i. The action (- (-log(p(un) wx;)) up to constant normalization factors) for each neuron's 's weights is then given by the following form
02
Siw]
((( o(w In)6(w xw),
20 W
2N
P.v
:00= QQ
4

---

Published as a conference paper at ICLR 2024
where Φ is the activation function and t is the discrepancy between the averaged network output and the target given by
8 (xxμ) (y (I (x,))//²²
(8)
Notably & is not given but determined by solving the following two mean-field self-consistency equations
f = Q @=Q[Q+02m1~
(9)
where the kernel Q is defined via
Qu 0 (eb (w xxp) Φ (w Bu))s|wo
(10)
and (1o)) denotes averaging over w with the probability implied by S[w] In this work we use the equivalent kernel (EK) approximation, allowing the sums appearing in eqs. 7 to be replaced by integrals. This approximation washes out the generalization phenomena associated with Grokking, while capturing underlying feature learning mechanisms. As demonstrated in App. D the feature learning effects hold also for finite datasets, which a generalization gap can be observed. Theo- retical corrections due to finite datasets can be made shown in Cohen al. (2021 Seroussi al. 2023
Naveh & Ringel (2021), here we focus on the EK limit for simplicity.
Even if 7 is given, the remaining action is still non-linear. Ref. Seroussi et al. (0233 proceed by performing variational Gaussian approximation on that action. Here we extend this approximation into a certain variational Gaussian mixture approximation. Specifically, we show that as one scales up d, N. n in an appropriate manner, and following some decoupling arguments between, S(w) has the form d$| ((r) where d >>> and S S[F(w)) has O(1) coefficients and ((w) is some linear combination of the weights. This allows us to treat the integration underlying (...) within a saddle-poing point approximation. Notably when more than one Global saddle appears the saddle-point treatment corresponds to p(w) (w)) having a Gaussian mixture form.
The three phases of learning described in the introduction correspond to the following behavior the saddles (see also illustration in Fig. 1. At first, single saddle cen- tered around Φ (w) = 0 exists and weights fluctuate in a Gaussian manner around these minima, as in GFL. In the next phase, the distribution is comprised of a weighted average of this zero-saddle and other [+()) dles. This marks new learning ability of the network and hence the beginning of Grokking. We name this mix- ture phase the first Gaussian Mixture Feature Learning phase (GMFL-I). This phase corresponds to the mixed ("droplets") phase. In this phase, picking random ron, there is a finite chance it is in the GFL phase and hence fluctuates around the target agnostic minima w O. Similarly, there is a finite chance it fluctuates around one of the non-trivial saddles. At some later point, after more data has been presented, only the saddles with w > 0 dominate the average (GMFL-II). The phase phenomenology shared by both models, however, the details of the new saddles that appear, as well as the decoupling scheme between the different components of w differ between both models. We next turn our attention these details.
Wiah-N
GFL
GMFL-I
GMFL-II
0 sad-
Figure Schematic phase diagram of learning phases. The inset plots in the different phases correspond to the nor- malized negative log posterior (5). transition between the different phases can be achieved by varying either n or N.
RESULTS
3.1 NON-LINEAR TEACHER STUDENT MODEL
Scaling setup and effective interaction. Consider the following two scaling variables (a, 8) which would soon take to infinity together and consider scaling up the microscopic parameters in the following manner
N + BN d V Bd 22 + o2/ VB 0²
a 0 n + an
(11)
B
5

---

Published as a conference paper at ICLR 2024
where comment that can be continuum approximation allowing replace data summation with integrals over the data measure and / is a combination of mean- mean-field-tye field-type scaling Mei et al. (2018) together with thermodynamic/ssddlee point limit. Notably the following combi- n'o" otaN. which we refer to as the effective interaction, is invariant
nation (u) of hyper-parameter under both a and 3.
Claim 1. Two relevant discrepancy modes before the transition. For 8 + oo and vB/a + 0, and u S uc(c) (Rs 30.2 for c -0.3) the discrepancy takes the following form
)²t()) bH1 (a) + cH3 (x)
(12)
where b, e E are some O(1) constant coefficients. We comment that wh² is proportional to the emergent scale of Ref. (Seroussi et al. (2023). For further detail see App. A.2
Claim II. One-dimensional posterior weight distribution. In the same scaling limit, u S uc the negative log probability (action in physics terminology) of weights along the teacher direction, decouples from the rest of the w modes, and takes the following form
Slw [hapww]] =d "wwww")
2n²o²
(w w²)²
2c (w =²)²
b
2
o + (w w=)3))
1 1+ 2 02 + (w w0)²
(13)
Notably, this expression reduces the high-dimensional network posterior into sealar probability involving only the relevant order parameter (6 w w²)) We further note that all the expressions in the brackets are invariant under and B. Heuristically, this will also be the action for small Au after the phase transition since the resulting corrections to the discrepancy have parametrically small effect. For further detail see App. A.2
Claim III. Exactness of Gaussian Mixture Approximation, In the same scaling limit, the proba- bility described by the above action exactly a mixture of Gaussians each centred around global minimum of S.
Claim IV. First-order phase transition. For u < ue the only saddle is that + Exactly at = uh three saddles appear two of which are roughly at Φ- +|w_|2 1 +1. For some finite interval u E [Me, Ue Au], these saddles stay degenerate in S value.
3.1.1
TEACHER-STUDRRR EXPERIMENTAL. RESULTS
(x)
0²=0.3
g² 0.1
0.01
0.00
8
2+38
4-330
0.01
6
0.02
0.03
0.04
2
0.05
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
0.5
1.0
0.06
0.10
0.15
0.20
0²
0.25
0.30
10' 12.
W 10
Figure 2: GFL to GMFL-I Theory and Experiment in the teacher-student model. Panel (a) and (b) show the negative log posterior before and after the phase transition induced by varying the noise o². A good match with the experimental values (crosses) and theory (solid lines) for rarer and rarer events obtained as we scale up model according to Table 1. (colour coding). Turning to network outputs, panel (c) shows the expected phase transition in learning the cubic teacher component and the inset shows the discrepancy in the linear teacher component. As our analytics holds before and at the phase transitions, discrepancies in f H3 close to the transition an expected finite-size effect made pronounced by its low absolute value.
To validate our theoretical approach, we trained an ensemble of 200 DNNs for different a = B values using our Langevin dynamics sufficiently low learning rates and for sufficiently long

---

Published as a conference paper at ICLR 2024
1) were: N 700, 3000, d = 150, of = 0.5, o = 8/N and we took c = -0.3. Our training ensemble consisted of both different initialization seeds and different data-draw seeds. These, together with the neuron indices associated N, provided 200.N draws from w, uo used for estimated p(w). The discrepancy was estimated using dot product of the outputs with H (a))(a) sampled over 3000 test points and then averaged over the 50 seeds.
(i.e. at time ensure equilibration, Our initial N, d, o², og
Our experimental results are given in Fig. 2 Panel (c) shows how the linear (inset) and cubic target components learned, as a function of 0²_ Notably, reducing o² is similar to increasing n, hence one expects more feature learning
lower o². As we increase a = (see color coding) sharp phase transition develops around o² = 0.18 where the cubic component begins its approach to the teacher's value (c). Panels (a) and (h) track - log(p(w w.)) at two points before and after the phase transition, respectively. As a 3 increases, the theory predicts a finite probability for ((w)= = w w. 1. This means that picking a random neuron has a finite chance of fluctuating around the teacher-aware minima. Such neurons are what we refer to our droplets. All graphs show a good agreement with theory as one scales up a = =3 before the transition. The remaining discrepancies attributed to finite size (i.e. finite a; () effects which, expected, become more noticeable near the transition.
The above experiment also points to some potentially powerful complexity aspects of feature learn- ing. Notably, an FCN NNGP kernel induces a uniform prior over cubic polynomials (i.e. all 3 O(d3) of those, it requires
hyper-spherical harmonics). As there to learn such target components in our setting (see also Ref. [ohen et al. (202)). ). Here this occurs two orders of magnitudes earlier (n 3000). This occurs because the complex prior induced by finite DNN learns the features from the readily accessible linear components of the target and applies them to the non-lineas ones. Proving that such "assisted learning" of complex features changes the sample complexity compared to NNGPs requires further work. In App. A.1 we provide an analytical argument support of this.
0.5e datapoints
Table 1: Scaling laws
Model
Width
Input-dimension
Data size
Noise strength
Weight decay
Polynomial regression
N + 3N
d d /Bd
+
3 a² 0 0218
g0t
a + / /B3
Modular Theory
N 82N
P BBP
p2 + 3p
a - 0
3.2 MODULAR ALGEBBA THEORY
Scaling setup and effective interaction.
Similar to the polynomial regression problem, we consider a scaling variable (8) which we later take to infinity together, and consider scaling up the microscopic parameters. The precise scaling given in Table I
N N P BBP d2 -0 02-0022000000 0213
(14)
Note that, here, we do not need additional scale for the continuum limit of the dataset, since the continuum limit is taken by considering all combinations of data points. As before, 3 is combina- tion of mean- mean-field scaling together with thermodynamic/sadddeee point limit. Notably, the following
203 p2
combination (u) of hyperparameter u = Not which we refer to as the effective interaction, is
invariant under B.
Problem Symmetries
The following symmetries of S(w) (which is also function of t) help decouple the posterior probability distribution into its Fourier modes and simplify the problem considerably:
I. Taking [n, m] + [[n + q) mod P. (m + 4') mod PI. and fp f(ptate') mod P with o' Zp
II. Taking (n, m] (gn mod P, gm mod P| and fp - Sap mod r for q E Zp but different than zero
Claim I. Single discrepancy mode. Several outcomes of these symmetries are shown in App- B.1. First, find that the adaptive GP kernel Q (given explicitly in Eq. 21) is diagonal in the

---

Published as a conference paper at ICLR 2024
basis Os., (Zi,,) P-1 222r((kknk=mmmm where k. k {0, 1-..,,-.1111 Considering eigenvalues, the second symmetry implies that ON would be degenerate with eck.ckr. ckr. For prime, P this implies, in particular, that all ΦN.E eigenvectors with & > 0 have the same eigenvalue (A). Notably, the target itself is spanned by this degenerate subspace specifically
Pim P-1
/2mkp/ P 22ak((+m)/
12kkp/ &R (Fn..)
(15)
Arm
As result, one finds that the target is always an eigenvector of the kernel and 2 = ayhm where E R. Thus there is only mode of the discrepancy which is aligned with the target.
Claim II. Decoupled two-dimensional posterior weight distribution for each Fourier mode. We decouple all the different fluctuating modes by utilizing again the symmetries of the problem and making& a judicial choice of the non-linear weight decay term (T). To this end, we define the following Fourier transformed weight variables (w): Uk): wn kk=
P-1 upe 2ntkn
Wm
Lk= ULE
which when placed into action yields
P-1
20 p2
P- Σ
S/ri] =
w w
Σ
UEUL
w AUEUK
+ [[w]
(16)
2
Not
k=
k=1
(see App. B.2 where apart from the non non-linear weight-decay term, all different is modes have been decoupled. For simplicity, we next choose, l'[w] = Σn P (ww.w.)² (vxv-x))) Here there U-g-kk but from the symmetry of the
are technically two order parameters- + w.-W-K action, we obtain that the saddles occur only at points where + = V. We comment that the analysis of a more natural weight decay terms such as Σ,, w + w using certain GP mixture ansatz of p(w:) as approximation, we obtained similar qualitative results.
Claim III. Exactness of Gaussian Mixture Approximation. factor of P in front of the action, the above non-linear action, per k-mode, can be analyzed using standard saddle point treatment. Namely, treating Q as a function of a, and with it A, can be evaluated through mndde--piint approximate
Following the presence of a large
the probability associated with this action. In the limit of B -> oC, we obtain that this approximation is exact. Following, this a can be calculated using the x² GPR expression x and in this limit the value of À can be computed exactly. Demanding this latter value of a matches the one in the action results in an equation for a.
Claim IV. First-order phase transition. Since the quadratic term constant in the scaled action, as long as no other saddles become degenerate (in action value) with the + = Tr = O saddle, the saddle- point treatment truncates the action at this first term. By increasing u the quartic term becomes more negative, and hence first-order mymmetr--breakinn transition must occur at some critical value of u. Past this point, a begins to diminish. If it will diminish too rapidly, the feedback on the action will be such that it is no longer preferential for to diminish, and thus probability distribution will have two degenerate minima at a zero and non zero value of & representing the GFML-I phase. Further increasing will break the degeneracy resulting in the global minimum of the log posterior distribution being non trivial. Notably a measures the test-RMSE here, thus the fact that remains constant and suddenly begins to diminish can also be understood as Grokking.
3.2.1 MODULAR ALGEBRA NUMERICAL SIMULATIONS
Solving the implied equation for a numerically yields the full phase diagram here (see App. B.3 for technical details of solution) supporting the assumptions in Claim IV. Fig. (3) plots the negative- log-pro oo--robabbiliyy of weights for an arbitrary taking + V for simplicity, since at the global minima this is anyways true. Here we increase by decreasing a² and find the action by solving the equation for a(a") numerically. The & + 0 saddles are exponentially suppressed in P. yet nonetheless become more probable as o² decreases. At around, o² x 0.227 they come within ( of the saddle at S& ② = 0, v = 0) (O(1/P) in the plot, given the scaled y-axis). This marks the beginning of the mixed phase (GMFL-D) wherein all action minima contribute in non-negligible manner. Further, in this phase, decreasing o² does not change the height of the & * 0 minima (see inset) in any appreciable manner. Had zoomed in further, we would see very minor change these saddle's
8

---

Published as a conference paper at ICLR 2024
height throughout the mixed phase, as they go from being O(1) above the + = saddle to O(1) below that saddle at o² z 0.175, this is shown in the inset graph in Fig. (3. This latter point marks the beginning of the GMFL-II phase, where it is the contribution of the minimum at Φb 0 which becomes exponentially suppressed in P. Notably a, which measures here the test-RSME, goes from -1 at the beginning of GMFL-I 0.7 its end. Over this small interval of o² we observe 30% reduction in the magnitude of the discrepancy which can be thought of as manifestation of Grokking.
Our analytical results are consistent with the experiments carried out in Ref. Gromov 2023). In- deed, as we enter GMFL-I, weights sampled near the Φ V * saddle, correspond to the cosine expressions for Wx of that work. As our formalism marginalizes over readout layer weights, the phase constraint suggested in their Eq. (12) becomes irrelevant, and both viewpoints retrieve the freedom of choosing cosine phases. In our case, this stems from the (1) x (1) complex-phase freedom in our choice of saddles 4, 0.
(8)
(b)
6, 00)
Figure 3: GFL to GMFL-I to GMFL-II (a) Probability distribution of weights as predicted by approach. The GFL phase is represented by the red graphs, where the minimum the action at zero (shared also by the GP limit) dominates the probability distribution. The GMFL-I phase can be seen in the inset graph, and the final GMFL-II phase is shown in purple.(b) The target component of the average network output. Here singularities can be observed at the σ² values where the phase transi- tions occur. The Parameters taken in this calculation are: N 1000, P = 401, of = 0.002/N, 0.0001
DISCUSSION
In this work, we studied two different models exhibiting forms of Grokking and representa- tion/feature learning. We argue that Grokking is a result first order phase transition. To analyze this analytically, extended the approach of Scroussi et al. (2023) include mixtures of Gaussian. The resulting framework led concrete analytical predictions for rich feature learning effects expos- ing, in particular, several phases of learning (GFL.GGMFL-I,,GMFF
in the thermodynamic/large scale limit. Our results also suggest that feature learning in finite FCNs with mean-field scaling can change the sample complexity compared the associated NNGP. Certainly, these describe very dif- ferent behavior compared to the recently explored kernel-scaling Li Sompolinsky (2021); Ariosto al. (2022) approach, wherein feature learning amounts multiplicative factor in front of the output kernel. A potential source of difference here is their use of standard scaling, however, this remains be explored.
As our results utilize rather general formalism Scroussi et al. 2023). we believe they generalize to deep networks and varying architecture. As such, they invite further examination of feature learning in the wild from the prism of the latent kernel adaptation. Such efforts may provide, for instance, potential measures of when a model close to Grokking by tracking outliers in the weight or pre-activation distributions along dominant kernel eigenvectors. As latent kernels essentially provide a spectral decomposition of neuron variance, they may help place empirical observations on neuron sensitivity and interpretability Zeiler & Fergus (2014) finmer analytical grounds. Finally, they suggest novel ways of pruning and regulating networks by removing low-lying latent kernel eigenvalues from internal representations.

---

Published as a conference paper at ICLR 2024
REFERENCES
S Ariosto, R Pacelli, M Pastore, F Ginelli, M Gherardi, and P Rotondo. Statistical mechanics of deep learning beyond the infinite-width limit. arXiv preprint arXiv:2209.08 04882, 2022.
Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers net- works. arXiv e-prints, art. arXiv:2302.05882 February 2023. doi: 00.48550/r2xi..330..55555
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research, 22(1))4788-444 -8388, 2021.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jac Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems, 1119768-7783 2022.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems, 35:32240-32256, 2022.
Supriyo Chakraborty, Richard Tomsett, Ramya Raghavendra. Daniel Harborne, Moustafa Alzan- tot, Federico Cerutti, Mani Srivastava, Alun Preece, Simon Julier, Raghuveer M. Rao, Troy D. Kelley, Dave Braines, Murat Sensoy, Christopher J. Willis, and Prudhvi Gurram. Interpretabil- ity of deep learning models: A survey of results. In 2017 IEEE SmartWorld, Ubiquitous In- telligence Computing, Advanced & Trusted Computed, Scalable Computing Communica- tions, Cloud & Big Data Computing, Internet of People and Smart City Innovation (Smart- BOOUUSCALCOMMUCCCTCCCCCC/C///// 8397411.
pp. 1-6, 2017. doi: 10.1109/UC-ATTC22011
Omry Cohen, Or Malka, and Zohar Ringel. Learning curves for overparametrized deep neural networks: field theory perspective. Physical Review Research, 3(2):023034, 2021.
Alain Durmus and Eric Moulines. Nonasymptotic woweerreeeccecceececoooeoe analysis for the unadjusted langevin algorithm. The Annals of Applied Probability, 1553 2017.
David Gamarnik, Eren C Kızıldag, and Ilias Zadik. Stationary points of shallow neural networks with quadratic activation function. arXiv preprint 19Xivv191,.01559,, 2019.
E. Gardner and B. Derrida. Optimal storage properties of neural network models. Journal of Physics Mathematical General, 21:271-284 January 1988. doi: 10.10880305-44002111181
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.0267. 2023.
Géza Györgyi. First-order transition to perfect generalization in a neural network with binary synapses. Phys. Rev. A, 41:7097-7100, Jun 1990. doi: 10. 10.1003Phy:RerAA41..099 URL. https:://inkk aps
7097
Lev Davidovich Landau and Evgenii Mikhailovich Lifshitz. Statistical Physics: Volume 5, volume 5.
Elsevier, 2013.
Qianyi Li and Haim Sompolinsky. Statistical mechanics of deep linear neural networks: The back- propagating kernel renormalization. Phys. Rev. 11:031059, Sep 2021. doi: 10.1103/PhysRevX 11.031059. URL htt.h::tttss.k.... wttpppp...///o///////p.o/.. http:////......11////////eeeeee
11.031053
Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural In- formation Processing Systems, 2022a. URL https:://oeeneeview net. forum?? ?id-
6at6rB312m
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. arXiv preprint 20.00:22 2022h.
10

---

Published as a conference paper at ICLR 2024
Song Mei, Andrea Montanari, and Phan-Minn Minn Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 1 115(33):E7665-ET)) 2018. ISSN 0027-8424, doi: 11.1073/pna.18800559909 content/115/33/ETT
URL https: www. pnas. org/
Alireza Mlousavi-Hosseing Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn fow-dimensional representations with sgd. arXiv preprint M.Xiv:2209.14863. 2022.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. 2023.
Gadi Naveh and Zohar Ringel. A self consistent theory of gaussian processes captures feature learning effects finite cnns. Advances in Neural Information Processing Systems, 34, 2021.
Gadi Naveh, Oded Ben David, Haim Sompolinsky, and Zohar Ringel. Predicting the outputs of finite deep neural networks trained with noisy gradients. Physical Review E, 104(6), Dec 2021. ISSN 2470-0053. doi: 10.1 103/physreve. 104.064331.. URL http://dx www.dxx.cci.ooooo wrg.:/....rr.... 11033
PhysRevE. 104. 064301
Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pp. 29-53. Springer, 1996.
Radford M Neal et al. MCMC using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Daniel A. Abolafia, Jef- frey Pennington, and Jascha Sohl-Dickstein. Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes. arXiv e-prints, art. 180.5:1810.05148, October 2018.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning). The MIT Press, 2005. ISBN 026218253X.
David Saad and Sara Solla. Exact solution for on- on-line learning in multilayer neural networks. Phys. Rev. Lett., 74:4337-4340, May 1995, doi: 10 URL https:
aps http://ooo....11////yeeeeeeee... 74.4433777744
Stefano Sarao Mannelli, Eric Vanden-Eijnden, and Lenka Zdeborová. Optimization and gener- alization of shallow neural networks with quadratic activation functions. Advances in Neural Information Processing Systems, 33:13445-13455, 2020.
Inbar Seroussi, Gadi Naveh, and Zohar Ringel. Separation of scales and thermodynamic descrip- tion of feature learning in some CNNs. Nature Communications, 14(1):908, 2023.
H. S. Seung, H. Sompolinsky, and N. Tishby. Statistical mechanics of learning from examples. Phys. Rev. A. 45:6056 Apr 1992. doi: 10.1 10.1103//hy.Ree4.5..00555 URL https://
1ink. aps or.. http://o//1i/ /100 100 1113//1yssee......m 45. 6056
Yan Shuo Tan and Roman Vershynin. Online stochastic gradient descent with arbitrary initialization solves non-smooth, non-convex phase retrieval. arXiv preprint arXiv:1910.12833 2019.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle, 2015.
David Tong. Statistical physics. University of Cambridge, 2011.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint arXiv:2309.02390. 2023.
Bojan Žunkovič and Enej llievski. Grokking phase transitions in learning local rules with gradient descent. arXiv e-prints, art. arXiv:2210.15435, October 2022. doi: 2.1
11

---

Published as a conference paper at ICLR 2024
Karl Weiss, Taghi M Khoshgoftaar, and DingDing Wang. A survey of transfer learning. Journal of Big data, 3(1):1-44))) 2016.
Christopher Williams. Computing with infinite networks. Advances in neural information process- ing systems, 9, 1996.
Zhao You, Xiaorui Wang, and Bo Xu. Exploring one pass learning for deep neural network training with averaged stochastic gradient descent. In 2014 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 6854 6858, 2014. 6854928.
doi: 10.1 10.1109/CAASPP20144
Matthew D. Zeiler and Rob Fergus. Visualizing and understanding convolutional networks. In David Fleet, Tomas Pajdla, Bernt Schiele, and Tinne Tuytelaars (eds.), Computer Vision - ECCV 20/4, pp. 818 833, Cham, 2014. Springer International Publishing. ISBN 978-3-319-10590.11
12

---

Published as a conference paper at ICLR 2024
SUPPLEMENTARY MATERIAL
A TEACHER- MODEL
A.1 SCALING LIMIT OF TEACHER-STUDEN MODEL AND SAMPLE COMPLEXITY ASPECTS
Here we detail the scaling limits in which we believe our results become exact. To this end, we first point the factors controlling the three approximations employ and discuss scaling limits in which these factors all vanish.
Our first approximation is the mean-field decoupling between the read-out layer and the input layer. Specifically we follow Mei et al. (0018); Seroussi et al. (2023); :: Bordelon & Pehlevan (2022), and introduce a mean- field scaling parameter X define of2 0./x (i.e. the microscopic variance each readout layer weight is 52 (x and x² = o² (i.e. the noise variance on the gradients is a²). Notably, since b, e xx o²/[o² 2 K] and K x o²a we finds that x has no effect on b or c. Examining the action, this yields an additional constant factor of x in front of the non-linear term, as expected since mean-field scaling increases feature learning. Taking X > 1 controls the accuracy of our mean-field decoupling between the read-out and input layer.
Our second approximation is the continuum limit, where we trade summation over data-points with integrals. This approximation is closely related to the EK approximation Rasmussen & Williams 20055;; Cohen et al. (2021). Taking n to be large while keeping n=((²²) =\t Nett fixed, corrections to this approximation scale 1/ō2. Notably taking large o² while keeping ner fix does not change b, c or the action of the weights as the latter contains only the combination n1/ō².
The third approximation we carry is the saddle saddle-point approximation colussian-miture approxi- mation for p(w). Examining our action
w
(w w² 2xm
2c/e] (w w=)
S [w)] = d
b[neft/
(I)
20
RdN
+ 100"
1 w
where we further indicated the dependencies of b and e which arise through GPR, comes out as controlling the saddle point approximation, provided the other coefficients in the action do not scale down with d.
We thus find that the following family of scaling limits controlled by two scaling parameters (a, 3 0o) ensures S(w) is exact and its saddle points dominate w
N + 8N
(2)
d / 3d
+ V Bx
2² a 02
0²
a
+ na
Notably increasing a makes the continuum approximation exact Cohen et al. (2021) whereas in- creasing 3 makes both the saddle point and the mean-field decoupling exact. Furthermore, it can be verified that both a and 3 have no effect on Š and in particular keep the coefficient in front of the non-linear term invariant. The only conflict between these two limits in x² hence we require a² >> 3, but otherwise one is free to scale these up differently.
We conjecture that the teacher-student model has a better sample complexity than the infinite- witth network described by the GP limit. Namely, the GP limit is able to learn the cubic target at, n oc d3 whereas the former can reach good accuracy at n ox d. Replacing n with the above, fet the latter is implied by the linear scaling of and d, Reef together with our analytical results, which show that the cubic component being learned. In general, reduction of complexity has been seen in the context of Gradient descent in Bietti et al. (2022); Sarao Mannelli al. (2020); Gamarnik et al. (2019) as well as in online learning up to logarithmic factors Arous et al. (2021) Mousavi-Hossemn et al. 2002) Tan & Vershynin 2019).

---

Published as a conference paper at ICLR 2024
Establishing this more rigorously requires showing that itself, and not n/0² can be kept propor- tional to d. Alternatively stated, that a can scale as vB while introducing corrections to the theory proportional to vB/a, which can then take be small as desired. This scenario seems to be consistent with the corrections to EK derived in Ref. [Cohen et al. (2021), which when evaluated in simple settings (random data on hypersphere and rotationally invariant kernel) indeed show that vB/a controls corrections to EK.
A.2 CONTINUUM ACTION AND MODE DECOUPLING
Our first step here consists of removing the randomness induced by the specific choice of training set. To this end, consider taking large limit while keeping n/o² fixed. This has several merits. At large, n it is natural to replace summations in Eq. 7 by integrals, which can be eval- uated directly. Furthermore, in this limit, one obtains an analytical handle on Gaussian processes regression via the equivalence kernel (EK) approximation Rasmussen & Williams (200): Cohen al. (202T). Notably, previous works have found reasonable convergence to this approximation, in similar scenarios to ours, for values as low as n 100 and o² = 0.1 Cohen et al. (2021). A key element the EK approximation is the spectrum of the kernel defined by
/
(x)
(3)
where 14 is the dataset measure, which corresponds here to i.i.d. standard Gaussians. Denoting by gx the spectral decomposition of the target function on Φ (2z), the GPR predictions (Eq. are given by
f(x) = E x o 9)0x(#)
(4)
where X+o the learnability factors. Notably, n only enters through the combination, n/o² which we treat as the effective amount of data. Hence, we can control the amount of learning two equivalent ways, by either providing more data points or by reducing the noise.
The resulting limit also enables us to identify the two relevant directions in function space, being H, (z) and Hz(æ). Specifically we argue that at large d the operator ((x, y) is block-diagonal within the space spanned by H1 (x) and Hz(x). Since it acts on a target having support only on Hi Hx(x), Hz),,, this limits f(x) to this space and hence also f(xx) [y(x) f Our claim is shown to be self-consistent in App. A namely, that making this assumption results in Q and which
obey the same assumption. Operationally, we thus take the ansatz
(x) bH, (a) + cH3 (æ)
where b,c E R are some unknown constant coefficients. Later, we derive self-consistent equations determining these constants. Given the above ansatz, we turn the action in Eq. [7] take the continuum approximation where 2p=1 + +n diu, and find that the integral appearing in the action can be solved exactly- & (xx) bHi, (w" x) cH3 (w* æ), we can compute the action in the continuum limit-
S
[w
1202
(b- 3c (20)))))))
dpz (wo a:) Φ (w x:) +o f dits (w" xx)* (w xc)
20 m
2No4
|w] 2
2n2 (w w²
2c (w w0y
20 m
adNot 1 + 2 /
1 + ww
the integrals To, I computed in App- A.4
In principle, one could perform saddle point analysis in two variables. A simplification can be made by noting that the dominant fluctuations are in the direction of w=. We therefore projecting
2

---

Published as a conference paper at ICLR 2024
into the subspace of w", such that, Pw- + P w, with Pw-
|lm w0 (w")? and
Pw- la ..-. Denote a = w" w and a = Pw- w| the action is given by
d aaaa ai
d
4n 4n
w 1 a²
2c w a² (a² aa
Slw) =
b
2
20
~2 1 2 (a + +af)
By concentration of the norm, we replace Tn by its average o yields the following action-
(w )²²²
2n
(w . w=)²
2c (w un"
Slw w = w=]=d
b
20 To
1 0Z + (w - w0))3)
1 1+ 2 0 + (w w
A.3 SOLVING THE EQUATIONS FOR b,c
Since the kernel does not mix polynomials of higher order, the mean network output in the function space spanned by the components of the target. Namely, any function in this dioddimeniioall function space is given by a linear combination of Hx = (x),), Hx) = 76 (z)). The factor of 1 results from the requirement that the basis of this V6
write the GPR expression for
space will be normal. In this basis, the discrepancy and the target is given by,
t
66 y V6E
(9)
In this basis, the dimensional kernel Q given by
S dipddp= Hr (a) (z. H1 (z) S - dpedppli (x) (x, 2) H3 (z)
dppdduu Hs (x) (x, i (z)
i = --
(10)
duzduu Hs (x) Q (x, 2) Hl (z)
where Pa is the standard Gaussian measure, Q is the mean continuum limit kernel defined as b,
(x. 03
Φ (xr un) & (z wr) dw
(11)
Z
Next, we proceed by calculating the matrix element of Q. As shown in the integral appendix A.4 we have
I du th (x) db (w- w)
2 (w" w)
VTv / (1 | 2 (w/l))
i ddy Hs (2)) (2) (x (w.
4 4ww" )))
Von / (1 + 2 |un
Since the fluctuations the directions orthogonal to wn are small, quantity |wo appearing denominator can be approximated to (w w²)² (d- e as explained in the main text. Since we are taking large values of then ليه Additionally, we denote g = wr wt, so that.
24 (.--st 7l1 + dq SA stmtke) Vex(1+2(p +0()))
- Slaihee Vo (2
o²
80 Slmbt.) dg
-J
(11++2(2+02)))
Where Z f dge S(gsb,))) These integrals are computed numerically, given, b, c to obtain self- consistent equations. Thus, Q is function of b, c, and on the other hand use Q obtain expression for b, c through the GPR formula in this two-dimensional subspace,
b
= -t=
[
σ² T
Q
x² +
(12)
VBc
Voc
7
n

---

Published as a conference paper at ICLR 2024
Using saddle point approximation, the action be expanded near it's minima to quadratic polynomial, in order to solve the integrals appearing in the matrix elements of Q. This results in two equations for b, e which can be solved numerically.
A.4 INTEGRAL CALCULATIONS FOR THE TEACHER-STUDEES MODEL
Here we compute the integral ( (ww" x=)ant
2 Φ (wv x)) where a ~ N (0, Ta). This integral
appears both in the computation of the action and in the calculation of the matrix elements of Q. Denote-
1
$ dxee- t (w x) (w* x)
2n+1
(13)
Where Z d2m) is normalizing factor. This integral is given by
Z / dxc 2 o (w x) (w" x))aa 2
1
In =
(14)
w VT det (1 + 2www"
2
(wo' ((0,1+2ww77)-1
2n 100 In-1
2
u w²
(wwo' 22)))))))) (( + 2n 12 In
π V 11+ w/
By Sherman Morrison's formula,
2ww
(1 + 2www")
=1
(15)
1 w
Therefore, recursive expression for In can be obtained in terms of w and 4n-1
In = 7, date 10² & (w x) (w0" x) 2n+1
/
(16)
2 w w
(ww ((0"=x)² x)
((0(1++wwwT)>)) + 2n /2012 In-1-
V π
11+2 2 /w/
The cases n = 0, 1 are the ones used in this work. For = 0 the integral is simply-
w
Io
(17)
V
1 + lw0
For the calculation is slightly more complicated. have
(ww' . xz)2) = E whus (avxz) = Σ w/w'' [(1 + 2wwo")-1
(18)
Σ w/ utj
W
1 2 |100
2 (w w²)
w"'
1 + 2 1
Which yield,
(w -0*)
2 (w w0") w=
11 V / (1
3/w²11
(19)
+ 2 /w0] ?
(1+22 + 2 2202))
4

---

Published as a conference paper at ICLR 2024
B MODULAR ALGEBRA
B. MODULAR ALGEBRA KERNEL AND SYMMETRIES
We are first concerned with the spectral properties of Qum,n'm' the latter being the average of num.n'm m' with respect S'w]. To this end note some useful model symmetries that occur when Traain spans the entire training set. By symmetries here we mean that the training algorithm, neural network, and the data-set are invariant remain invariant under the following transformations:
- [[n + q) mod P. (m + q') mod Pl, and fo + f(ptoto') mod with g, q' I. Taking (n,, m] Zp
II. Taking (n, m] + [qn mod g.m PJ and fp + sap mod p for g E Zp but different than zero.
Notably, Q does not contain any reference to a specific label (p). Hence, both these symmetries symmetries of Q (both before and after any potential phase transition). Specifically, let us denote
by T1 (T2) the orthogonal transformation taking [n, m [n + 1, 1.m] mod P ([n, m] - [nz, +
mod P) and by Cg the transformation taking [n, ml commutation relations holds,
- [qn, qm mod P. Then the following
= - = Q2, Cgl = [Ti T2]
(20)
This implies that Q, Ti and T22 can be diagonalized on the same basis. The Fourier basis (Φk, )
discussed in the main text, clearly diagonalizes Ti and T2 simultancously. Furthermore, as no two ŌK, share the same eigenvalues of Ti and T2, each Øx, must be an eigenfunction of Q.
Next, notice that since [Q, [Q, 0. Cyok, must also be eigenvector of Q with the same eigenvalue. However, since [T122 Cg] * 0 Cg shifts us between, Ø hence implying degenerates. Based the definition of the Fourier modes,
find CyoRR Øyk.,k hence all pairs of momen- tum (k, k' -0-00 g') such that = 49, k' qg' must have degenerate eigenvalues. Notably for prime P for any k, g # 0 there exists a such that gk = g. We furthermore note that an additional "flip" symmetry exists, implying that k, K' are associated with the same eigenvalue as K,, k.
The above results for Q hold exactly regardless of network width, noise, other hyperparameters. Next, we focus on Q at the NNGP limit (N + oo), which according to our saddle-poin treatment, also the kernel at any point before the phase transition (at large P). Here we find an enhanced symmetry which explains its poor performance. Indeed, before the phase transition, The kernel Q is given
Qum,n'm = 0a [Kam K Kwmm + 2K² u'm
(21)
K Kum,n'm = ((W, + W Wm+p)(W, + Wm ww (. (0.
On + 8m-P.m + m + 8m+Pa =
-m,n'' + 6m,m'
where we recall that m E [O...P each index. This extra symmetry implies, via the above arguments, that all k, k', and q. with k = ck' and c'q' are equivalent. For prime, P this results in only 3 distinct eigenvalues the one associated with k; k'' 0, O. the one associated with all 0, # 0 and i 0, modes and the one (A) associated with & * 0, k' * 0. Notably. the target is supported only on this last (P - 1) dimensional degenerate subspace. The equal prior on (P 1)2 eigenfunctions then implies O examples are needed for the GP to perform well on the task here, namely the entire multiplication table.
1]. Given that, one can act with symmetry II separately on
B.2 FOURIER TRANSFORM OF THE ACTION
Here compute the action Fourier space without the weight decay term which appears in the main text. To construct the Fourier transformed action, start from the fluctuating kernel term (Q)
5

---

Published as a conference paper at ICLR 2024
given by
2P-1 Σ
0 Q
Σ
(wm + wr) (Wme + =wn) n, m) (n', ml|
(22)
71,0' -0 m2 P
where have used physics-style bra-ket notation denote the basis which the operator Q is presented. For instance, 2ommo.n,m is given by acting with (no,mo) from the left and |n6- m)) from the right and using (no, 1 mi) nn mo.
Based on the symmetries of the problem we consider the following coordinate transformation
kp
Ami km
Wm
WRe
Wn H
U&e
(23)
Substituting this in the equation for Q:
② = Σ Σ
2m kn + uue 2pt km) Wprr 2pik'n + uwe km |n, m)
(24)
ksk' qu'
m.n
W 2 gn' + uge vye Fm (wge Fannn + uge " (n', m'l
Wge
m'n
Where we can simplify further
J=
Σ (wwwwe (*+N')) + u 2g3 (ko+1)) + ww uke 2p3 (K'++km) + wp u
2ps (pottkmmm)
In, m)
mn
(25)
= P (wy + 0) + 2 0, D. + k²) + waew |k, k²) + wk Uk [k',)))
(which also serves as the definition of the Fourier transformed bra-ket notation) and similarly for the conjugated term:
J' P ((-q, - Wi + (-d'.-ql Wy + (- (q +0) w. wy + (0, - + 4')l (g))))
(26)
We next turn to the quadratic term
P-1
P-1 P-1
Σ
Σ
Wk w 275 (k+k()n
(27)
==0
n=0 k
P-1
P
Wp Wp
k=
using ΣP- given by
2s (k+k')n
P5
The quadratic term in the action (main text equation 3.2 |w0 2
P-1
2P-1
P-1
1w0 2
=
Σ
w
n
+
Σ w = Σ
WKWFk
+
Σ
ULU-K
(28)
==0
m=P
k=
The target can be also expressed terms of the Fourier basis, noting that the mode zero zero for P prime
p-
22kkp/P (2rk(u+))/P (n,, ml Σ
2rigp/
1g. 4)
(29)

---

Published as a conference paper at ICLR 2024
The remaining term in the action (main text equation 23 following the anzats o²t ay, and considering the whole dataset. the involves the kernel and is given by
E
VanT Fa Quann'm
(30)
p.
P-1
E
mn m
Σ
2rigp/ 1q.9) (n', m') Σ 2 1d. ) Qumn'm m'
p.mnn.m'n
q=1
Σ
(n, m 9.99) (n', ml ...9) Qran' mame Σ 10)p//
mn.m'n' 't' q-9
Σ
Σ
(n, m 9.99 (n m ( Qrn' 'mer Ea.gg
mm 9:0"
P Σ al
Σ
Queeme (n, m) (n', mm1 lg.9)
m'nmmnn
P-1
P
Σ
la al @la.9)
Note that the chosen basis orthonormal, so that (t kt) = 6x, and thus all terms with (0) vanish and, and we left with-
2 t,, tt.tt
(31)
P-1
= p3
(t. El (wwuw |k, k²) + u (K, (wgugg (-9, -d1 d1+ + Wg (-q', -9l) [t,
Rklag
pa
(w + uk (w, ( gr a + Wg ug 6 9. -g'rt)
kk q4'
P-1
4p'
Σ
WoU-W-U--
The action in Fourier space then given by
P-1 WKWW
p-1
a0 p2
P-1
S(mb = P
+
UKU-K
WkW U
(32)
2
2
No
k=0
Note that since the weights are all real, the Fourier transforms obey WA rewrite in terms of the parameters |uxl. [uk]
k, we can therefore
P-1
20 p2 P-1
S [1ا] = P
|wx 12 + | |ux(²)
Wwk
(33)
2
Not
0=
So that the action can be separated into independent actions S₁. that each depend on different values of l
S(l0 = Σ P
[wx
2+ lexl²)
200 p2
|w
1
|ual
1
+ 1
2 (lool + +luol2)
2
N
0
= So(lwol ,/0ol)
Σ
Sa /1ve))
(34)
kze0
7

---

Published as a conference paper at ICLR 2024
B.3 SOLVING THE EQUATION FOR
As done in the teacher-student model, obtaining a self-consistent equation for the discrepancy re- quires averaging the kernel with respect to the weights. Here, Q is given by-
Q
(wrwww Ik + 0) + UaUe 10, k + k²) + waes (k; k') warek A:)) 1e aa1""
(35)
-((-q -4' Wallea + (-d's 1-- - u + (- (q + z), Wighe + (0, - + a u.p
Since the action in Fourier space depends only on the magnitude of the weights and not the phase, we deduce that the phases are uniformly distributed on the unit circle and are independent of the magnitudes. There are multiple types of terms that appear in Q and we will calculate each part individually:
1. For terms of the form- WhUp
have-(ugugrwy (wke-W
(Iweel [nne//mer |wag ( (e (clorttontwtt0)))) Since the phases of n, are indepen- (elloetoo)) (c((
dents, then (e (el(tnttia)) = (e!+e) (elOu) and since the phases are uniformly distributed over [0. 2n these averages vanish and if = & then (ellos+tou)) = (c2iie) which vanishes too, that (cllortool) = 5k.-g- Similarly, { (d((uttx)) (wieerrnggg)) = dw, o |wEE 12))12)
If q * # then
-v..-gTThus Thus, we have-
2. For terms such as- (wwwwww Wyyg)) ) by averaging over the phases of uy the term vanishes, and similarly terms like Ugy')) vanish.
3. We are left with terms of the form- (wx www.wrrrgu)))) for this not to vanish, we require Φb + by + th + Diy' O so that either k = k and -q'- k = -9 and k' = -16 or
k = q', k' = -q. The same is true for (,
In conclusion, the average kernel is given by-
|ww |ow (() |k: (k: k
(36)
K
Σ (Iwx) |ww 22) (2]1 (| (0,
Σ
(leal? |uw ())x)))(21.=))))))))))) (2 00, A1 + 10, (0.
Since the target has no 10) component, it is orthogonal to the vectors (k + k,, 0]. (0, h + k'| and (0, 0| that these terms don't contribute when Q acts the target. We are left with-
P-
e
M 3 (1), =
Ami pé
[ww 2 [k, k') dr,
(37)
=
2p pk
|ww (ux?? 1k,k) o-e 2
Since for all & + 0 the actions are identical, the mean
= 02 [Iwm1 loxl2) = o² ! (Iwrl 2 |mxe 2) e- - Su(lwelloolla
ʃ e Sa( - llve [sa)
8

---

Published as a conference paper at ICLR 2024
constant function of for every k. Thus, it can be deduced that the target in eigenvector of Q, with an eigenvalue A. A colfssonsistent equation for a obtained using the GPR formula, where-
(38)
x(a)+ 6
B.4 CALCULATING THE EIGENVALUE THE TARGET
In the limit of > 1, the quantity (Iwall 2 |e (2)
can be estimated using saddle point approxima- tion. The minimum of the action can be found by minimizing with respect to |wx| |
S0
= |wcal
[wkk
+ |wske
(39)
No
80.a²² a
|wpl
1
|ux| + |wwl
Not
$ P
0
|c2|
|ank | | + clu
(40)
Not
a0
|ual
|w 1 + c c/vx|l
Not
If [wwl 1 |0x are both nonzero, then from symmetry we assume that [wx = |0x 2, so that
1
20
(20
2
|cx| 2 |wn!
7 Not + (Not Not -7
= w
If [wal = 0 then either |uxi = 0 or (-7)- but since we chose positive y this is not possible. Similarly choosing |ux| 0 results in |wa]| 0 so that either both are zero, or we have- |ea] |wx|² = wL. Using saddle point approximation, have for le + 0
4e Su ++) ([wxl w+ ) (lux) w4) + 5 ([wx)) ((ex))
P ([wpl (uwl)
(41)
4e (uw
In principle the leading order fluctuations near zero should appear in the probability distribution. However, this term be discarded since the fluctuations in Wk; Ux of order P-1 and this contributes to À terms of order p-2 so these are higher order corrections and can be ignored. With this discrete probability distribution, the matrix elements of Q can be computed as follows
4w e (x(144004)
(wwll lux 1)
(42)
4e -8
C EXPERIMENTAL SETUP
Networks were trained for 360000 epochs, with 10-3 with varying noise and as detailed in the graphs. The networks were trained on 50 different randomly generated datasets, with at least different networks trained on each dataset. Parameters of the networks for the B 1 case were given by-
N = 700, n = 3000, d = 150, o = 0.5, = -033
(43)
with の² (0.1,0.15, 0.2, 0.25, 0.3). This was repeated for 3 = 1 1/8, 1/2, and in each case took = 3. The error bars on the negative log posterior plots were calculated by dividing up the set of 200 trained networks into 5 subsets and computing the histogram of the weights in the w" direction in each subset. The error was obtained by dividing the std the different histograms by the square root of the number of samples per ensemble. In the calculation for the average network output, the value for b, c was computed as an average over the different datasets, and the error is given by the std divided by the number of ensembles.

---

Published as a conference paper at ICLR 2024
D TIME DEPENDENT GROKKING
The equations appearing in section 3 rely on the equivalent kernel (EK) approximation, where the number of data points taken to infinity, allowing the substitution of summations with integrals. This approximation captures the underlying feature learning aspect Grokking, but phenomena resulting from differences between the train and the test data such delayed generalization remain obscured. Here we provide numerical evidence that, away from the EK limit (i.e. for finite datasets and finite GPR noise) the phase transition behavior we predicted using EK persists and becomes accompanied by the expected generalization behavior of models which Grok.
To showcase this, we demonstrate experimentally in the teacher student setup, that for finite sized datasets the qualitative phenomena described in the EK limit hold as well. Namely, first order phase transition into a GMFL phase observed also for networks not in the EK limit. Moreover, in this case delayed generalization between the train and test data is observed during training.
Here we study the teacher-student setup, we train an ensemble of 50 networks for 10 epochs, each with:
n 600, d 150 112
(44)
a² 0.05, 0 = 0.011, o 0.5
We drive a phase transition by reducing the width of the network, N,, from = 2800 to = 700 done in Gromov (2023). As the width of the network approaches infinity the network would learn the cubic component as these require n x d3 Cohen et al. (2021). Lowering the width of learn the cubic component. We comment that opposed to the
the network allows the network experiments shown in fig. (3). we do not use o² to drive the phase transition, since taking vanishing values of o² puts us in the EK limit, and for large values of o² the network would fail learn both train and test, artificially minimizing the difference between the two.
10
N- 2800
1400
GP Timit
8
6
+
S
4
2
0
1
W w
Figure 4: Phase transition GMFL-I phase for networks trained on finite sized data sets. Here we use N to drive the phase transition, and observe that for the smaller Ns two new local minima appear in the action, which become more significant by decreasing the width.
Our experimental results are given in Figs. and Fig. tracks the action S, where S - -log(p(( w- w.)) for different values of N: before the phase transition (N = 2800) shortly after the phase transition (N 1400) and into the GMFL-I phase (N 700). The error bars were computed by calculating the action for five subsets of the ensemble separately. The phase transition appearing here is also first order phase transition, demonstrated above for the EK limit.
10

---

Published as a conference paper at ICLR 2024
4000
3750
3500
Loss
3250
3000
2750
Test
Train
N=700
N-1400
2500
N-2800
8
10
12
14
16
In(epoch)
Figure 5: Dynamic delayed generalization. Here the train and test loss of an ensemble of 50 networks was tracked over the course of 7 - epochs, for networks of three different width. The average loss shown in color, and the individual networks in gray. At N 700 a sharp drop in the test can be observed, is commonly associated with Grokking Nanda et al. (2023) Gromov 2023).
Fig. [5] plots the train and test loss and Fig. tracks the projection the average network output on different normalized components of the target, as function of the number of training steps for the different network widths and averaged over an ensemble of 50 networks. Where we define the linear component of the network to be- H (w" x), and the cubic component as- ella(w" æ), and x taken either from the test or the train dataset. Thus the projection of the network on a certain normalized component is one, it has been learned entirely by the network.
Both figures represent different ways in which there is delayed generalization. the delayed generalization associated with Grokking in Nanda et al. (202)), where after extended period of no change to the test loss or even a slight increase, a sharp drop appears. Another mea- sure of generalization is improving how well the test and train compare in terms of learning output components. In Fig. 6 delayed generalization is observed in the sense that the narrower networks (ie after the phase transition), continue learning the cubic component of the target for the test data, whereas the learning of the train data saturates and slows down.
Fig. 5 we observe
In conclusion, we have demonstrated that feature learning phase transitions discussed in the main text are accompanied by delayed generalization associated with Grokking. The above result, to- gether with the known Grokking behavior the second model Gromov (2023). supports our con- jecture that first-order phase transitions are potential underlying mechanism of Grokking. We note that our formalism naturally extends to the finite size dataset regime by applying Cohen et (2021) to augment the EK results as also used in Refs. Naveh & Ringel (2021); (Seroussi et al. (2023) in the context of kernel adaptation. Our focus on the EK limit in the analytical analysis was mainly for simplicity.
11

---

Published as a conference paper at ICLR 2024
0.30
1.0
0.25
0.8
0.20
Trom
J.H.
|H:11
H3
N=700
N-1400
0.15
GPP
0.6
GP unin
0.10
0.4
0.05
7.5
10.0
12.5 15.0
7.5
10.0
12.5 15.0
In(epochs)
In(epochs)
Figure 6: Learnability of target components. Here we track the output projection of network onto linear (Ihs) and cubic (rhs) components of the target, both for test and train data, as function of the in of the epochs. The linear component of the target is given by the vector in data-space, whose components are given by- H1 (x)) with xy drawn either from the train cubic component is given by- e HH(w" :mmu), where the factor is from the definition of the target.
test. Similarly,
The average output as a function either the train data or the test data can was be projected onto these normalized components. This computation preformed an ensemble of 50 networks, for networks of three different width. The average quantities are shown in color, and the individual networks in gray.
12


Reviews:
[{"review_id": "tk8igMfCsc", "reviewer": "Reviewer_t2DF", "summary": "This paper makes attempts to explain the grokking phenomenon in deep learning via an adaptive kernel approach (Seroussi et al., 2023). Two settings are studied: a student erf-network learning a single index teacher; and a two-layer net with quadratic activation learning modular addition.\n\nThis paper doesn't introduce the background knowledge and related work at all, making it a very hard task to understand what they are exactly doing.", "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.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "* This paper seems to have made some efforts to explain grokking, though I cannot fully understand them.", "weaknesses": "* This paper is poorly written. It claims that it is using an \"adaptive kernel approach\" to explain grokking, but they never explains what this method is. I tried to read the previous works, but they are not easy to read for ML audience, either. I urge the authors to introduce the background better: What is the \"adaptive kernel\"? Why is it important to study \"action\"? Why does the approximation in the paragraph beginning with \"Next,\" in Page 4 make sense? What is the theory of phase transition in physics? How are phase transitions connected to bifurcation and saddle point equations? The current version of the paper contains too many unexplained jargons. Even if I could have spent hours reading previous works to understand the background better, I don't believe the current version of the paper is ready for ML audience to read.\n* If this paper focuses on theory, then it is better to organize the claims into theorems and lemmas. In the current version, there are no theorems and lemmas at all. Everything seems to be stated in a very informal way, making it very hard to check the correctness.\n* In my current understanding, the whole \"adaptive kernel approach\" builds on top of the Langevin dynamics, where isotropic noise is injected into the weights. If there is no noise or the noise is not isotropic, then the theory in this paper may not hold anymore. However, in the paper that proposes grokking (Power et al., 2022), the noise should be from the random sampling of batches and thus can hardly be isotropic, but the grokking phenomenon can still be observed. The adaptive kernel theory in this paper cannot cover this case at all.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_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": {}}, {"review_id": "ZxHkrma9Bl", "reviewer": "Reviewer_PXdZ", "summary": "The paper studies Grokking by building on prior work on adaptive kernel approaches for feature learning and claims that Grokking can be understood as a phase transition between different internal representations of the network. The authors explore feature learning in the two models studied (cubic teacher and modular addition) through this mapping to phase transitions.", "questions": "- “We set the gradient noise level … under the equilibrium ensemble of fully-trained networks.” Feels ambiguous. Could you clarify?\n- Could you clarify Figure 2? Specifically, the caption implies a qualitative change at the phase transition point, but I’m not sure what I’m supposed to be looking for. Also, how is it that increasing the strength (in abs terms) seemingly improves the linear component?\n- Why is there a -1/P term in the definition of the target labels in the modular addition case?\n- Is there a reason why the modular addition task is referred to as a student-teacher setting? It doesn’t feel like a teacher is involved here; the targets are simply the modular addition labels.\n\n**Nits:**\n- Typo: In section 3.1.2, “effecting amount of data” should be “effective amount of data”\n- Better flow if Figures 1 and 2 were moved closer to the text explaining them. Otherwise, the reader might have to keep moving between pages.\n- I think the paper could have a broader reach if improved intuition is provided (e.g., explaining how the terms in the action stem from the loss, regularization, and added noise more explicitly).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The formal treatment of Grokking using the approach proposed appears to be novel and compelling.\n- Understanding the phenomenon through the lens of feature vs. lazy learning and casting it in the language of phase transitions is a direction of research that many prior works speculated about. This paper presents much interesting work in this direction.", "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.\n\nThe general feeling I have from this paper is that the approach somewhat obscures the contributions and implications. The formalism developed here and in the referenced work seems compelling, yet the paper does not go beyond two simple toy models. It’s not immediately clear which aspects can be generalized to other problems and what insights a reader can take away to their particular settings. It would have been great had the paper explored some of the directions suggested in the discussion section. For instance, providing a measure for Grokking on problems in the wild or studying whether using this formalism for pruning/regularization could be a fruitful application. Without these directions, I find it difficult to recommend a strong acceptance. 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_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": {}}, {"review_id": "bTowxUgD7F", "reviewer": "Reviewer_MgBg", "summary": "This paper studies phase transitions in a model of deep neural networks. The authors focus on two distinct tasks: polynomial regression and modular arithmetic. Using an approach to study large width neural networks trained to equilibrium with Langevin dynamics, they derive an action that defines the Bayesian posterior distribution of read-in weight vectors given the training data. The authors show a competition of the Gaussian weight prior proportional to $\\log p(w) \\propto - \\frac{1}{2} |w|^2$ and a data dependent likelihood term which arises from integration of the readout layer. Because of the competition of these terms, the action can transition from having one saddle point at $w=0$ (giving the lazy NNGP kernel + small feature learning corrections) to having multiple saddle points which all contribute to the final predictor. Beyond a certain point, the outlier overlap values dominate the action and the $w=0$ saddle point no longer contributes. These define the 3 phases of what the authors refer to as Grokking transitions and the three phases are Gaussian Feature Learning (GFL), then Gaussian Mixture Feature Learning GMFL-I and GMFL-II.  For tasks with a target function that is controlled by a single direction $w^*$ , they show that the overlap $w \\cdot w^*$ has a distribution which becomes multimodal and that this can be made dramatic by scaling up the likelihood term. The authors argue that this phenomenon can lead to sudden changes in the learned representations and generalization error in neural networks as a function of $\\sigma^2$ , dataset size $n$ or model width $N$ since all three of these quantities appear in the likelihood term, in analogy to first order phase transitions.", "questions": "1. I am a bit confused about the scaling limits. What is the exact order of limits taken? Is width $N$ and data $n$ scaled together in some way? Is the kernel $Q$ always full rank? What about the scaling of $\\sigma^2$ with various quantities? In the Appendix, this is discussed in a high level ($n$ large first, then $d$), but the resulting action still contains factors of $N, d, \\sigma^2$ , etc rather than $O(1)$ quantities, which makes it a bit difficult to interpret. Is it thought that the $n$, $d$ limit commute? I suspect not, as the authors discuss various joint scalings like $n \\sim d^3$ or $n \\sim d^{1.5}$ etc in Appendix. Which (if any) of these are actually adopted when deriving the action? \n2. In the 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$. \n2. Could one study mean squared error so that the likelihood term was intensive in dataset size corresponding to a loss $\\mathcal L = \\frac{1}{n} \\sum_{\\mu=1}^n \\ell(x_\\mu, y_\\mu)$ and parameter distirbution $p(\\theta) \\propto \\exp\\left(  - \\beta \\mathcal L - \\frac{1}{2} |\\theta|^2 \\right)$? If this choice was made does the dataset size still control the phase transition? The reason I ask is because this would correspond to Langevin dynamics on a reasonably scaled and regularized cost function $\\mathcal L + \\frac{1}{2\\beta} |\\theta|^2$ with added Brownian motion with variance $\\frac{2}{\\beta}$, which at large $\\beta$ is closer to how networks are trained in practice ( mean loss, rather than extensive sum over data points). 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.\n4. 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. Why must a saddle point approximation over the density of weights $w$ be taken? One could also imagine computing observable averages by sampling this non-Gaussian density as in works on mean field neural networks like Mei et al 2019 or Bordelon & Pehlevan 2022? Is there something lost by approximating each ``well” around a minimum of the action $S(w)$ with a quadratic? Is this the approach taken by the authors when they predict macroscopic quantities like scale of cubic component in the predictor? Could the authors comment on these different approaches?\n5. Could there be a connection between the phase transition reported here and the transition in https://arxiv.org/abs/2210.02157 Figure 5c where as the feature learning strength $\\gamma_0$ increases, the training dynamics transition from having a convergent perturbation series in $\\gamma_0^2$ to a non-perturbative regime where the power series diverges?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "This paper extends a promising approach to neural network theory, known as the adaptive kernel approach, which studies how the kernels of deep networks adapt to data after feature learning. This paper provides two interesting case studies (polynomial regression and modular arithmetic) where they make progress on deriving an effective action which depends only on overlaps with the teacher direction $w^*$ or $v$. They show that the derived theory is accurate in simulations of networks on these learning tasks.", "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 as it is usually understood where training loss decreases much earlier than test loss during gradient based learning dynamics.  I do not see this as a fundamental limitation of the paper (which I quite appreciate) but mainly as an issue of framing. In my opinion this work is a more fundamental phenomenon than grokking since it pertains to fundamental questions in deep learning such as how width, data and parameterization affect feature learning. The paper is not completely rigorous and relies on various approximations, but I think that is completely fine since it supports its claims with experiments and computations at a physics level of rigor. 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). \n\nThat said, if these issues/questions are addressed I would be happy to raise my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_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": {}}, {"review_id": "ie0ISWDtzx", "reviewer": "Reviewer_64Ja", "summary": "The authors provide a heuristic analysis of two settings (a teacher-student model and a modular arithmetic model), which they show exhibit the grokking phenomenon. They relate this phenomenon to a first order phase transition, and provide a characterization of the effective action, with grokking corresponding to the transition from one minimum to the apparition of local minima.", "questions": "- In 3.1, \"$\\sigma^2/n$ is kept fixed\". Do you mean $\\sigma^2 \\sim n$ in scaling? Wouldn't the result of Cohen et al. rather hold for $\\sigma^2=O(1)$?\n\n- 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. Does it refer to the ansatz (12)? Since the target only depends on $H_{1,3}$, I do not understand why the model would depend on higher order Hermites, nor why this would be a sign of feature learning. Could you elaborate?\n\n- 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?\n\n- (Minor) In Fig. 2, solid lines correspond to the theory? It should be explicitly stated for better readability.\n\n- (Minor)  It would be more compelling to complement Fig. 3 with experimental simulations, similarly to Fig.1 middle and right.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "The question addressed (grokking, and on a higher level feature learning) is an important one, and the authors provide an insightful and intriguing viewpoint  into how it arises. I also acknowledge the fact that the settings the authors set out to analyze are theoretically very challenging, making the provided insights valuable.", "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.\n\nGeneral remarks:\n- 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.\n- 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). \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_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": {}}, {"review_id": "I9Yp4FC3la", "reviewer": "Reviewer_iD3W", "summary": "This paper attributes grokking to a phase transition described by a Laudau-like theory using techniques of Gaussian processes. In the grokking process, representations form more and more \"droplets\", until these droplets come together to form water, completing the phase transition. The physical intuition is clear and makes sense. However, I'm not sure whether this theory really explains grokking or something else, given that there is no empirical experiments that compare NN training results to theory.", "questions": "See weaknesses above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper is well-motivated and well written.", "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.\n* The link to NN training is unclear. How to translate the theory to NN training results?\n* The fact that the analysis only applies to two-layer networks is not satisfactory but understandable.\n* Missing some literature review on grokking, e.g., \"Omnigrok: Grokking Beyond Algorithmic Data\" and \"Explaining grokking through circuit efficiency\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_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:

Decision: Accept (poster)