Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model
Authors: Zhiwei Xu
Zhiyu Ni
Yixin Wang
Wei Hu
Venue: ICLR 2025 Poster (2025)
Keywords: Grokking
feature learning
deep learning theory
Primary area: learning theory

Abstract: ''Grokking'' is a phenomenon where a neural network first memorizes training data and generalizes poorly, but then suddenly transitions to near-perfect generalization after prolonged training. While intriguing, this delayed generalization phenomenon compromises predictability and efficiency. Ideally, models should generalize directly without delay. To this end, this paper proposes GrokTransfer, a simple and principled method for accelerating grokking in training neural networks, based on the key observation that data embedding plays a
crucial role in determining whether generalization is delayed. GrokTransfer first trains a smaller, weaker model to reach a nontrivial (but far from optimal) test performance. Then, the learned input embedding from this weaker model is extracted and used to initialize the embedding in the target, stronger model. We rigorously prove that, on a synthetic XOR task where delayed generalization always
occurs in normal training, GrokTransfer enables the target model to generalize directly without delay. Moreover, we demonstrate that, across empirical studies of different tasks, GrokTransfer effectively reshapes the training dynamics and eliminates delayed generalization, for both fully-connected neural networks and Transformers.

Published as a conference paper at ICLR 2025
LET ME GROK FOR YOU: ACCELERATING GROKKING VIA EMBEDDING TRANSFER FROM A WEAKER MODEL
Zhiwei Xul; ZZiyu Ni Yixin Wang" 9, Wei Hu .
University of Michigan, University of California, Berkeley izhiweixu, xu, vixinw, vvh) @umi ch. edu,, whi.uni@berrelle edu
ABSTRACT
"Grokking" (Power et al., 2022) is a phenomenon where a neural network first memorizes training data and generalizes poorly, but then suddenly transitions near-perfect generalization after prolonged training. While intriguing, this delayed generalization phenomenon compromises predictability and efficiency. Ideally, models should generalize directly without delay. To this end, this paper proposes StokTransfer: simple and principled method for accelerating grokking training neural networks, based on the key observation that data embedding plays a crucial role in determining whether generalization delayed. GrokTra fer first trains a smaller, weaker model to reach nontrivial (but far from optimal) test performance. Then, the learned input embedding from this weaker model is extracted and used to initialize the embedding in the target, stronger model. We rigorously prove that, on a synthetic XOR task where delayed generalization always occurs in normal training, GrekTrans fer enables the target model to generalize directly without delay. Moreover, we demonstrate that, across empirical studies of different tasks, GrokTransfer effectively reshapes the training dynamics and eliminates delayed generalization, for both fully-connected neural networks and Transformers.
INTRODUCTION
"Grokking" is an intriguing phenomenon recently discovered by Power al. (2022), where a neural network first memorizes the training dataset but has poor test performance, and after much longer training, suddenly transitions to near-perfect generalization. Initially reported for Transformer models trained modular arithmetic tasks, the grokking phenomenon has since been observed in other settings such as learning group operations (Chughtai et al., 2023), sparse parity (Barak et al., 2022), and image classification (Liu et al., 2023).
While grokking is an interesting phenomenon, it introduces unpredictability into the training process and compromises its practical efficiency. When the model has interpolated the training data with small training loss but still performed poorly on the validation set, it becomes difficult to predict whether or when the model will eventually generalize. Ideally, we would like the model to make continuous progress during training, keeping the gap between training and validation errors minimal. This raises the question:
How can we effectively modify the training dynamics so that the model generalizes without delay?
In this work, we show that data embedding plays a crucial role in determining the training dy- namies; an informative embedding enables continuous progress during training. To obtain such an informative embedding without excessive computational cost, we propose a novel method called GrokTransfer, which leverages the embedding learned by a weaker, smaller model to accelerate the generalization of larger target model. See Figure la for an overview of frokTransfer.
Specifically, GrokTransfer involves two main steps: (1) Train a weaker model until it groks to non-trivial test performance; (2) Extract the weak model's learned embedding and use linear
"Equal contribution; < Equal advising.

---

Published as a conference paper at ICLR 2025
1.0
IGneEii star (hurdar))
4.injcct
Tol (Slindertt
New Embedding
Target Model
Accuracy
Train
3. Lineaar Map
Loss
Troin
Training Data
0.0
100
1000
1000
Weaker Model
Embedding
Epochs
Epochs
a
(b)
Figure I: (a) Overview the GrokTransfer framework. (b) Comparison of the training dynamics of model trained using BrokTransfar versus one trained from scratch. There is clear phase transition between memorization and generalization we train the model from scratch (blue lines). GrokTransfer (red lines) enables the model make continuous progress, significantly reducing the gap between memorization and generalization. See Appendix A.3 for the detailed experimental setup.
mapping of this embedding to initialize the embedding of the target model. Then, proceed to train the target model. We theoretically study GrokTrans fer the setting of a two-layer neural network trained high-dimensional XOR classification task, where normal training exhibits grokking. We prove that GrokTrans fer enables the target model to directly generalize without delay. We further empirically verify the effectiveness of GrokTrans show grokking. This is done for both fully-connected neural networks with trainable embeddings and Transformers. Figure 1b shows typical training curves GrokTransfer vs. training target model from scratch, on a modular addition task. It shows that grokking and significantly improves efficiency.
on typical algorithmic tasks that
effectively climinates
In summary, our contributions are as follows:
We propose a novel method, GrokT rrokTransfee fer, which leverages the embedding learned from a smaller, weaker model accelerate grokking in the target model.
We theoretically justify validate our method on several algorithmic tasks that exhibit grokking in normal training, demon- strating that Grok frer can effectively eliminate delayed generalization.
in an XOR classification task. We further empirically
1.1 RELATED WORK
Our work draws two themes around grokking and weak-to-strong knowledge transfer.
Grokking. Liu et al. (2022) reported that the model starts grokking when learns the hidden structure of the data. Gromov (2023) showed that grokking is robust different optimizers such as vanilla gradient descent and Adam; and regularization methods including no regularization, weight decay, and dropout. Davies et (2023) hypothesized that grokking and double descent, another surprising phenomenon, are caused by the same hidden mechanism. Nanda et al. (2023) reverse- engincered grokked transformer model for modular addition and reported that the learned algorithm is a composition of trigonometric and inverse trigonometric functions. Merrill et al. (2023) and Varma al. (2023) contributed the occurrence of grokking to the competition of sparse (generalizing) and dense (complementary) subnetworks during training. Zhu et al. (2024) showed that models only grok when the training data exceeds some critical size. Liu et (2023) attributed grokking to large initialization scale and induced grokking real- world datasets such as MNIST and IMDb by initializing models with large weight norm. Further work (Miller et al., 2023; Humayun et al., 2024) showed that grokking also be observed in other scenarios such Gaussian Process regression and multi-class classification with adversarial samples. series of theoretical papers have established rigorous results for grokking/delayed generalization in several settings outside of algorithmic tasks: linear regression with linear models (Zunkovië & llievski. 2022), and binary classification with neural networks (Lyu et al., 2024; Xu et al., 2024). Lyu al. (2024) proved that grokking can be induced by sharp phase transition from kernel regime rich regime. Mallinar al. (2024) trained Recursive
2

---

Published as a conference paper at ICLR 2025
Feature Machines algorithmic tasks and found its training dynamics similar to neural networks, showing that grokking is not restricted to neural networks. He et al. (2024); Wang et al. (2024) found transformers achieve -ittr--distriittionn generalization (2024) provided analytical solutions for complex modular arithmetic tasks and hypothesized that some complex modular polynomial tasks cannot be learned by shallow neural networks. Mohamadi et al. (2024) showed that learning modular addition is fundamentally hard for neural networks in the kernel regime. A related phenomenon, termed "sudden drop in the loss" (Chen et al., 2024; Gopalani et al., 2025; Yang et al., 2025), describes an abrupt drop in loss after an extended plateau during online training.
some tasks through grokking. Doshi et al.
Recent work has proposed several methods to accelerate grokking. Liu et al. (2023) explained grokking through the concept of "Goldilocks zone", a spherical shell of weights, and found that restricting the weight norm to sphere of the appropriate radius during training can accelerate generalization. However, this method introduces instability in the training process and still involves a phase transition. Furuta et al. (2024) suggested initializing the model with weights embeddings from another model that has already generalized on a different task may accelerate grokking, which needs to train the same model on additional data, while our method do not need additional data. Lee et al. (2024) decomposed the gradient at each step and accelerated grokking by amplifying part of the gradient. Interestingly, Minegishi et al. (2024) demonstrated that the gap between memorization and generalization can be nearly eliminated a lottery ticket, a set of sparse mask matrices, is applied to the model during training. However, this lottery ticket can only be obtained by first training the same model under the same initialization till generalization. In contrast, our approach can nearly eliminate the phase transition without requiring additional data or pretraining on the same model.
Weak to strong knowledge transfer. Burns et al. (2023) proposed a method where a small model teacher model. This teacher model is then used to generate pseudo-labels to fine-tune larger student model. Surprisingly, the student model can outperform the teacher. Wang et al. (2023) designed learned linear growth operator, which uses Iearnable linear map of a pretrained small model's weights as the initialization for the large model's weights, to accelerate the training of large models. contrast these works, our method focuses on transferring the embedding layer from a weaker model to the target model and reshaping the training dynamics to accelerate grokking.
is first fine-tuned
1.2 NOTATION
For set S with finite elements, we denote its cardinality by |S] and use Uniform(S) to represent the uniform distribution over S. We denote the set (1. nn by [n]. We use sgn(x) to represent
denote by A,, [Ai,,.... .A....AAi, A..m the i-th row, Aug. [ AJ.1" TT 1ū-+++)xx the i-th j-th rows, and A| the Frobenius norm. We use (xx) max (0, x) represent the ReLU activation function. We denote the inner product between two vectors a, b by (a. b). For two sequences {z,, and {yr, }. we say In Ogg,) if there exists some constant C such that Fn I Cy, for all and 22, 22y,.) if yn O
the sign of scalar æ. For matrix A E R=
ACCELERATING GROKKING VIA EMBEDDING TRANSFER FROM A WEAKER MODEL
2.1 MOTIVATION: THE ROLE OF DATA EMBEDDING
To demonstrate the pivotal role of data embedding in shaping training dynamies, we examine the modular addition task + b mod p. Following settings in Nanda et al. (2023) and Liu et al. (2023), we take p = 113. The dataset consists of {((a, b), y) Josa,tp-11 with label y = (a + b) mod p. 25% of the dataset is randomly sampled the training set. We evaluate four types of embeddings:
One-hot embedding:
Each integer a E [0.p p 1] is represented by its one-hot encoding.
Binary embedding: Each encoded in binary, padded with zeros to the maximum length
[logz(p - 1)] +1.
Fourier embedding: Each a is encoded as a vector of trigonometric functions: [cos( 2ri1e), )in 2miia) frequencies.
P "),"" ",""" , 2miza sin( 2miia )}, where ,,, .4. E N are predetermined
p
3

---

Published as a conference paper at ICLR 2025
1.0
12.
10.
Accuracy
0.6
8-
0.4
0.2
0.0
10
10a
10
10a
10a
104
10
10a
Epoch
Epoch
Train Test
onehot
kinary
fouder
Figure 2: FNN training dynamics using different embeddings for the modular addition task (p = 113). The training dynamics vary significantly across different embeddings. The one-hot embedding and GPT embedding exhibit sharp phase transition. See Appendix A.3 for details of the experimental setup.
GPT embedding: Each embedded using OpenAl's text deexxx-emedddine - 3-small a model (OpenAl, 2024)
One-hot embeddings contain no prior information about the data, while binary embeddings capture the ordinal information of integers. Fourier embeddings, inspired by the analytical solutions learned by neural networks (Nanda al., 2023; Morwani et al., 2024), encode tpsk-ipe information. GPT embeddings encode general information about integers. Figure 2 shows the training dynamics of foed-forward neural network using these embeddings. The training dynamics with one-hot and GPT embeddings exhibit clear grokking behavior, whereas those with binary and Fourier embeddings show continuous generalization progress. Notably, Fourier embeddings enable the model to simultaneously achieve memorization and perfect generalization. We observe that general embeddings like one-hot and GPT embeddings suffer from generalization delay, while embeddings encoded with task-related information allow the model to generalize continuously.
A series of works (Liu et al., 2023; Kumar et al., 2024; Lyu et al., 2024: Mohamadi al., 2024) found that the default initialization scale is relatively large and causes generalization delay. They observed that reducing the ini- tialization scale can accelerate grokking and hypothesized that grokking arises from a time gap between the Neural Tangent Kernel (NTK) regime and the feature-learning regime. However, our empirical findings indicate that grokking persists even after carefully tuning the initial- ization scale (see Appendix 1.3. This suggests that grokking occurs even when the model is not initialized in the kernel regime, implying that the kernel regime may not be the sole cause of grokking. In Figure 3, we compare the changes in the empirical NTK (Mohamadi et al., 2023) corresponding the dynamics in Figure 2. The change of empirical NTK evolves similarly across all four types of embeddings (see Appendix A.3 for details).
coco
Tpoch
Figure Change of empirical NTK.
In conclusion, the choice of embedding significantly impacts training dynamics, and an informative embedding can close the gap between memorization and generalization. However, finding such an informative embedding for specific tasks is not always straightforward. Binary embedding, for example, reduces the sharp phase transition for modular addition but fails to do so for modular multiplication. In the next section, we will show that constructing a task-specific embedding from training data can be a promising approach to obtaining an informative embedding that can accelerate grokking. The embedding construction can be achieved by training a much smaller, weaker model. Here "small" refers to smaller model expressivity. This weaker model can learn informative embedding without achieving optimal generalization. This embedding can then be used to positively influence the training dynamics of the larger target model.
4

---

Published as a conference paper at ICLR 2025
2.2 OUR METHOD: GROKTRANSFER
We propose GrokTransfer, a simple and principled method for accelerating grokking in training neural networks, In more detail, given specific task and training set G. we consider a target model fr that has trainable embedding layer Ey with vocabulary size d, and embedding dimension dr. Our proposed method GrokTransfer works as follows:
1. Train a Weaker Model:
Train weaker model Jw with trainable embedding table Ew Rddxxdw on G, where dw is the embedding dimension in the weak model. Train fw until it groks non-trivial performance on the validation set.
Train the Target Model: Initialize A Ew and randomly initialize matrix B Rdw xdr, Train the target model with an embedding layer set to Er = A B, where both A and B are trainable.
By training a weaker model, the first step aims to obtain an informative embedding that aids the training of the target model. In practice, the weak model can much smaller than target model or even have different architecture (e.g., the weak model can be fully-comnected network when the target model is a Transformer; see Section 4). As a result, training a weak model greatly reduces the computational cost of acquiring informative embedding. This contrasts with the method proposed Minegishi et al. (2024), which requires the target model to be trained till perfect generalization first. In the next sections, will demonstrate, both theoretically and empirically, that even if the weak model only partially generalizes (i.e., has non-trivial but non-optimal test error), its embedding still allows the large model to generalize optimally without delay.
In the second step, we impose low-rank structure A B on the embedding Ey while training the target model. This constraint alters the empirical risk landscape and provides favorable initialization for the embedding table. The intuition behind our method is as follows: by initializing with an informative embedding from the weak model, the target model can bypass the initial phase of pure memorization. Instead, it can start generalizing almost immediately as it begins to optimize the training loss.
3 CASE STUDY: GROKTRANSFER ON XOR CLUSTER DATA
In this section, we theoretically study an XOR classification task and prove that GrokTransfer can eliminate grokking for this task.
3.1 THE SETUP OF XOR CLUSTER DATA
We study the setting where the data [x] 22. ,.p]' signal Fooss E R", Zsignall
Uniform( {+1]}) ()) Tnoise Uniform((=t=)=- and the label y = x 2122 Here is the parameter that controls the scale of the noise. We denote this data distribution by P and consider training datapoints {(r,, v,)) drawn i.i.d. from distribution P. We assume the sample size n to be sufficiently large. specifically larger than any universal constant mentioned in this paper. The data distribution comprises four feature vectors (see Figure 5a for a projected visualization), and the model need learn all four features to achieve perfect generalization.
We denote a width-m two-layer neural network by f(x) = =ITh a a (my,, x)), where wy E ER, E [m] are neurons in the hidden layer and 0j E RR E [m]] are second-layer weights. The model is randomly initialized by
Wy 1.dd
(O, Ip), aj lid N(O, a jE [m].
Define the empirical risk with the exponential loss as: T.(f) = En Uy,, fx:)) where 1(y, 7) =
exp(-vğ). We use gradient descent (GD) with weight decay of(til) = (1- 2)0{0) = to update both layers (wy, az) m where is the coefficient of L2 regularization.
Vo ao,
Setting = 80000, n = 400, E = 0.05, this configuration approximates one of the distributions explored in Xu et al. (2024), where grokking was observed. Under this setup, we train two-layer neural network on {(,,, 9: 9:))² Pi with default Py Torch initialization. We observe grokking, as shown in Figure 4(a), where overfitting is achieved by the fifth epoch and generalization begins around the
5

---

Published as a conference paper at ICLR 2025
80-th epoch. Below we will show how method FrokTransfer constructs embedding and eliminates the observed delay in generalization subsequent sections.
1.0
Accuracy
0.8
Accuracy
0.8
0.8
L685
0.6
Welght
0.0
0.6
0.6
0.6
0.4
-0.2
0.2
0.4
10a
10a
10'
10'
10a
10
10a 105
Epochs
Epochs
Weight Index (c)
(b)
Train
Test
Neuron
Neuron
Neuron
Figure 4: (a) Training dynamics of a two-layer neural network with hidden width of 2048, where grokking is observed. (b) Training dynamics two-layer neural network with a hidden width The model can only achieve around 75% validation accuracy and phase transition near 100th epoch observed. (c) Visualization of individual neuron weights from the model trained in (b). shows three distinct patterns and each corresponds to feature direction of the XOR data distribution. See Appendix A.3 for details the experimental setup.
3.2 EMPIRICAL ANALYSIS OF THE WEAKER MODEL
Applying GrokTransfer, first train small two-layer neural network with only 3 neurons fs(x) Σラ= ())((W till convergence (Figure 4(b)). Denote the first-layer weight matrix by W [w,, W2, Wz] E RPx3, the number of training steps by T, and the model after training by f Due to the complexity of the training dynamics, it is hard to derive the closed form of ,(T) and W(T), Below we empirically investigate what information the model has gained and how well it learns.
Figure 4(b) shows that, after training, this weak model has non-trivial performance with test ac- curacy around 75%. The neurons {w (T) (T)) ]=
J== are visualized Figure 4(c), displaying patterns 0) and [-1,-1,0, 0, ..... Oj. Note that the specific features learned by the model are sensitive to its initialization. Nevertheless, we find that empirically, the learned features are always three among the four features [=1, 0,....o) O]. provided the test accuracy is around 75%.
1,1, ... 0]. [1. 1,,0,
Notice that an optimal function for this classification task is
f(x) sgn( + +22) + 6(-21 22) o(-x1 222) 6(01 x2)),
which needs four neurons represent all features [=1, =1]. It thus follows intuitively that the weak model fs cannot achieve better generalization with only three neurons. Formally, we establish the following lemma regarding the expressive power fs-
Lemma 3.1. For any f(x) Σj= ayo(w] x), where & is the ReLU activation function, have
P P(a,))p(y sgn(f(>))) s 75%.
Although the model ( fails to generalize perfectly due to the inherent limitation of capacity, has correctly selected the subset that contains features after training shown in Figure 4(c). Consequently, for any input P. W z becomes high-quality embedding for in much lower dimensional space. Figure 5a shows that, with this new embedding, data points are well- separated in three-dimensional space with relatively high signal-to-noise ratio (SNR) compared to the original embedding.
Next, empirically examine the order of the ratio between the norm of the complementary sub- network and the norm of the generalizing subnetwork. This will be used to estimate the SNR of P with the new embedding. Given the structure of the XOR cluster data, the first two rows of W(7) correspond to the generalizing subnetwork. We define the norm ratio between the complementary

---

Published as a conference paper at ICLR 2025
and generalizing subnetwork as follows:
vle/vp22 [[ 3.p.
rw
[ 12.11///22
Figure 5b and show that the norm ratio is proportional to e and 1/ T/ i.e, rw x e/vn. We will use this property to show that, under mild assumptions, the target model can learn this low- dimensional XOR task with just one step of gradient descent.
spsiler
opelon
1/40
1/150
0.6
Norm Rado
10-1
Norm Ratio
20
0.4
0.2
x3
Class +1
02
6x 10-4
10
-0.6
0.2
X2
0.5
0.4
0.5
9
4''10"
Feature Dimension
6x
4310 6X12
Sample Size (n)
2x 3X331x10
(b)
(C)
Figure 5: (a) 3-D Visualization of the distribution P with the embedding from the weak model. The clusters are well-separated under the new embedding. (b) Norm ratio rw for different values p and e with fixed sample size n, indicating that rw does not depend on (c) Norm ratio rw for different values of n and E with fixed feature dimension p. For each c, the slope around 1/22, indicating that rw is proportional to 1//n. See Appendix A.3 for details of the experimental setup.
3.3 THEORETICAL ANALYSIS THE TARGET MODEL
In this section, we theoretically analyze the behavior of CrokTransfer on the XOR clus- ter data, We consider the target model as a large model with width of the form f2(x) Σn= ayd((uy, U xT),,, where U = [[1, U2, u3] E RRx3 comes from the first-layer weight matrix W(T) learned by the weak model (visualized Figure 4(c)). Here, is the embedding matrix being transferred from the weak model fs. which will then go through another linear transformation (given by by 's) to form the embedding in the target model. Following our observation in Section 3.2, we can write
"1 = [n (rTT,T]].
= [-MuT [ 6.
=
where 141 [1, u/2 = [-1, 1] are two orthogonal features P. and 6, = [6,,1 -5.--2] 2] RP 2j [3)..' Here we let 5 = [61 82, dgl 3.p. Given universal constant C assume
(A1) The noise scale S (n/(plog*) w/
(A2) The norm of the complementary subnetwork satisfies ||ō|| CEvp/n.
(A3) The initialization scale Vinit < Co (Cl)-_((n))
(A4) The step size VriWini// CC S a S VITWinit.
(A5) The number of neurons satisfies 2 2 log
Here Assumption (A2) corresponds to the finding that rw O 6/ VT. T in Section 3.2. Assumptions (A1) and (A2) together ensure that the SNR of the distribution P in the embedding space is large enough. Assumption (A3) controls the initial weight norm of the target model such that the empirical risk starts within reasonable range. Assumption (A4) guarantees that the step size appropriately balanced: it is neither too small to prevent meaningful updates after a single-step gradient descent nor
We assume that the weak model learned three features [1, 1], [-1, 1]. 1].[-,1, 1-,,- -1] without loss of generality. Our result will hold the same for any three features among the four features [+1,+1]
7

---

Published as a conference paper at ICLR 2025
too large cause overly drastic movements. Assumption (A5) ensures that the model's 's width is large enough to ensure certain concentration results about the random initialization, All assumptions are satisfied in the empirical setup discussed in Section 3.2.
and V We denote a = [a [a1... aml T R follows:
[01 [ E R3x,,
We initialize a and V as
di.d ~
Lid Uniform( ((11//))) 1/ ~ Uniform (mnn E [m]]
and keep a and U fixed during the training process. Following the training method outlined in Section 3.1, we use gradient descent V(++1) = V(+) Vv ( at step t to update the linear layer V. where a is the step size and the empirical risk Σ(-) is defined in Section 3.1. With the assumptions and initializations, we state the theorem that characterizes the train and test error of the target model after one step.
(1//n// over the generation of the training data and initial weights of L after one step of training, the classifier sg (1) (x)) can correctly classify all training datapoints and generalize with test error no greater than 2 n)))
Theorem 3.2. Suppose that Assumptions (A1)-(A5) hold. With probability at least 1
Theorem 3.2 shows that with GrokTransfer, after just one step of gradient descent, the target model overfits all training data and achieves near perfect test accuracy. Notably, this is not in a kernel regime but a feature learning regime. Since models with normal training cannot achieve generalization in one step (Figure 4(a)), this result indicates that our method brokIranse fer effectively boosts the generalization speed of the target model and eliminates the time gap between overfitting and generalization. Empirically, the model continues to generalize with further training (see Figure in Appendix A.3). Given that the weaker model is has only three neurons, the computational cost of training fs is negligible compared to the cost of training the target model fr with sufficiently large width. This implies that TrokTraasffr may reduce the overall computational cost. In the next section, we will compare the computational cost of our method to that of standard training procedures.
EXPERIMENTS
This section empirically studies GrokTransfer in modular addition and multiplication, as well as the sparse parity task. Our experiments verify that GrokTransfer effectively reshapes the training dynamics and eliminate delayed generalization for both fully-connected neural networks (FNN) and Transformers (TF). The AdamW optimizer (Loshchilov & Hutter, 2019) is used in all experiments in this section.
4.1 FNN- > FNN
We first consider a three-layer FNN as the target model and conduct GrokTrans on tasks including modular addition, modular multiplication, and (g, kpariiyy (Barak et al., 2022). These results are compared to training a target model from scratch. The modular addition task is introduced in Section 2.1, and modular multiplication defined similarly with the label y = ab mod p. The (q. --)-parity task consists of dataset {(x, y) x E {+1]", = Ihes Tı, [S| = k). Following the setting in Merrill et al. (2023), we choose 40, k = 3. and S = [1,
For the modular addition and multiplication tasks, we employ two-layer neural network with trainable embedding as the weak model, which we train for 10 epochs. We then initialize the target model by setting its layer training dynamics of the weak model, the target model trained via GrokTrans and the target model trained from scratch. Notably, GrokTrans fer nearly eliminates the sharp phase transition observed in normal training. Here all training hyperparameters (initialization scale, learning rate, weight decay) are selected by grid search, and the best configuration is defined as the one that reaches 99% test accuracy the quickest. The oscillations of accuracies in the second row of Figure 6 are related to the "slingshot mechanism" (Thilak et al., 2022) and training instabilities associated with large learning rates (Wortsman et al., 2024). Since large learning rate and this kind of oscillation
the embedding learned by the weak model. Figure 6a and 6b show the
?Our result will not be affected if and U are also trainable. We set them fixed to simplify the analysis while still conveying the main ideas.
8

---

Published as a conference paper at ICLR 2025
HR Oneil
H40 Steal
0.6
Accuracy
Accuracy
Accuracy
0.8
0.7
0.2
0.6
0.5
10000
10000
100
1.0
(204 Tau Sangel
0.8
0.8
0.9
Accuracy
Accuracy
Accuracy
0.6
0.4
0.4
0.7
0.2
0.2
0.6
0.0
100
100
1000
10
100
1000
10
100
1000
Epochs
Epochs
Epochs
(a) Modular addition
(b) Modular multiplication
(c) (40, 3)-parity
Training dynamics of FNNs on various tasks. The rows represent different models/training methods: The first shows the dynamics of the weak model used in GrokTranst fer, the second row shows the dynamics of the target model trained using GrokTransfer and the target model trained from scratch. The columns represent different tasks: the first column for the modular addition task, the second column is for the modular multiplication task, and the third column is for the (40, 3) 3)-parity task. The comparison between the first and second rows shows that the target model trained via GrokT fer can surpass the weak model's 's performance. The comparison within the second row shows that GrokTrans fer eliminates the sharp phase transition and enables the model to make continuous progress. See Appendix A.3 for details of the experimental setup.
Figure 6:
are believed to help generalization (Damian et al., 2023; Lu al., 2024), we do not change our configuration selection criteria.
For the parity task, we use three-layer FNN as the weak model, as empirical evidence suggests that two-layer FNN without bias terms cannot generalize this task. The weak model is trained until achieves 70% test accuracy, after which the first layer's weight matrix is transferred to the target model. As shown in Figure 6c, the weak model undergoes generalization delay, but the large model inheriting its embedding generalizes continuously.
Ablation study: To further understand the empirical effectiveness of GrokTransfer perform an ablation study by varying the training epochs of the weak model in the modular addition task. We extract the embeddings of the weak model at epochs 100, 500, 800, 900, 1000, 1100, 1500, and 2000. For each embedding,
Accumcy [Smal]
Accurec (aagee)
apply GrokTrans fer to the target model and train it for 10* epochs. To measure the gener- alization delay of the target model, we define Time Gap as the difference between the first epoch that achieves 95%
Accurazy
0.8
:/ Time Cap
C20
Trmm
02-
training accuracy and the first epoch that achieves 95% test accuracy. If the target model fails to reach 95% accu- racy, we set 1/Time Gap = 0. Figure 7 shows that the test performance of the target model, initialized with the weak model's 's embedding, is positively correlated with the test performance of the weak model. A grokked weak model essential for the target model to achieve near-perfect generalization with minimal generalization delay. We hy- pothesize that the target model can only generalize well after the weak model has grokked.
100
300
1800
2000
Fpochs
Figure 7: Ablation study showing the effect of the weak model's performance on the test accuracy of the target model (initialized via GrokTram fer and trained for 10 epochs).
4.2 FNN- + TRANSFORMERS
Interestingly, we find that the embeddings extracted from the weak FNN model can be transferred to the target model even when the target model is Transformer comparable to the scale of GPT2- small (Radford et al., 2019). Under this FNN + TF setting, GrokTr aransferr still mitigates the generalization delay of the target model. Specifically, choose the target model to be Transformer

---

Published as a conference paper at ICLR 2025
1.0
(Grnffrrasffer
1.0
(GrokTransferr
Test. (GrokTransfer)
Test (GrokTransfer)
Train (Large]
Train (GrokFast)
(Large)
(GrokFast)
Accuracy
0.6
Accuracy
0.6
0.2
0.2
0.0
10'
Epochs
Epochs
Figure 8: Training dynamics of Transformers Modular Addition Task. The weak model is a three- layer FNN. (a) Dynamics of the target model (an 8-layer transformer) trained via and the target model trained from scratch. (b) Dynamics of the target model (a two-layer Transformer) trained via GrokTrans and the target model trained via GrokFast (Lee et al., 2024).
with attention layers, (dembed: dmlp: Thead) = (512, 512, 4). For each sample, the input is sequence with two tokens (a,b) b). We extract the embeddings of the weak model in Figure at the point that first reaches 30% test accuracy. Figure 8(a) shows that Grok T rans mafee fer enables the target model to generalize much faster than training from scratch and exhibits little generalization delay. Here both method suffer from training instability of large learning rates.
In terms the computation cost, we use wall-clock time as the measure. The computation cost of GrokTran ans fer comprises the training of weak model and the training of target model. Table shows the total wall-clock time for weak model, target model with GrokTrane fer, and target model trained from scratch. The time spent training the weak FNN model is negligible compared to training the target transformer model. The total wall-clock time of GrokTrans fer r is approximately five times faster than training from scratch.
Model
Weak
Target (GrokTransfer) Target (scratch)
Total Wall-clock time (ms)
2828
71079
392667
Table Comparison of total wall clock times (forward and backward passes) for different models. The weak model is three-layer FNN, The target/large model is an 8-layer transformer.
Lee et al. (2024) proposed a gradient amplification algorithm GrokFast to accelerate grokking. We with GrokFast ast in Figure 8(b). The weak model embedding we transfer
compare GrokTrans is the same as the one used in Figure 8(a). For the target model, we follow the model used in Lee et al. (2024), which is two-layer decoder-only transformer with (dembed dmlp- Thead) = (128, 512, 4). The Time Gap of ProkTransfor is 46 while the Time Gap of GrokFast is 1119.
CONCLUSION
To eliminate the unpredictability associated with grokking, we proposed GrokTransfer, novel method that effectively accelerates grokking by transferring the embedding from a weaker model. Our method was inspired by the key observation that data embedding critically shapes training dynamics. We theoretically justified GrokTran rookrransfee fer on an XOR classification task. We also empirically evaluated on various algorithmic tasks known exhibit grokking under standard training. Our results showed that GrokTrans fer can effectively modify training dynamics, enabling continuous progression model performance.
One limitation of our work is that the theoretical result only considers relatively simple XOR task. For this task, after transferring the embedding from the smaller model, one step gradient descent suffices for both memorization and generalization. Theoretical justification for more complex problems an important future direction. Furthermore, our method focuses solely on accelerating grokking and was only investigated on problems where grokking occurs. It would be interesting to study whether similar ideas can be applied to improve training dynamics or enable seak-to-strong generalization in broader context.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported in part by the Office of Naval Research under grant number N00014-23-1- 2590, the National Science Foundation under Grant No. 2231174, No. 2310831, No. 2428059, No. 2435696, No. 2440954, and a Michigan Institute for Data Science Propelling Original Data Science (PODS) grant.

---

Published as a conference paper at ICLR 2025
REFERENCES
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit, Advances in Neural Information Processing Systems, 35:2 00:21150-22164 2022.
Collin Burns, Pavel Izmailov, Jan Hendrik Kirchner, Bowen Baker, Leo Gao, Leopold Aschenbrenner, Yining Chen, Adrien Ecoffet, Manas Joglekar, Jan Leike, Ilya Sutskever, and Jeff Wu. Weak- to-strong generalization: Eliciting strong capabilities with weak supervision. arXiv preprint 10X19:2312.09390. 2023.
Angelica Chen, Ravid Shwartz-Ziv, Kyunghyun Cho, Matthew L Leavitt, and Naomi Saphra. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in MLMs. In The Twelfth International Conference on Learning Representations, 2024. URL net/ Eorum? 1d=M0SPiKHHRR
Bilal Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations. In International Conference on Machine Learning, pp- 6243-6267. PMLR, 2023.
Alex Damian, Eshaan Nichani, and Jason Lee. Self-stabilization: The implicit bias of gra- dient descent the edge of stability. In The Eleventh International Conference Learning Representations, 2023, URL hettss//openrevieee
net forum nidnnki
Xander Davies, Lauro Langosco, and David Krueger. Unifying grokking and double descent. arXiv preprint drXiv:23033.06/73 2023.
Darshil Doshi, Tianyu He, Aritra Das, and Andrey Gromov. Grokking modular polynomials. arXiv preprint arXiv:2406.03444 2024.
Hiroki Furuta, Minegishi Gouki, Yusuke Iwasawa, and Yutaka Matsuo. Interpreting grokked trans- formers in complex modular arithmetic. arXiv preprint 2.X.::22022.172666 2024.
Pulkit Gopalani, Ekdeep S Lubana, and Wei Hu. Abrupt learning in transformers: A case study on matrix completion. Advances in Neural Information Processing Systems, 37:55053-550855 2025.
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02677 2023.
Tianyu He, Darshil Doshi, Aritra Das, and Andrey Gromov. Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks. arXiv preprint arXiv:2406.02550, 2024.
Ahmed Imtiaz Humayun, Randall Balestriero, and Richard Baraniuk. Deep networks always grok and here is why. arXiv preprint arXiv:2402.15555 2024.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint 2rXiv:2001.08361 2020.
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics. In The Twelfth International Conference on Learning Representations, 2024.
Jaerin Lee, Bong Gyun Kang. Kihoon Kim, and Kyoung Mu Lee. Grokfast: Accelerated grokking by amplifying slow gradients. arXiv preprint arXiv:2405.20233 2024.
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651-34663. 2022.
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In
The Eleventh International Conference on Learning Representations, 2023.
12

---

Published as a conference paper at ICLR 2025
Iyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019. URL https:://///. http://pppneeeeeeeee/eeee net forum? id= Bkg6RiCqY7.
Miao Lu, Beining Wu, Xiaodong Yang, and Difan Zou. Benign oscillation of stochastic gradient descent with large learning rate. In The Twelfth International Conference on Learning Representa- lions, 2024. URL https::/ /tpp/ppeeeeeeeeeeeeeeeeeeeeeeeee reenrevien net. / forum? 1d==WmwN3 sQp..
Kaifeng Lyu. Jikai Jin, Zhiyuan Li, Simon Shaolei Du, Jason D Lee, and Wei Hu. Dichotomy of early and late phase implicit biases provably induce grokking. In The Twelfth International Conference on Learning Representations, 2024.
Neil Mallinar, Daniel Beaglehole, Libin Zhu, Adityanarayanan Radhakrishnan, Parthe Pandit, and Mikhail Belkin. Emergence in non-neural models: grokking modular arithmetic via average gradient outer product. arXiv preprint arXiv:2407.20199 2024.
William Merrill, Nikolaos Tsilivis, and Aman Shukla. A tale of two circuits: Grokking competition of sparse and dense subnetworks. arXiv preprint arXiv:2303. /1873, 2023.
Jack Miller, Charles O'Neill, and Thang Bui. Grokking beyond neural networks: An empirical exploration with model complexity. arXiv preprint arXiv:2310.17244 2023.
Gouki Minegishi, Yusuke Iwasawa, and Yutaka Matsuo. Bridging lottery ticket and grokking: Is weight norm sufficient to explain delayed generalization? In ICLR 2024 Workshop on Bridging the Gap Between Practice and Theory in Deep Learning, 2024.
Mohamad Amin Mohamadi, Wonho Bac, and Danica J Sutherland. fast, well-founded approxima- tion to the empirical neural tangent kernel. In International Conference on Machine Learning, pp. 25061-25081 PMLR, 2023.
Mohamad Amin Mohamadi, Zhiyuan Li, Lei Wu, and Danica J Sutherland. Why do you grok? theoretical analysis of grokking modular addition. arXiv preprint arXiv:2407. 12332, 2024.
Depen Morwani, Benjamin Edelman, Costin-Andrei Oncescu, Rosie Zhao, and Sham M. Kakade, Feature emergence via margin maximization: case studies in algebraic tasks. In The Twelfth International Conference on Learning Representations, 2024. URL /ttetts://ppereeeee fert/ id 9wDX850
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, 2023.
OpenAl. text-embedding-3-nnl -3-mmlll model. model 1 seseembedding ng 2024.
https://platffot openai cotimaacccomm/osssoeoeooo
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177 2022.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8);9, 2019.
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprini 27Xiv:2206.04811 2022.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprini arXiv:2309.023390 2023.
Boshi Wang, Xiang Yue, Yu Su, and Huan Sun. Grokked transformers are implicit reasoners: A mechanistic journey the edge of generalization. arXiv preprint arXiv:2405. 15071, 2024.
13

---

Published as a conference paper at ICLR 2025
Peihao Wang, Rameswar Panda, Lucas Torroba Hennigen, Philip Greengard, Leonid Karlinsky, Rogerio Feris, David Daniel Cox, Zhangyang Wang. and Yoon Kim. Learning to grow pretrained models for efficient transformer training. In The Eleventh International Conference on Learning Representations, 2023. URL https: /ooeerveeeeeeeeeeeeeeeeeeeeeeeee ew net/ forum wet.rrmmmiiiiiiieeeeeo d=cDYR551 6f.
Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie E Everett, Alexander A Alemi, Ben Adlam, John Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl- Dickstein, Kelvin Xu, Jachoon Lee, Justin Gilmer, and Simon Kornblith. Small-scale proxies for large-scale transformer training instabilities. In The Twelfth International Conference on Learning Representations, 2024, URL https: / wowporpeeeeeeeeeeeeeeeeeeeeeeeee ew net./ forum ? www.dd/wwwwww@e000mmmmmmm
Zhiwei Xu, Yutong Wang. Spencer Frei, Gal Vardi, and Wei Hu. Benign overfitting and grokking in ReLU networks for XOR cluster data. In The Twelfth International Conference on Learning Representations, 2024.
Yongyi Yang. Core Francisco Park, Ekdeep Singh Lubana, Maya Okawa, Wei Hu, and Hidenori Tanaka. Dynamics of concept learning and compositional generalization. In The Thirteenth International Conference on Learning Representations, 2025. URL https: /otpep/e//peeeevie.. net net/forum?s forum? i d=s1 ZO0YBEFB.
Xuekai Zhu, Yao Fu, Bowen Zhou, and Zhouhan Lin. Critical data size of language models from grokking perspective. arXiv preprint arXiv:2401.10463 2024.
Bojan Žunkovič and Enej Ilievski. Grokking phase transitions in learning local rules with gradient descent. arXiv preprint 2Xiv:22100.1435,, 2022.
14

---

Published as a conference paper at ICLR 2025
A
APPENDIX
A.1 Proofs
15
A.1.1 Proof of Lemma 3.1
15
A.1.2
Proof of Theorem 3.2
15
A.2
Additional Experiments
25
A.3
Experimental details
25
A.3.1
Experiments in Section and 2.1
25
A.3.2
Experiments in Section
26
A.3.3
Experiments in Section 4
26
A.4
Additional Discussion
27
A.4.1
Forward Pass FLOPs estimation for Models in Figure Left
27
A.4.2 Additional Experiments
A.1 PROOFS
A.1.1 PROOF OF LEMMA 3.1
Lemma 3.1. For any f(x) Σÿ= ( x)), where o is the ReLU activation function, we have
P P(z,))p(y = sgn(f(x)) S 75%.
Proof. For any (I, v) define x' (x1, - I2, F3, ... Fp) and y/ sgn(x1 =y y. is sufficient to show that if sgn (f) sgn (gn(((( (--r)) v/ sgn(f( then y sgn()(-x') with probability 1.
Assume sz)( and sgn( [gn(s(-20) (-x)) Given d(=) 2 0, Vz, sgr(((x)) implies that there exists at least one E [3] such that aj has the same sign as y and w, x > O. Without loss of generality, assume sgn(a1) y, wi x > Then for (-x,-1 -y), it follows that
f(-z)= = Σ a. a2d(-w3 x) + aso(- x))
has the same sign Again without loss of generality, assume sgn(az) y.
If y' = (xz) and y' * sgn(f(-zz)) hold, following the same discussion, we have that at least two a, 's have the sign as y' -y, which contradicts the previous assumption that sgn( = sgn(az) y.
A.1.2 PROOF OF THEOREM 3.2
Additional notations: For training dataset {(x, 31)) / we denote the signal of x; by E; [x22 #p2). For each μ {##1) 1,2), define
In = (i E [n] E. 1
and n. In,l- Denote the new embedding of the i-th datapoint by 24 UT 2,, [n]. Define
14 1|2, - 112,-1u -p1] p4: i == 1,2.
15

---

Published as a conference paper at ICLR 2025
Then 11
0. 12, -], and { = becomes the features for P with the new embedding. Denote the signal of 2; by 2; = |12. - 12, -41] E, Define the set of training data
Gdata = (1(
|| = s 8² V logn, n, for all i e [m]]
By Lemma A.2, (((,,, y: (data) 2 c ( n)). We further define sets to separate the second-layer coefficients for the ease of discussion:
Tos { E [m] aj > 0};
JNeg = {j E [m] : aj < <0}-
divide the index of neurons by its initialization and define Te = {j E [m]] ej Uinite for E U We further define
Joos,e = Tos n Tei
JNeg.e = INeg N
For each initialization of Uj (0) it by
we denote the set of datapoints which have positive inner product with
= =
0}. E Uniform((=1])).tt
[=#1+#p2)
Theorem 3.2. Suppose that Assumptions (A1)-A5) hold. With probability at least O(1 over the generation of the training data and initial weights of fL, after one step of training, the classifier sg (x)) can correctly classify all training datapoints and generalize with test error no greater than exp( (2002-220g' n)))
Proof. For brevity, we omit the subscript Sr. in the proof below.
At step / = 0: for each (r,, Vi), we have
(x) = Σ (a .4)):
[m]] are bounded random variables with zero mean. The absolute bound is where (o =;)); E
(0)
/3vinit Vm
(1,0((15) =))
S
(max ||
6² Vppn log n) 50 Vmm
where the first inequality uses Lemma A.2 and the second inequality uses max; Assumption (Al). Then by Hoeffding's inequality and law of total probability,
|| = 2 and
F(I t) < (( (x:)) t[Gaata) P(Gaata) I 2 exp (
2t
+ exp((=2(log' n)).
25
init
Let Vinit log.. n. follows that
P(max ((0 (x+) < t 21 2113
(E(x]
2log² 2
2 2n exp(-
") eexp(-2Q(ll" n)) exp(- n)).
25
(1)
We define set of training data and initial weights:
(1l( ((((=,)) v(0): ((,,, H E Gusta condition (1) and
all conditions in Lemma A.I and A.4 hold
Combining (1), Lemma A.I, A.2, and A.4 then applying the union bound, we have
1 - -
1
P(([(z,, w)),, v(0) E g 2 exp( 22oog2 m)) - 0
=
0
n2
n
n?
16

---

Published as a conference paper at ICLR 2025
Denote (8)) Myy, ( (x:)) (xp (x:)). Conditioning the ratio between the maxi- mum and minimum loss bounded by:
R maxie(n] ((0) S exp(2Uinit logn). n).
(2)
minie[n] (
For each j. below we will analyze the gradient descent update for all possible combinations of (0) (0) conditioning the event G. aj
(1) When a, 0: If oj Unit [1, 1), then according Lemma A.I, have
+1.1))+ Ø; -1 I-mri [
Vv n, t12-
2
Recall that the gradient descent update of u is
0(++1) = =vj²)
ag
y expl (x:)) 22))2
(4)
It follows that
J3
333
(0)
+
a
yi
o ((11)) =3)
n
133
(0)
yi (0)
233 +
aj
Σ
yid 21,3
n
n
(EI
Vinit +
20
(0)
0
a
max (10) 2
P
log n)
(5)
n 1E7
m
1.9a/I-pi
1.9a (1
IV Vnnit
+
exp(- Uinit log n) 2 Vinit +
Unii
)(1
Unnit log n)
m
4/m
log n
2a
Unit
5
where the first inequality uses Lemma A.2 and , 0, Itps: the second inequality uses Assumption (A1), (A3) and (A4); the third inequality uses Lemma A.4 and exp(x) 2 1 x. Further for = 1, we have
|0 2/1
Σ
(0) Zi,1
a aj +
Σ
(0) y.d
Vilk
2i,1
IEI
F1
EE 1. 11.
Σ
yit (zi,1 Zt,1)
(0)
(1.
VEI ps UZ (...1.)).
UI..1.)) 43
Σ
(0) 24
Σ
(0)
2d,
iET
iET..1.3).
exp(eini log n)
log n +
2
100)
Σ
(0)
(6)
mm
nym m
iEZp,1.1)--22
exp(Vinit log s² log n2 + Vm
a
2a
exp(Vinic logn))
7 +
72
n/m
8
2
+ Vnlogn exp( - 20 log logn)(
n
n
/olgg))
2
< C 2 VP
log n + C a n
+ Vnnin log n) S C
mmn
logn nm
Vmllnn
where the first inequality uses (2) and ž. -E for E I1. IpI.1.).-paa the second inequality uses (2), (3), and (B4) in Lemma A.4; the third inequality uses Assumption (A1)-(A5); and the last inequality uses Assumption (A1), (A3) and (A4).
17

---

Published as a conference paper at ICLR 2025
For datapoint (z. P. define 2 [21, 23] = U x. Applying Lemma A.3 obtain
P([ 2|| 1 E² V log n) 2 - exp( 22(lg_ n)).
Conditioning on
- 2|| S 22
P log
(7)
if Fsignal = M1,, we combine (5) and (6) and have
(of)) .2) = (v))) 8) + (of) 2 - 2) 2 2(Vinit
20x ) C'iiins 5
L V
log n 2
3
Vnntt +
20 ) (8)
nl
2
5mmm
Further for any pair j1, j2 with (0) = 1/j Vinin [1, 1, 1] and an
0, 0z2 < 0:
(o 2) < O, follows that
2
23 Byy m
Σ vilh
= ,
(0) =;))======>>
(1)
22 + 23
J2.3
(0)
(1)
2 Uini
-
22.1
[=1
2 Unit
2=²
P V n
log Uiiit + C a
23
Vinit
Vmllg n
(0)
where the first inequality uses /33 = Vimit and ž = 0,1 1, 2; the second inequality uses (7); and the last inequality uses condition (7), 2, and Assumption (A1), (A3) and (A4). Combining (5)
and (9), we have
yel (0) " (10(0)
(0) 21))2r1 2 maxf Vinit 2a
55mm }.
m
2
which together with (6) yield that
aj ((( 21 2)) ai (1 (1)
(1)
2)) = ajj (v 2)
(1)
2
vm (1)00 2) + a
(0) 22)) 2i,) + 23 Σ yil (0) Ø' (((11/00)
j.. - 0 11
n/m
1=1
Vinit + max
Vinit
2a
} Coinnt? /
/ P logn - C
=²
logn
(10)
m
2
5mmm
v m log
Vinit
+
Vinit
Cumme"
V n
P
logn C
2²
P n
logn
Vm
4
5
vm log n
Uniit
V m
10m²
where the second inequality uses max(x, 2 (a + v)/2 and the last inequality uses the fact that n is sufficiently large.
If (1) 2) > have
an ((
(1)
z)) + وزz ((() (1)
.)) = mm (1/00 0/ z)
=))
(1)
(1
(0)
(1 (1)
(0)
1 Ji
11
ja 2)
Vm
V m
(1I)
n
2
23
yet
(0)
0'(10
(0)
24,3 +
(u hu
(1)
(2 a)]
Vm
n n
m
2
4a
C
a
2 log n
>
2a
5mmm
m
logn
5m
18

---

Published as a conference paper at ICLR 2025
(0) (0) the third equation (4); the first inequality (5); where the second equation uses "/ and the second inequality uses Assumption (AI). Combining (10) and (11), follows that
2a
On ( 2)) + a ((0)2 z2)
(12)
5m
when Uj1 V)2 Vinit [1, 1, 1] and Fxignal p1.
(0) (0)
If Tsignal +/11 following the same procedure, we obtain that (vf (of), 22 < 0 for ay
For aj
similar to (5), we have
233
(1)
(0)
0j
Σ
Fr,
+
Σ
vel 2,,
iEI-1i
p. UZ..1)) P2
Unit
20
(0) 0
max
100) e²
P
log n)
ny m
m
n
iET
(13)
2.1a m1
2.1a 4/m
4
Vinit
exp(Vinit log n) 2 Vinit
(1
(1
20 log n)
n//m
log
30
Unntt
Vinit
2
4/m
4
where the last inequality comes from Assumption (A4). Then (u 2) also hold for az following the same analysis. Thus we have
a, M(
(1) 2)) .)) as o(
(1) 2)) = 0
(14)
(0)
=1) (0) Un 1, 1] and Tsignal = +H1-
when jj
1/
Tsignal E { }, combining (5) and (6), we have
[(05),))1= < ((gg =) I(, -- +21))1+(05"" - ,,,0, +(05 -,,0== ===
V L Plogn+0+0
c_ mm V22 Llogn S 2vinit 2
L V n
S 0 + Vinits
log n,
where the last inequality uses Assumption (A3) and (A4). Thus
2)) < 2iini
P nm
log S
2
(15)
Vm log n
Note that neurons initialized with Vimi k), E {+1) share very similar dynamics and following
the same procedure, specifically, le +1 (resp. -1). ) the neurons align well with //1 (resp. +/11) Additionally, the neurons do not align well with =22 for both i +1 and i - -1. For brevity, we
omit the analysis for o
(0) Vinili i, k). ls e (=1) (eema [1, 1, 1]3.
Next we analyze the one-step update of neuron Uy with initialization Vinit [ini[1, 11, 1].
(2) If (0) 1 = Vimit [1, -1, 1]. then according to Lemma A.1, we have
I 11,11,1.+2 =0; Ip,-1.11,-Mi Mi Ip,-1111, Itwzi Ip,- Ø. (16)
19

---

Published as a conference paper at ICLR 2025
Similar to (5), we have
(I)
5,3
(0)
2
(j)
a
n2
0j
yilk
(0)
d' 6'(11)
(0)
21) F1,
)
2
a
(0)
Σ
(0)
aj
yil
3
yiu;
31,3
2 aj
n1
iEZ
iEI+1²
(0)
Σ
(0)
nym
3,,
2
2/m
HEI-NI
$EZ+p2
(17)
(0)
21,
+
(0) (1,
21,3)
(213 =2,33
a
n/m
2/ml m
iEI-11
iEI iEI-11 m1
iEI+u2
2a
71
n exp( Vinit l)g + a
(n-p + n+pp) exp (hii log n){²
p V
log
nym m 4
nym
n
= =0( a
(62 V/ + mmn) lo = 0( vm log
a
vm
where the first equation comes from the GD update; the second equation uses (16); the third equation
uses la,l = 1/ /mm the fourth equation uses 2iS = 0 for E I+pzi the first inequality uses
3,3 2, E I - (2) and the definition of G; the fifth equation uses [ n/4 I n/ log n and
Vinit logn) 1| < 2v log n S 2/ n by Assumption (A3); and the last equation uses eep( Uinit
Assumption (Al) and (A3). Further for the first entry of eyy we have
(1)
(0)
(0) a
a
j
Yell
(0)
8' ((e,
(0)
a
2 aj
2
y 21,1
Σ
Vily
(0)
2,11
+ 5
IET M1
teI
(0) 21,1
Σ
(0)
+ a
nym
24,1
2/m
VET
(18)
a
(0)
(0) (31,1 -2,1)
Σ
( (2 - -2,1
+ a
2,,1
+
4
nym
2/m
SE
122
VET m1
20
n+ exp(- Vinic log n)]
mm-m1 m1 + n+va) exp(Vinic log n)²
L logn V
n/m
nym
n
= 0
Q
P
+ Uniit) logn) = =0( Vm logn
Vm
n
where the inequality uses ε₁, that
2 for i E Itps. And for the second entry of uj. it follows similarly
2/ - (r1
(1)
(0)
2
a 0))
=
ny m
Σ
10 2122
Σ VEItuz
100 21,2
2/m m
0
v m log
a
IEZ
(19)
Unifying (17), (17) and (18), we obtain
(1)
(v)) 2 sgn sgn(r,))
(0)
(0) 15i 0
(20)
Vm n
for I = 1,2,3 3. Here {&} are defined as & = -1, = 1,2 and E 1.
For datapoint (x,y) P with z [21, 22, 23] TT]T= = U z. We condition on the event
===
P log n.
n
If Tsignal = - = for each pair ji, with (0)
= 1 (0)
= Unnit [1, -1, 1] and aj > 0, 0, < 0, we have
(0 .2 > 0,1 001-10 1,2, and
[1 - 0/2 vm
α [-1,,
[(0
- 1
10)
(1
- 2/
(0)
a m
= 0
a
m log
20

---

Published as a conference paper at ICLR 2025
by (20). It follows that
a, (0 (1) 2)) a2 (2) 22)
vm ") 52 -2
1 (1)00
1 α
1-1, 1, 11.
(0 1
(1)
(1)
a
1-1,1, 11, + (0/) (1)
1/ 2 -z))
(1)
(21)
Vm
Vm
Vm
2a
0
0
8²
P
logni)
2
Vm
vm
Vm log
vm log n
n
m
If Tsignal +/i we have (0)) 2) < 0,1 1, thus
aj (
(1) .z)) 0320((1)) (1) z2)
If Ssigmal +22 we have ( .) 0,l 1, Applying (20) and Assumption (A4), we have (2) 2) > 0, = 1.2. It follows that
3, (5 (1) 2)) { + 4イæ 0(1 (1) 2)) 1 (1)
.2)
(1)
11 12
vm
/m( (1
[-1, 1, 11, 2) (0/0) "2
(1)
(1)
a
11,,,
+ (00 - (1)
12 --2) -2)) (22)
mml
m
mm
2a
0
0
p V
log m)
--
m Vmllgn'
Vm log
n
m
for sufficiently large n. Here the last inequality uses Assumption (AI).
If Esignal -p2g we have (0) 2) < 0,7 1,2. Applying (20) and Assumption (A4), we have (0) (1)
2) 0, - 1,2. It follows that
H) (1) 2)) + 02 ((1 ,)) 0.
(23)
In conclusion, for datapoint (z. y) with Tsignal
14, conditioning on (7), the output of f()) is
(x) = Σ (1)) ,2)) = Σ Σ Cap((1f)))) 22)
.2)
j=1
eEUniform( jE.)
Σ
Σ
a)(((1))) .²))
,))])))))))-))-))))())))))))(( - Σ () .2))
eleg=1
jEJnon
jEJvng
(24)
2a
4vm
min JinnllPeeeelll |NNeg.el)
(Uiit +
5m
logn n
vm
a
4/m m
(C
>
0
40
log
Vmn
mm
6:63=1
for sufficiently large n. Here the first inequality uses (12) and (21), the property that [J
|JNaggell I 2m/ logn n from (B3) in Lemma A.4 and the property that ((()))) ,2) < 2( a/ vm); the second inequality uses (B3) and Assumption (A4). Similarly, we have that for datapoint (xz, y) with Zsignal +111 conditioning on (7), the output of f()) is
(x) = E Σ
ago((1 (1)
.2))
Σ
ayp((u > 0.
(25)
e:6==1 jEJron.
JEJNeR
21

---

Published as a conference paper at ICLR 2025
For datapoint (a, y) with Esignal +u2, conditioning on (7), the output of f()) is
(x) Σ
(11 =))==
Σ
Σ
agdl ())()))))))))))))))))))))
1=1
UEniformm( (=)) (+1}))
= Σ
Σ
Σ
ap0(1)) ,)) - E
a//1,02/= [1]-1]
JEJpunec
JEJNee.
Σ
| min( e: [NNeR a 4/m (Uinit +
Σ
20mi JNeg.el
log
vm logn
==11-11
e:e1=e2
< 2
a
+
5mmm
(Uinit +
a )
+
Umnt V m
a
10(C + 1)ax
Ca
< 0,
16
log n
Vlogn 8
logn
Vlog n
(26)
where the first inequality uses (15), (22). (B3) and the property that of (( ()) 2)) 2)) < 2( + a/ / m the second inequality uses (B3); and the third inequality uses Assumption (A4). Similarly, we have
that for datapoint (Ir y) with Tsignal -2,, conditioning on (7), (z) 0, which combined with (24), (25) and (26), yields that
sga (gn( =
for any (x. v) ~ P, UT satisfying
2 El| S E² / logn.
||-
According to the definition of Gaata all (Ii, vi) satisfy this condition. Thus conditioning on the event G. the model (() can correctly classify all training data points. And applying the law of total probability, we obtain that the test error is bounded by:
F(a.))r(V # sgn(f()) & P (()))~(( (v * sgn(s (x)))
=ll logn
+ (() (12- E11 = V
P logn)
-P(ee)> (13- E|| S 2
_ Pogg m) eexp(=2(log' n)), n
where the last inequality uses Lemma A.3.
Lemma A.1. Suppose that Assumption (A2) holds. With probability at least - (2 the following conditions hold:
I,s.-1)itus = =Itpi 11b.--))--1 = =2, i,j E {+1);
Iij+M1t = Ø; Ipsjt1)- = II-p
i,j E (+1};
(27)
I141,11M)tt = Itwai I1+1-1NN1-0 =0 E (+1);
I1-, = Ø: I1-,+11,, = =I p2
k E (=1};
[Z 2
n2 s Vulogn, i (+1), {
(28)
Proof. For simplicity, we denote P(- ((r,,))) Ganta) as P(.) in the proof below.
For = [Pinits Vinit Vinie), we first show that for {(I, w) E Gasta.
(1,
(0)
2x) 0, Vi II
According to the definition of Gaatas || E 22 lognoonn for all i I- Thus
(0 (0) 21)
(1/ -x)
(0)
+ (1,)
zت 2,) > Vimi (2 ||| =+ll)) 0,
22

---

It follows that Published as a conference paper at ICLR 2025
where the first inequality uses 2 [0,0,2] when i Imp and the second inequality uses Assumption (A1). Similarly, we have
(o)), 21) < 0, Vi E I+#i-
Thus conditioning on {(z,, yolli= E Idata: we have
-1,1,11,--1 I mi I11111). = B.
For i I=ag. recall that 24 [giiinall o TTnnisse with Rigsignal
[ T
and ni,no
[10,3) Fi.pl T. We have
3
3 51)
3
(0) 21) = Vinit
Vinit
u/) un)T 1. = Vimit - T
11
Viiit
5t)
Im
1=1
P((0) 2.) > 0) 2
Applying Hoeffding's inequality, we obtain
P
I I+wal
> 1) - 2 exxp(-
(-)-
2t
2
Similarly we have
I I-wal
2t
PP( -u
1) S 2 2exp(
2
Let t = logn. We have
S nn loogg
J E 1=
2
with probability at least - 4/n?. Following similar discussion, we have that
11.j.-1)+tm I+mis Inj.-1,,-M 2, i.j (+1);
11,.-11) = Ø; - II-m i,j e (+1]:
1141,11,M,tt2 = Itrsi - = k (=1);
I1-1,41,6,+0 2: -1-,+11,),a =I- k (#1)
hold with probability given (1)) E Gdata- And
n2
S Vnllon, i, (+1), E {
2
hold with probability least 1 16/n2. In total, the conditions above hold with probability least 1 expo-S2(log' n)) 0 0(+)) (
Lemma A.2. Suppose that Assumption (A2) holds. Let the training duta { ys ys) for model Sr. be sampled ii.d from P. With probability at least 1 - exp(-S(log" 12)), we have
||
E,
P
logn, for all E [m].
(29)
n
Proof. Applying Lemma A.3, we obtain
P[ z,1l S 2 V
L Lign, Vi [n]) 1- 1- Σ P(I Zill > 22 V
L log
1=1
n exp (o n)) exp(- n)).
23

---

Published as a conference paper at ICLR 2025
Lemma A.3. Suppose that Assumption (A2) holds. For (2,,22,'" Exl ~ P with [x1, H. E (p2), we have
P([ Vllmax 82 P 2 logn) S eex(( <exp(-S(log'1)) n)),
(30)
where 112,-114 -2: M1] T ER.
Proof.
We start our analysis with [21, T M41 Note that
p-2
" x
8
1,44
a summation of independent bounded random variables with zero mean. By Hoeffding's 's inequality, we have
t
t
P(I x 2 S 2exp
S 2eex
2 2 87
2|101|382
Similarly, the concentration for uz and Ug + 2 are as follows:
P( x t) s exx
2110
12
P([ng x + 2 2 1) S 2exp
2 21
Combining these inequalities yields
PU x mmll > t) S 6
t²
S 6 exp
(-
202
nt²
(31)
where the last inequality uses Assumption (A2). The proof concludes by letting E² Vp/n log n. The analysis for other values of [21, x2] follows similarly.
Lemma A.4. Suppose that Assumption (A5) holds. Then the following conditions hold with probabil- ity at least O(1/n²):
(B1) maosk(Pos,Neg) Neg) ||T
log
(B2) maX=EEnnfom((+11)) (113)) ITe
m log
(B3) (uxx (=1)a) Jkk,e
<
m log "5"
16
(B4) (1 np.
13 <
12 log n
Proof. Note that [Peel Bin(m, 1/2). Applying Hoeffding's inequality, we have
P [Posl
m
<
m
< 2 exp
(-
2m
<
2 n
log
log
where the last inequality comes from Assumption (A5). And similarly
P([JTees
P
2m log"
< eexp
<
2
n
which completes the proof of (B1). Note that | ~ Bin(n, 1/4). Applying Hoeffding's inequality, we have
P
P(1
|
2 exp
log" n )
2n
0 1 - )
Ap (ttn+++2)
log
(B2)-(B3) can be proved following the same procedure. We omit the proof here.
24

---

Published as a conference paper at ICLR 2025
14
-12
15
Tannrruyyy
1--2 Ancuray
Toans
Figure 9: Training dynamics of the model Sr. discussed in Section 3.3.
A.2 ADDITIONAL EXPERIMENTS
A.3 EXPERIMENTAL DETAILS
All experiments in the paper can be run on a single NVIDIA A100 GPU. The loss function for modular arithmetic tasks is cross-entropy loss and for (40, parity task is logistic loss. All models used the paper, unless stated otherwise, set dmipp = 4dembed. where dmip is the MLP dimension and dembed is the embedding dimension. All FNN models used are in the paper homogencous and do not have bias terms. Code is available at https://github.co com thiveiix//rokkrrasseer
A.3.1 EXPERIMENTS IN SECTION 1 AND 2.1
two-layer FNN with trainable embedding layer as the weak model. We choose (dembed, width) (4, 16) for the weak model. The target model is a three-layer FNN with trainable embedding layer. We choose (dembed, width) = (128, 512) for the target model. The hyperparameters (init scale, learning rate, weight decay) are selected by the following grid search:
In Figure 1b,
init scale: (0.1,0.2) 1.5]
learning rate: [10 5 10 10 5 10 1000002000 10
weight decay: [10-1,10 10 2,-
We select the configuration that first achieves 90% accuracy on the validation set. The best configu-
ration for GrokTrans is (0.3, 0.005, 3). For standard training, only learning rate and weight decay tuned. They are selected by the following grid search:
learning rate: [10 3 5 x 10 10.2,102.0 10 10-2,11
1,,01x100011111 1,10'1
weight decay: [10- 10 1,2,3,,,,
and the optimal configuration is (0.05,3).
In Figure 2, we set the dimension of the GPT embedding to be 128. For the Fourier embedding. we choose & 7 frequencies, and let ij to be the j-th smallest prime number. For each type of embedding, we normalize the embedding of each integer to 1. The FNN used in Figure is three layer dense neural network
= (
where W1 xewied dim W2 E wwiddh W3 Pxxiddt width 512. The hyperparameters (init scale, learning rate, weight decay) are selected by the following grid search:
init scale: [0.1, 0.2. 1.5)
learning rate: [10-4 10 10-2 1001,111
weight decay: [10 4 10 10 1,,0,,220 1,0 1, 10).
We select the configuration that first achieves 90% accuracy on the validation set. The best configura-
tion (init, Ir, wd) for the four embeddings are:
One-hot: (0.2, 0.01,
Binary: (0.3, 0.01, 1);
Fourier: (0.5,0.1,0.1)
GPT: (1.3, 0.01,
25

---

Published as a conference paper at ICLR 2025
In Figure 3, the distance between two empirical NTK is estimated following the method in Mohamadi al. (2023). We denote O as the pseudo-NTK of the model epoch f, i.e.
(2( = oo
(11
Vo
10 (22)l/D
f=1
We estimate the distance between the empirical NTK at step and t - 1 by | θ
A.3.2 EXPERIMENTS IN SECTION 3
For experiments in Section 3, we let the sample size n 400, feature dimension = 80000, and noise level 0.05. For Figure 4(a), the model is two-layer neural network with width 2048. The optimizer is full-batch gradient descent with learning rate 0.1 and weight decay 0.1. In Figure 4(b), we train small model with only three neurons. The initialization of the hidden layer follows i.i.d (0, 0.01), ), and the initialization of the second layer follows i.i.d ((0, 10 10-1). The learning rate is 0.1 and weight decay is 0, 1. Figure 4(c) visualizes the hidden layer of that small model after training.
Figure 5a generates 4000 i.i.d datapoints from the distribution P. and visualizes Uz for each x. Figure 5b fixes n = 1000 and train the weak model for p = [4 x 10 x 10 16 x 10 ..e [1/40, 1 /800, 1/160). Figure Sc fixes p 8 10 and train the weak model for n [400, 800, 1600, 3200), [1/40, 1/80, 1 1/160). Figure takes Vini = 0.4, learning rate 2.0 and zero weight decay.
A.3.3 EXPERIMENTS IN SECTION 4
The attention layer used in this paper follows the same structure that in Nanda et al. (2023). While Nanda et al. (2023) also suggested to set the precision to be float 64 to mitigate the Slingshot phenomenon (Thilak et al., 2022), fluctuation of accuracy and loss during training process, we still use float32 to control the computation cost.
All modular tasks set p 113 and the fraction of training data being 25%.
For the (40, 3arrtiyty task, we set the sample size n 1000.
Unless otherwise specified, we use the A optimizer (Loshchilov & Hutter, 2019) for all experiments; we initialize the weights using the default PyTorch initialization scaled by a factor init scale > 0 to control the initial weight norm, as proposed by Liu et al. (2023).
The hyperparameters (init scale, learning rate, weight decay) are selected by the following grid search:
init scale: [0.05, 0.1, ... 1.5]
learning rate: [10-4 5
10 10 10 110 10 3,00,,00222000000000000000000300
100100011000 2,000 10.5, 1.0]
weight decay: 100 1,,,,,,,,4,
In Figure 6(a),bb), the structure weak and target model are the same as those in Figure 1b. In Figure 6(c), the weak model is a threc-layer width 16 FNN, the target model three-layer width 512 FNN with dembed = 128. In Figure 6(a), the optimal configuration for GrokTransfer fer is (0.3, 0.001, 1) and the optimal one for training from scratch is (0.1, 0.1, In Figure 6(b), the optimal configuration for GrokTrans fer is (0.3, 0.005, 3) and the optimal one for training from scratch is (0.1, 0.1, In Figure 6(b), we have the number of training samples n 1000. The model trained via GrokTrans Trans fer uses learning rate 10-3 and weight decay 10 10-3; the model trained from scratch uses learning rate 10-2 and weight decay 1.
In Figure 8(a), the weak model is two-layer width-4 FNN, and the target model is an 8-layer trans- former with dembed = 512, dmup = 512, Thhead = 4, dhead = 128. The optimal configuration for tar- get model trained via trained from scratch is (0.4, 0.0005, 1). In Figure 8(b), the weak model remains the same, and the tar- get model becomes 2-layer transformer with dembed 128, dmip 128, Mhead 4. dhead 32. The optimal configuration for target model trained via GrokTrans fer is (0.6, 0.005, 0.1). The optimal configuration for GrokFass st.it.i is (Ir, wd) (0.01, 1.0).
is (0.7, 0.001, 1). The optimal configuration for target model
26

---

Published as a conference paper at ICLR 2025
A.4 ADDITIONAL DISCUSSION
A.4.1 FORWARD PASS FLOPS ESTIMATION FOR MODELS IN FIGURE 8 LEFT
For the weak model, two-layer width-4 MLP:
forward 2 * (8 4 + 8 4 113 + 226) = 1436 10²,
For the target model, an 8-layer transformer, following Table in Kaplan al. (2020), we have 6291456, Cforward 2(N 10 Two-layer FNN takes around 2000 epochs to generalize. Target model trained by GrokTransfer takes around 1000 epochs and target model trained from scratch takes around 10000 epochs. Thus, the total flops of GrokTrans fer is around 00 + 2 10 and the total flops of training from scratch is around 10
N 2dlembedRMager (2dattn (2dattt (dmip) 2+51248+
A.4.2 ADDITIONAL EXPERIMENTS
030
0.30
A.0
0333
Figure 10: Left: Training dynamics of the one neuron weak model. Middle: Visualization of the neuron in the weak model. We can see it has learned the feature [1, 1]. Right: training dynamics of target model with embedding transferred from one-neuron weak model.
44
Cios
14
Epecho
Figure 11: Left: Visualization the distribution with the embedding from the one-neuron weak model. Middle: Visualization of the distribution P with the embedding from the two-neuron weak model. Right: Training dynamics of the target model with embedding transferred from a two-neuron weak model.
When the weak model only has one neuron, the target model has the form:
fx(x) = Σ agoli u x)x)
Denote uT.. Equivalently, we have
f((z) = Σ a,oly 2).
Note that sign(fL(2))) sign(fz(22) if sign(21) 22gn Since both positive and negative classes locate on both sides of the original point (Figure 11 Left), it is impossible for fr. to correctly classify these points (Figure 10 Right).
27

---

Published as a conference paper at ICLR 2025
20 PCA exmmed Embeiding
14M leamed Umbwidegs
dre
Figure 12: 2D Visualization of the 4D embedding from the weak model trained on modular addition task.
Tain (CrokTaansferi (orklaaaneee)
Train (4-Layeri Test (4-Layer) (2-Layer) (2-Layer)
Train (Large)
Test (Large)
Accuracy
0.8
0.7
0.4
0.6
0.2
0.5
0.0
10a
10'
10'
10"
10a
Epochs
Epochs
Figure 13: Left: Training dynamics of an 8-layer transformer modular multiplication task. Right: Training dynamics of a 2-layer transformer and a 4-layer transformer trained on the (40, parptr task. Both do not have delayed generalization.
100
Accuracy
Train [GrokTaanserr) (Grokiiarsffft
60-
[Largel
Test (Large)
40-
20
10
10'
30
10"
16'
Epochs
Figure 14: Training dynamics of a four-layer MLP on MNIST dataset. 200 samples are randomly selected the training data, and 400 samples randomly selected the test data. The weak model a two-layer MLP with 40% test accuracy.
28

---

Published as a conference paper at ICLR 2025
T.e-tDA Thansarpe
Thet (VLP)
6..
Teet (TF)
6.5-
7.5
xeep
Accuracy
6.5
Than Lare!
63
00
13
-27
12
32
131
Epocha
Epochs
Epocha
Figure 15: (a)Training dynamics of the target models used in Figure 6a and Figure 8 Right, but with low-rank embedding A B. both and B are randomly initialized. MLP represents a three-layer MLP and TF represents two-layer transformer. (b) Merge A and B into E, at 100-th epoch and keep training. (c) Set the embedding dimension the weak model to be the same as that of the target model.
1.0
Train (GrokTransfer)
Test (GrokTransfer)
0.8
Train (GrokFast)
Test (GrokFast)
Accuracy
0.6
0.0
10a
10
10
10a
Epochs
Figure 16: Dynamics of the target model (a two-layer Transformer) trained via fer,, and the target model trained via GrokFast on modular multiplication task.
29


Reviews:
[{"review_id": "UL6xEnEIJl", "reviewer": "Reviewer_GjdD", "summary": "The paper studies ways to reduce the \"delayed generalization\" in grokking, by first learning embeddings in a smaller, faster model where generalization isn't delayed, and then transferring the embeddings to a larger model. This is shown empirically in various examples including modular addition, multiplication and parity with fully-connected nets or transformers. The authors also show theoretically that fast generalization occurs in XOR when applying the GrokTransfer technique on top of a weak solution that was found with only 3 neurons (this solution is shown to be found empirically).", "questions": "* Could you comment on why you chose a weak model with 3 neurons in Section 3.2? Would the story change if you had 4 of them and achieve perfect accuracy? My impression is that the delayed generalization mostly happens when you are heavily over-parameterized and can very quickly memorize all the training data in early epochs, which still wouldn't happen with 4 neurons. Also, would the transfer still work if the weak model had fewer than 3 neurons?\n\n* For the modular addition/multiplication problems, what would be the smallest weak model that can lead to good transferrable embeddings (in practice, but also in theory with an appropriate construction?)\n\n* For Theorem 3.2, is this in a kernel regime? Would training only the second layer instead give similar guarantees? Regardless, a brief discussion after the statement would be a good addition.\n\nOther, minor:\n* Figure 4: Should I infer from the caption that the test accuracy becomes high around the same point (near 100 epochs), while the larger model also displays the property that training accuracy is high very quickly? Also, is it reasonable to say that the small model does not exhibit grokking here? Please clarify these points since they seem important for the story.\n* L 298: Do you mean that whenever the test accuracy is around 75% or above, you found empirically that the neurons consist of three such features? Please rephrase as this is not very clear.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Grokking is a surprising phenomenon that leads to seemingly unpredictable outcomes. The proposed GrokTransfer method seems like an empirically effective to mitigate this. The theoretical claims also justify why good embeddings obtained from a small model can help accelerate generalization in a larger model. Overall, this makes the work interesting and significant.", "weaknesses": "* The theory on XOR would benefit from a better understanding of the \"weaker model\", which is currently mostly empirical. It would great to have even partial results for this part, even in a simpler setting (e.g. training a single neuron on this data might be tractable and similar to previous work?)\n\n* The presentation could be improved in various ways, in particular the notion of \"small model\" considered isn't very precisely defined and would benefit from clarification. Other aspects that could be discussed further include trade-offs between width and iterations needed to generalize: it seems that smaller models reduced the \"delayed generalization\" but also slow down convergence in terms of epochs. Perhaps this changes when looking at \"compute/flops\", and it would be good to include such plots.\n\nSee also the questions below.", "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": "vXoPLJU8Ri", "reviewer": "Reviewer_g3Q3", "summary": "This paper proposes a method to accelerate grokking via embedding transfer from a week model. The authors first observe that data embedding plays a crucial role in determining whether generalization is delayed. Then they initialize the embedding weights extracted from  training a smaller and weaker model to the target, strong model. Finally, the authors give rigorous proof on a synthetic XOR task with simple network settings, and provide empirical studies showing the effectiveness of the proposed GrokTransfer method on fully-connected network and transformers.", "questions": "1) In section 3.2, could you add more details on getting the SNR with different embeddings?\n\n2) The intuition on why the data embedding  make a difference on the generalization speed is still not very clear. \n\nIf we take all layers beyond the last linear layer (which includes data embedding layer and some hidden layer) as doing feature engineering, as mentioned in [1], will this claim (in section 3.2, ``with this new embedding, data points are well separated in a three-dimensional space with a relatively high signal-to-noise ratio (SNR) compared to the original embedding'') still hold if we pass this data embedding into above hidden layers? Does the separation among the feature vectors got from the last hidden layer make more sense than the one among data embeddings?\n\n\n\n[1] Papyan, V., Han, X. Y., & Donoho, D. L. (2020). Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40), 24652-24663.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) This paper is clearly written and the idea is easy to follow.\n2) The choice of the task (XOR) and the setting of the model (very simple network) is easy to interpret for the purpose of the finding of this paper.", "weaknesses": "1) the key observation which is that the data embedding plays a crucial role in the generalization delay is not well studied. As the proposed accelerating method is based on this observation, it would be better to provide more evidence to say this observation is indeed the correct one. In addition to the modular addition task, providing more results on other tasks would make this observation more convincing. Also, instead of data embedding layer only, is the initialization of other layers a possible reason to cause the delay? An ablation study on other possible reasons to cause the delay would be better.\n\n2) In addition to simple settings, verification on more complicated tasks would make the finding more general and solid.\n\n3) In the theoretical analysis, the proof seems a little bit limited because the choice of very simple setting.\n\n4) It would be better if the authors make a clearer study on how to choose the \"weaker and smaller\" model to get the acceleration with good test performance.", "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": "IWnH2Lk06S", "reviewer": "Reviewer_opL6", "summary": "This paper investigates \"grokking,\" a phenomenon where neural networks initially struggle to generalize, only to suddenly achieve near-perfect generalization after extended training. To accelerate generalization, the authors propose GrokTransfer. GrokTransfer leverages the embedding from a smaller, preliminary model that achieves moderate test performance, transferring it to a stronger target model to speed up generalization. This approach successfully eliminates delayed generalization on a synthetic XOR task and performs well across various algorithmic tasks on both fully connected networks and Transformers. The results suggest that GrokTransfer could reshape training dynamics, enabling models to generalize without delay.", "questions": "- Would a “smoother” version of the proposed method work even better? For example, one could start from a single embedding layer with a small number of neurons and gradually deepen/widen it, i.e., adding more layers and adding more neurons to each layer.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The proposed method is interesting and novel to my knowledge. It first trains an under-parameterized model that is incapable of perfectly interpolating the data, and then uses the learned data embedding to facilitate the grokking of an over-parameterized model. It’s like approaching a problem by first crafting a simple but general solution that works well for most cases and then refining the solution to cover all the rest of the cases. Intuitively, this eliminates a lot of unnecessary competition among various not-so-simple solutions in an over-parameterized model, hence accelerating the learning process.\n- The paper is very well written. It includes an extensive discussion of the related work, a clear presentation of the motivation behind the proposed method, and a detailed case study that explains how and why the method works.", "weaknesses": "- It is unclear how well the method would generalize and scale to more complicated problems where such acceleration can make a real impact. The algorithmic problems are useful for analysis but are too simple in comparison with real-world problems which often involve high-dimensional inputs. For high-dimensional inputs with many redundant features, it is likely for the weaker model to lock onto degenerate solutions that would hinder the stronger model from grokking. It would be interesting to see if this is really the case or not.\n- The paper does not provide much insight on how to design or choose the weaker model. Clearly, the weaker model can neither be too weak nor too strong, i.e. there is a trade-off. If it is too weak, the solution may degenerate and thus make it harder for the stronger model to grok. If the weaker model is too strong, then it may take too much time to train the weaker model. Therefore, for this method to be truly useful, there should be some general rule of thumb for choosing the weaker model, otherwise, much time could be wasted on trial and error.\n- As the authors have mentioned in the paper, the theoretical result only considers a relatively simple XOR task. There does not seem to be any clear indication if the analysis could potentially be applied to more general problems. Therefore, the significance of this result is in doubt.", "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": "cT7Hhwvssu", "reviewer": "Reviewer_uoKT", "summary": "The paper introduces GrokTransfer, a method that expedites grokking by transferring embeddings from a weaker model. Through a simple XOR classification task, the authors offer both theoretical and empirical justification for GrokTransfer. Experiments on other algorithmic tasks show its effectiveness in embedding transfer from FNN to FNN/transformers.", "questions": "1. Regarding W1, the FNN->transformer setting was tested on only one task. I would like to see GrokTransfer's performance on more algorithmic tasks in this setting, especially in comparison to GrokFast. According to Table 1, training is very fast with GrokTransfer, so it should be quick to perform experiments on additional tasks (including those not tested in the paper). \n\n2. Grokking is known to occur beyond algorithmic data [1], which is already cited in the paper. I would like to see how GrokTransfer performs on non-algorithmic tasks, such as image classification with MNIST, as explored in [1].\n\n3. A new paper [2] has been recently released on accelerating grokking through weight transfer based on the Kolmogorov-Arnold (KA) representation theorem. Your approach seems simpler, but the approach in [2] can even handle two non-standard arithmetic tasks—composition of operations and systems of equations. How does GrokTransfer compare to [2], in terms of theoretical basis, performance, and practicality?\n\n[1] Ziming Liu, Eric J Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In The Eleventh International Conference on Learning Representations, 2023.\n\n[2] Yeachan Park, Minseok Kim, & Yeoneung Kim. Acceleration of Grokking in Learning Arithmetic Operations via Kolmogorov-Arnold Representation. arXiv preprint arXiv:2405.16658, 2024\n\nMinor\n- Typo at line 255: “need [to] learn”\n- There are reference mistakes in the paper, concerning arXiv citations that should point to conference proceedings. Here is an example of wrong references:\n\nTanishq Kumar, Blake Bordelon, Samuel J Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics. arXiv preprint arXiv:2310.06110, 2023.\n(Published in ICLR 2024)\n\nNote: the minor points did not influence my final score.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The authors provide a sensible justification for the core idea of their method through a motivating experiment in Section 2.1. \n2. They demonstrate its effectiveness both empirically and theoretically, highlighting improvements in computational costs and performance.\n3. The proposed method is simple and straightforward, and if it can generalize to broader tasks and architectures, it has significant potential.", "weaknesses": "1. The authors tested three algorithmic tasks—modular addition, modular multiplication, and the (40,3)-parity task—for FNN->FNN transfers. However, they provided results for only one task (modular addition) in the FNN->transformers setting. Given that FNN->transformer transfers are more practical and high in potential, there seems to be no reason to exclude modular multiplication and the (40,3)-parity task. Without these experiments, I can't help but think that the generalizability of GrokTransfer is limited.", "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": "RGYG3Hx69V", "reviewer": "Reviewer_s3cV", "summary": "The paper introduces a method to mitigate grokking, which describes the scenario in which models overfit to the training dataset long before it begins to generalize to unseen test data. This is done by training a smaller model with a much lower embedding (number of features in first layer) dimension, which converges quicker and without grokking. The architecture of the original model is then modified by replacing the first embedding layer with the product of two matrices $E_T = AB$ where $A$ and $B$ are of much lower rank. $A$ is initialized with the embedding layer of the smaller model. This model is evaluated on several synthetic tasks -- modular addition, modular multiplication, and a (40,3)-parity task defined as learning $y = \\Pi_{i \\in \\{1,2,3\\}} x_i$ where $x \\in \\{\\pm 1\\}^{40}$. Experiments show that compared to the original model that is randomly initialized, the resulting model mitigates grokking and converges faster. Theoretical analysis is also provided for a specific case involving learning $y = x_1$  XOR $x_2$ for a 80K dimensional vector $x$.", "questions": "- Do the conditions in the theoretical analysis, under which grokking is mitigated for the XOR task, have any implications on what are the causes of grokking? For instance, it possibly suggests several several causes of grokking, including the SNR, overparametrization (A5), initialization (A3), etc. \n- Are there failure cases of the method? I do not see mentions of them in the limitations.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "- The method is clearly described and easily understandable and implementable.\n- The combined time for the proposed method appears much faster than the baseline method of training the original target model from scratch. Table 1 suggests it is more than 5x faster. However ablations are needed to determine whether the speedup comes from architectural modifications, or from the proposed method.\n- For the tasks studied, the model performs well in terms of reducing grokking and accelerating convergence (although comparisons are not apples-to-apples -- see Weakness 1)", "weaknesses": "- My main concern is that all the numbers reported for GrokTransfer and the target model (Large) are not directly/fairly comparable. The model used for GrokTransfer parameterizes the first layer with $E_T = AB$ where $A$ and $B$ are (very) low-rank matrices. On the other hand, the original target model uses the full $d_v \\times d_T$ matrix which is significantly larger. The number of parameters present in this layer is now much larger than that of $E_T = AB$ . A fair comparison should not even consider target models with the full $d_v \\times d_T$ embedding layer, but only those parametrized with $AB$ , where, for instance, $A$ can be randomly initialized instead of being initialized from the weaker model. As such, it is difficult to conclude anything regarding whether the improvement comes from weak model initialization, or simply from architectural modifications, from any of the existing results.\n- This defeats the purpose of studying/mitigating grokking in the first place, where the goal is to be able to train over-parametrized models with minimal grokking. This method instead replaces the over-parametrized model with one of more manageable complexity, which ties to the next weakness\n- Method also greatly reduces the expressivity of the original target model, by adding a very low-rank bottleneck to the first layer. While this works for the synthetic closed-form tasks considered in the paper, it suggests that (1) the original target model architecture is clearly ill-suited, since it converges perfectly even after introducing a very low-rank bottleneck, and (2) ability of the method to generalize to more complex tasks which require greater expressivity is limited and questionable.\n- The general motivation of the method appears to be applying dimensionality reduction to pre-process inputs that have very low SNR. This opens up questions regarding how the method compares to classical techniques like LDA?\n- Method introduces an additional layer of complexity in the training pipeline, due to having to train an additional model which would require its own architectural and hyperparameter sweeps.\n- Minor comments on notation: Inconsistent notation switching between $d_{embed}$ and $d_W$  / $d_F$, which makes it hard to look up. Notation on line $y=x_1 x_2$ is also ambiguous, I assume the multiplication operation here is XOR. Also overloaded notation for $p$ for modulo and dimension of input in XOR task.", "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)