Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks
Authors: Tianyu He
Darshil Doshi
Aritra Das
Andrey Gromov
Venue: NeurIPS 2024 oral (2024)
Keywords: In-Context Learning
Grokking
Modular Arithmetic
Interpretability
Primary area: interpretability_and_explainability

Abstract: Large language models can solve tasks that were not present in the training set. This capability is believed to be due to in-context learning and skill composition. In this work, we study the emergence of in-context learning and skill composition in a collection of modular arithmetic tasks. Specifically, we consider a finite collection of linear modular functions $z = a  x + b  y \text{ mod } p$ labeled by the vector $(a, b) \in \mathbb{Z}_p^2$. We use some of these tasks for pre-training and the rest for out-of-distribution testing. We empirically show that a GPT-style transformer exhibits a transition from in-distribution to out-of-distribution generalization as the number of pre-training tasks increases. We find that the smallest model capable of out-of-distribution generalization requires two transformer blocks, while for deeper models, the out-of-distribution generalization phase is *transient*, necessitating early stopping. Finally, we perform an interpretability study of the pre-trained models, revealing highly structured representations in both attention heads and MLPs; and discuss the learned algorithms. Notably, we find an algorithmic shift in deeper models, as we go from few to many in-context examples.

Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks
Tianyu He a.t (tianyuh, ddoshi, aritrad] @umd. edu
Darshil Doshi
Aritra Das
Andrey Gromov a.a
gromovand@meta.co com
Abstract
Large language models can solve tasks that were not present the training set. This capability is believed to be due to in-context learning and skill composition. In this work, we study the emergence of in-context learning and skill composition in collection of modular arithmetic tasks. Specifically, we consider a finite collection of linear modular functions
a x + by mod p labeled by the vector (a, b) E Zz We use some of these tasks for pre-training and the rest for out-of- distribution testing. We empirically show that a GPT-style transformer exhibits a transition from in-distribution to out-of out-of-distribution generalization as the number of pre-training tasks increases. We find that the smallest model capable of out- of-distribution generalization requires two transformer blocks, while for deeper models, the out-of-distribution generalization phase is transient, necessitating early stopping. Finally, perform an interpretability study of the pre-trained models, revealing highly structured representations both attention heads and MLPs; and discuss the learned algorithms. Notably, find algorithmic shift in deeper models, as we go from few to many in in-context examples.
Introduction
Large language models (LLMs) can perform simple tasks absent from their training data. This ability usually achieved via in-context learning [5). More importantly, LLMs can perform an even larger variety of very complex tasks upon appropriate prompting or fine-tuning. The latter ability to perform complex tasks usually attributed to the following mechanism. First, LLMs learn large variety of simple tasks and, then, how to compose those skills to form very complex skills 3 Furthermore, LLMS also exhibit "emergent capabilities" of scale (either model size, compute or data) [29] 9. It is plausible that these sudden performance improvements are due to one or both of these mechanisms. For example, LLMs show grokking algorithmic tasks 244,, which results from the model learning very structured representations Once these representations emerge, the model abruptly learns how to perform the task.
sudden emergence
new complex skill as function
18
20
In this work, we set out to examine skill composition both empirically and mechanistically. Inspired by the prior work that investigated emergence of in-context learning on linear regression tasks I 25). we introduce finite collection of discrete modular arithmetic tasks [24] generalized to the in-context learning setting. Each task corresponds to learning linear function z = a J + by mod over Z.p from the examples provided in context of the autoregressive (Adel(AM). In the bi-variate case there are p² such functions labeled by the vector (a, b). The main objective this algorithmic datasct is to probe how AM utilizes the tasks it has learnt during training to solve the new tasks.
'Department of Physics, University of Maryland, College Park
'Meta AI
Conresponding author
'Souree Code: http::////ggthhhhcc//////aaaaaaa hrtp://aaaaaaaoooooouaraaatttaer
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

DOGG
(c)
task (1.1)
task (2.3)
2
4
12
(e)
14
Figure 1: (a) The dataset. The tasks are labeled by vectors (a, b) E Z Each table contains examples of a. by mod p. A fraction 1 of the examples is blacked out; while the remaining examples are flattened into a single "document" in the batch. Each document is organized as a collection of triples (Ir, y, oz by) for æs, from the training set (i.e. not blacked out in the table). Our training is similar to the traditional next-token prediction (autoregressive): with the main difference that we predict every third token, which are marked in red (æ and are uncorrelated). Every task appears exactly the same number of times in each batch. (b) Phase diagram for six-layer model. We find four different phases. (I) in-distribution memorization: The model only performs well on tasks (a, b) and examples (x, y) from the training set - it does not generalize on unseen examples or tasks. (2) in-distribution generalization: model generalizes on unseen examples (Ir y) but not on unseen tasks (a, b). (3) out-of distriiitibttonn memorization: model generalizes on unseen tasks (a, b) but only for examples (1r, y) has seen during training. (4) ant-of-diisributtinn generalization: model generalizes on unseen tasks (a, b) for seen as well as unseen examples (x, v). We focus on investigating phase (4) more detail. (c) In-context sample complexity. Accuracy of the model in phase (4) as a function of the number of few-shot examples. (d) Representations developed by one of the attention heads in the first layer. These are projections of the embedding of a pair of numbers onto the two largest principal components (PCs) of the internal representation formed after passing through the attention layer and projection matrix. (e) First 3 PCs embeddings separate (og27-annotan numbers into even/odd planes, with 0 sandwiched between them.
Our analysis shows that the solution found by the AM after optimization is qualitatively different from the linear regression cases studied before [1] In those cases, due to the continuous nature of the task, AM develops an emergent first-order optimization method that minimizes an emergent quadratic loss function. Furthermore, it shown in M. single linear attention layer solve the regression problem, while adding extra layers and non-linearities slightly modifics the gradient descent. In the modular arithmetic case, AM first learns how to solve the pre-training tasks and later (assuming enough different tasks) develops a generalizing solution by combining the solved tasks.
Our main findings as well as the structure of the algorithmic dataset are illustrated on Fig. I Our main findings are: (i) there are four different phases in the end of pre-training depending on the number of tasks, ni.d. and number of examples per task, (ii) At inference time, there is generalization transition the number of few-shot examples, as number of examples grows, the models starts to generalize. This effect is somewhat similar to the transition in sample complexity for the modular arithmetic found in 24 (iii) model develops a striking circular representation for all of the tasks that naturally generalizes the circular representations found in the original work (24). We further find that the deeper models are easier to optimize, but much harder to interpret. The optimization is discussed in more detail in the main text. Here we will highlight that optimization for these tasks is challenging and the AM tends to prefer the minimum that just solve handful of tasks and memorize the training set. To avoid such minima make that every batch contains equal number of tasks (meaning that no tasks is over- or reder--epeeeented each batch). We further find that for larger models early stopping is necessary because the generalizing solution is transient.
We organize our paper as follows. Section notations and discuss the experimental details. In Section 4 we demonstrate empirically that the
2 contains the literature review. In Section we explain

---

out-of offoofsttributiin ICL ability emerges as the number of training tasks increases. We also study the effects of model depth and task difficulty. In Section we carefully examine minimal setting, i.e. two-block transformer: we compare the representations learnt in four different phases and show that in the generalizing phase the representations are highly structured and generalize the original modular addition case of [24).
2 Related Works
In-Context Learning (ICL) Brown et al. [ first demonstrated that large models performance improves substantially when a few examples of the task at hand are provided inference time, in the prompt. Akyürek et al. (2). Ahn et al. [I, von Oswald et al. (28) showed that the AM implements emergent first-order optimization on an emergent objective function to solve linear regression tasks. Furthermore, [2] showed that larger models learn to perform Bayesian estimation. Garg et al. [IO] demonstrated that transformers can learn several simple classes of functions in context. Kirsch al. [15] presented how task diversity and model size would affect the ICL performance for unseen tasks using a mixture of modified MNIST datasets. Raventos et al. [25] investigated the relation between task diversity and uuttoo-diistrbuttio ICL ability on linear regression tasks. Lin and Lee [[6] identified two operating modes in ICL using a mixture of linear regression tasks, where for the first several shots, the model tries to figure out the correct task vector and later uses it predict the correct results. Boix-Adserà et al. [4] showed theoretically and experimentally that with enough wre-training data, transformer model can perform abstract reasoning that MLP cannot do. Guo et al. [I3 showed that transformers can use lower layers to memorize and upper layers to perform ICL in feature regression setting. It was found in 26 that ICL is transient phase from the optimization point of view: it goes away once the model is over-trained. Hendel et al. 14]. Liu et al. [17 showed that language models form in-context vectors, which can be used to steer model predictions.
Modular Arithmetic Power et al. [2 discovered Grokking, where models trained on modular
arithmetic datasets have an abrupt change from random guessing to generalization on the test set way after the model memorized the training set. Gromov [11], Nanda et al. (20), Gu et al. [22 showed that for modular addition tasks, models learned to map integers to Fourier features to solve modular arithmetic tasks. Liu et al. [18) showed that grokking is related to learning highly structural features, and the grokking transition can be explained by a toy model. Zhong et al. [30] showed that there is more than one algorithm that a model can implement to solve modular addition. Doshi et al. [6] showed that corruption of the label does not prevent the models from finding a generalizing solution. Doshi al. (7) showed that MLP and transformer models can solve specific family of modular polynomial tasks by bijectively mapping them to modular addition tasks.
Interpretability Elhage et al. [8] Olsson al. [22] showed that transformers can form induction heads that predict the next token in a sequence by identifying and copying patterns from earlier in the sequence. With several indirect empirical evidences, they showed that those heads might constitute the core mechanism of ICL. Nichani et al. [2I] showed theoretically and empirically how disentangled transformers learn causal structures from in-context Markov chains by forming induction heads.
Preliminaries
Linear Modular Functions We consider modular arithmetic tasks of the form: = al B'y mod p. We will refer to the coefficients (a²,8) E Z as the task vector. The superscript t e (1,... , labels the p? possible tasks. We will refer to (It, U:) ) E ZZ as the input vector, which is labeled by the subscript
In-Context Learning with Transformers We use GPT-like transformers (5] with ReLU activation function and Rotary Positional Embedding (RoPE) (27. The model has d consecutive blocks, H/ attention-heads, and embedding dimension dembed. Each number is tokenized as an independent token. The pre-training done following slightly modified next-token prediction setup, with sequences of the form:
s¹ = (I1
1/1
T2
1/
2/
Fnax
Inax
=n E Ep
3xx
(I)
3

---

where nexx is the maximum number of in- nn-context examples. The model is asked to predict only the labels {zf,- z }. We emphasize that we do not explicitly provide the task vectors (a' b³) the model (see Fig. this information implicit the in-context labels {a{}- In order for the model to generalize, must determine the underlying task vector (a', b²) from the few-shot examples.
Generalization There are two notions of generaliza- tion in this setup. (i) In-distribution: Generalization to wnseen input vectors (z,, y:). but on task vector (a',b") the model has seen during pre-training. (ii) Out-of-distribution: Generalization to task vectors the model has not seen during pre-training. To clearly separate these regimes, we split the task vectors into in-distribution (i.d.) set Tia == ((a'bb)); (b) and out- of-distribution (o.o.d.) set To.o.. == {(a', (o,,6))o.o. Jo..d.. Similarly, we split the input vectors into train and test sets: Muain = {(rz ya) } Jrsin: Peatt == {(2,, y4 (:)]tess This results in four distinct sets of sequences constructed from those sets; we name them Strinn Sl.d. .0..o.d and So.o.d. The set St used for pre-training, while the other three sets are used for evaluations.
performs well
X
performs poorly
Strain
Syrurs
In-disributtinn
memorization
In-distribution
gencralization
Out-of-dismibution
memorizatioo
Out-of-distriuuion
generalization
Pre-Traie Pre-Training Task Selection and Sequence Design We always sample the pre-training task vectors T.a. in sets of 4, following the rectangular rule, shown in Figure ([a). Additionally, each batch contains an equal representation from all the task vectors in the T.a Moreover, all the tasks share the same sequence of inputs. For example, batch with two different task vectors (a"1 oft,,0 ,,1);;; (a²2 bta) and two distinct input sequences per task (resulting four total sequences) is shown Figure bb)
(a)
(b)
Pre-training Sequence Design:
Pre-training Task Selection
Jids
Vin
Fig
1/1
Rectagular [Schomatic]
Mis
Fig
Vig
95i
231
Jija
Figure 2: Structured selection of pre-training tasks and sequences.
This structured approach creates a coherent signal from the sequences within each batch: ensuring that the model learns multiple task vectors with reasonable batch sizes. Alternatively, if the batches are sampled i.i.d., then the model is confused by the batch noise and cannot learn any tasks.
Detailed discussions task selection and sequence design are presented in Appendix D
Default Setting Unless stated explicitly, we will use p = 29, the number of heads H = 4, and embedding dimension dembed 512, with netx 32 in-context examples. All models are trained with AdamW optimizer [19] with batch size B = 1024 for 200k steps. We have also tied the embedding layer of the model with the readout layer.
4 Emergence of In-Context Learning and Task Composition
In this section, we demonstrate that a transformer model with depth 2 2 can develop ICL and out-of-distribution generalization modular arithmetic tasks. We delve deeper into the two notions of generalization (i.d. and o.o.d.), and discuss the relevant factors. We find that the model's ability to generalize -itt----iittiittionn is predominantly determined by the number of pretraining tasks nid.-
4.1 Transition driven by the number of tasks
In Figure a), we show the accuracy of d models vs the number of training tasks ni.a. and the number of Tew-shot examples quantified by the fraction of the total number of data points, at on sets Si.d. St.d, Stood and S o..o.d. The phase diagram in Figure
constructed by merging the last
shot accuracy version of these four diagrams shown in Figure 27 (a) of Appendix G

---

train
1. test
o.o.d train
0.0 lest
0.8
100
0.1
89
05
prediiiinn inputs
89
6A
02-
25
20
2 2
2
2
2 27
2
2
2
22
# of pre-traiiing taaks (A)
of pro-rakking laoks (n)
# pre-traiing lasks (n.
a of pre-training banks (%) d
(a) Accuracy phase diagrams of all four sets
8
GOOUISCY 19
24
93
12.
1508
1000
28 rumber of chote
number shots
(b) Loss step
(c) Accuracy step
(d) Loss of shots
(e) Accuracy of shots
Figure 3: Phase diagram for the depth d = 6 models. (a) Accuracy all four sets used plot the phase diagram, with an early stopping applied. Notably, in the regions when models generalize to o.o.d. sets, the pre-training performance degrades; (b, c) 0.6 training accuracy and o.o.d. test accuracy (dotted line). For n4.. = 2 we notice that the o.o.d. generalization ability of the model first improves then degrades as train longer; (d, e) 0.6, loss and accuracy context length, measured on So.o.d. at the end of training, where for A..d. 2 case the ICL ability fades away.
The ability of the model to generalize in-distribution increases with as can be seen by comparing the first two panels of Figure Bta. This behavior in correspondence with the original work on grokking, where the transition generalizing solution is driven by the amount of data. Further, we observe that an increase in ni.d. enhances the in-context sample efficiency, i.e. the model generalizes at inference time with fewer few-shot examples. This indicates the onset of the transition from mask-memriiiing solution to one that generalizes nut-of-distribution. The model switches to a new algorithmic way of solving the task and the solution is more few-shot-sample-effiiie
Shifting our focus the last two panels of Figure B)), a), can solve new tasks that were absent in the training set. Notably, there appears to be trade-off between memorization and generalization when the model attains this o.o.d. generalization ability. As the o.o.d. performance increases, the pre- trainiiinng performance simultaneously degrades. This phenomenon indicates shift the algorithm implemented by the model. Prior this transition, the model primarily needed to select possible vectors from the list of memorized tasks and apply them. However, post-transition, the model adopts more universal approach to solve the task in-context. We emphasize, that the model learns to perform ICI. in both scenarios. The difference lies in the approach to generalization. When the model can only generalize in-distribution it's task is to classify the sequence one of the seen tasks memorized task vectors, it does well for x, y pairs that only appear in the test set. However, as the number of tasks vectors grows the model fails to store them all and is forced to find method of determining the task vector algorithmically at inference time. In that case model performs equally well seen and un-seen tasks alike. In fact, the small two-layer model study has such a low capacity that it entirely skips the in-distribution generalization phase and immediately jumps from pure memorization to distriiuuui generalization.
see that when ni.d. 2 256, the model
unknown. Once it matches the sequence to of the
Next, to further illustrate the effect of task diversity, we plot the pre-training accuracy (set Stdd.) and the o.o.d. test accuracy (set So.a.d)) as function of training steps (Figure Bbb c)); for various values of nt.d.- We observe a clear memorization-te tirrrizatiio-t-geeenalizzioooono transition as task diversity increases. Interestingly, for fi.d. = 28, the ICL ability on S0.9. set exhibits non-monotonic behavior, where the o.o.d. performance rises and falls along the training process. This phenomenon is likely due to a competition between the memorizing and generalizing circuits inside the model. Note that this phenomenon is akin to the analyzed in Singh et al. (26). albeit with different setup.
5

---

id train
0.O.d. test
o.o.d. test
130
pretrainn fncttn (o)
44
3.0
6--
10
15
20
25
bodos
number of shots
(a) d
o.o.d.
0.0.d.
pe teling lrees oo
03
6.
22 2.0
21
10
8.d1pro toining 1
waining a
(b) d=
Figure 4: Phases depth d = 4 and d = models. With decreasing model capacity, the performance on both sets degrades. At the same time, the transient nature of ICL does not appear. (a, b) from left to right: accuracy phase diagrams on pre-training set Si.d. and on o.o.d. test set go.o.d. with early stopping; loss and accuracy vs context length on o.o.d. test set So.o.d for a 0.6.
test
Further evidence supporting the two-circuit competition can be observed in panel (d) of Figure The loss curves show "monotonic -> non-monotonic - monotonic" transition as the task diversity increases. With minimal number of pre-training tasks, the model primarily engages in memorization, resulting a monotonically increasing o.o.d. loss curve. As the number of pre-training tasks increases, the loss curve exhibits non-monotonic behavior, indicating competition between two distinct neural circuits. This transient nature of o.o.d, generalization for 7i.d. 2 is a peculiar case where memorization circuits are initially suppressed but eventually prevail. With substantial task diversity, the circuits responsible for generalization take over, culminating in monotonic loss curve. Similar insights be derived from examining the monotonicity of the accuracy curves in panel (e).
4.2 Effect of Model Size and Task Difficulty
A natural question to ask is if similar phenomena be observed with different model sizes task difficulties. Here we present our results with and in Figure and leave the results for other prime p values in Appendix
When comparing phase diagrams in Figure with Figure we observe that those phase diagrams across different depths are qualitatively similar, where the o.o.d, generalization only emerges with large enough number pre-training tasks. As model capacity decreases, performance on both the pre-training set and the o.o.d. test set degrades. This is particularly evident in the d = 2 case, where pre-training accuracy falls drastically as model gains o.o.d. generalization ability.
Interesting observations can be made by comparing loss and accuracy on the o.o.d. test set as function of context length the end of training. First, it is evident that as the model depth decreases, the 1-shot loss surge attributable to memorization becomes milder. Notably, for models with n.4. 29, there is loss surge in the 1-shot case across all three depths. Furthermore, the d = 4 model with ni.d. = 2 behaves significantly differently from the corresponding one with d = case, where the model fails to perform ICL for the o.o.d. test set. This is also distinct from the d case, where the model tends heavily toward memorization due to its excessive capacity. Instead, the model manages to maintain a better balance between memorization and generalization end of pre-training. Consequently, the model has 1-shot loss surge followed by notable drop in ICL loss. This suggests that d optimally leverages the available model capacity facilitate effective learning dynamics for o.o.d. generalization.

---

depth = model
depth model
1-shot
13-chot
31ishor
BOAL sechor predictions
Mark corrcctff Xf alas Wt holds for any i.
connect
now 3: Subtrecioaar from pow2
Figure 5: d 4 and d 2 models' performance on k-shot inference, on the grid of inputs (xe, v) E 23 (task vector = (6, 6)). row 1: Models' predictions on o.o.d. task of the type (x1 91 21 based on Modular Regression algorithm. row 3: Subtract row from row I, by using correct=1 and incorrect=0. The red points correspond to the examples where Ratio Matching does not give the correct predictions but the model predicts correctly. The blue points are examples that the model missed despite Ratio Matching being applicable. This row tells us about the model's ability to implement Modular Regression by combining the in-context examples. Note that d model readily learns to combine previous examples, while its d = 2 counterpart struggles due to its limited capacity.
x ?). row 2: Analytical plots showing predictions solely
5 Interpreting the Learned Algorithms
We now explain the algorithms implemented by the models to achieve o.o.d generalization. We find that the models implement two distinct algorithms depending on the depth and the number of in-context examples during inference.
Ratio Matching: If there exists s.t. cI,,, 34) (x, y) mod p then qæt mod p. This algorithm only works when y/r = y//x; mod p holds for at least one of the in-context examples.
Modular Regression: Find {ci, (e,,c,,..., } s.t. Σi= (x vi) = (æ, v) mod p. Then the predic- tion is 2 = Σ C222 mod p. This can be viewed as discretized circular regression over GF(p).
We find telling signatures of these algorithms upon analyzing model predictions with varying numbers of in-context examples (Figure 5). With very few in-context examples, the models implement the Ratio Matching algorithm. As a canonical example, consider the 1-shot scenario (æi V1 21 x y ?). In this case, Ratio Matching will only solve the task for inputs that obey (Ir y) e1 c1(11, y1) mod p for some e₁. Indeed, this exactly what we observe in Figure for both d = 2, 4 models. With many (~ 10) in-context examples, d = 2 and d 4 models implement different algorithms. The d = 4 model can combine the in-context examples using the Modular Regression algorithm, leading to near-perfect o.o.d. generalization. On the other hand, the d = 2 model still uses Ratio Matching and shows sub-optimal performance. We ascribe this to the limited capacity of the d = 2 models. In other words, the d = 4 models exhibits algorithmic shift from Ratio Matching Modular Regression as them number of in-context examples increases. Whereas, the d 2 models shows no such shift.
To implement these algorithms, the model needs to perform linear operations over GF()), which can be broken down into the following essential skills (in addition to copying information, which is readily implemented by attention heads).
I. Modular Map: Encode the tokens such that operations over GF(p) can be easily implemented II. Multiplication: A necessary skill rescale in-context examples in both algorithms
III. Addition: A necessary skill for combining in-context examples in Modular Regression algorithm
While both d = 2, 4 models learn skills I, Il perfectly; the d = 4 model outperforms its d = counterpart in skill III which helps it combining the in-context examples.
In the remainder of this section, we show the special attention heads that implement skill (Section 5. and MLPs that implement skills II. III (Section 5.2) We explicitly show the deterioration in this

---

Dua = 2 layer 1.head!
Due =
layer 1, head
layer 2, head
layer 1, head
Quary
Query
Quary
Quay
Key
Kay
pilee coamp--nue
0.00
-00--d ale sle shi prineppal component
princcpplllcccccc component
prinnccipal component
principaal component
(a)
(b)
(0)
(d)
Figure Models that generalize o.o.d. (left) exhibit more structured attention maps. Additionally, the top-2 principal components of the features from the corresponding heads also show more structured patterns. The features are computed for sequences with an 0.o.d. task vector (at, b') = (6, 6), loop fixed. We annotate each
over (Iz, 3/) at specific shot while the previous parts of the sequence (a) The principal components form a circle of circles where the position of the outer circle is controlled by Z;. This pattern remains the same for different task vectors or the shot choices; (b) Only plotted pairs with even log27 Tú, with each log27 24 circled with different colors. The PCA pattern forms similar double circle those in (a), with the key difference that those circles depend on task vector choices and the shot choices; (c, d) Models without o.o.d. generalization ability. We pick heads from the first block that corresponds to the first column of (a). Clearly, the structure of attention maps and PCA patterns deteriorate as the task diversity decreases.
PCA plot with the (logg7 24, log27 y4) mod p pairs.
structures as pre-training task diversity (nt.a.) decreases. We also elucidate the algorithmic shift in the d = 4 models as the number of in-context examples increases.
5.1 Attention Heads Implement Skill
So far we have presented "black-box" evidence suggesting that the model is implementing the pro- posed algorithms (Figure 5). Now, we turn to wwhite-box' interpretability, to identify the components within the transformer model that responsible for the essential skills. In Figure a. b), we analyze the important attention heads in the d = model. Specifically, we compare the emergent structures in models trained with different pre-training task diversities.
In Figure 6 a)), we show the attention head from layer that implements skill I. In the top panel, we see that each query only pays attention to itself and the two preceding keys. This pattern likely stems from the fact that each example in the sequence contains three tokens 24, yi, z/; and suggests that this head is mostly focused on the information within the example.
In the bottom panel, we perform principal component analysis (PCA) on the outputs of this head. Specifically, we feed batched k-shot sequence of the form (x1 V1 21 where the first k inputs are fixed and the last input (Ir y) is scanned over all p² possible pairs. We concatenate the resulting features from æ and y, resulting in p² x 2dambed batch of features - and perform PCA on this matrix. We project all these 2dembed dimensional features onto the first two principal components. Annotating the pairs (z. y) with (log27 v. log27 y we find "circle-of- circles" where circles of period 28 are arranged in bigger circle of period 28. Number 0 located
26 yk 2/
2).
We use 27 the base of logarithm, which is primitive of GF(29)
8

---

depth = 4 model
depth model
layer 3, 2-shot
layer 3, 5-shot
layer 3, 16-shot
layer 1, shot
wan
1.00
waun
widgy
Sprr po
0.78
wol
Sold
0.50
modu
Red
025
yox mod
y0a yia mod
702 m you mod P
y/x 4xk-yye mad
0.00
layer 4, 2-shot
layer 4, shohot
layer 4, 16-shot
layer 2. 2-shot
g
0.50
undky)
0.75
redal
andy
ye you, y1t mo%.mmdd
by ox mod (am2 bvd)
yicmyle med p:: vls yoh 2/k'=y/t, mod pt y/M 110-ayhe yyht
Vic Wob mad p: yix -y/.mmdp
Ratio Matching
Modular Regression
Ratio Matching
Figure 7: Cosine-similarities (Equation between layer outputs h' at token 2 position (cos(h', (x,, y), h2 (a', s/),, first row) and token y position (cos(h', (2, v), h², (3',, v')), second row) reveal patterns in the models' internal representations. For clarity, we only show selected x and x' values, where y and y values range from 0 28 between each tick. For the d 4 model, kaleidoscopic patterns in the third layer indicate the generation of all possible yi /i/x; features for subsequent computations. The last layer shows algorithmic shift from Ratio Matching Modular Regression. The d = 2 model also shows the kaleidoscopic pattern in the first layer, while the second layer identifies the relevant y/ features from the in-context examples for matching.
at the center of the circles We observe similar circle-of-circles for concatenated features from (a, and (y. 2) as well. We refer the reader to Appendix F for further details.
In Figure b)), analyze the head from layer 2 that effectively "copies" in-context examples. The upper panel shows the highly structured attention map that focuses on the current as well as the preceding examples. This pattern aligns with the Ratio Matching algorithm where the model compares (z, y) pairs across different in-context examples.
In the bottom panel, by conducting a PCA analysis, we again identify circles when annotating examples in the (log27 x, loggr y) format. However, unlike the previously discussed pattern. the specifics of this "circle-of-circll vector (a*, 1.). This variability suggests that the head in question possesses information about the specific task from the context.
arrangement vary depending on the position and the choice of task
We further highlight the importance of the structure we find in the above paragraphs via comparison with models, trained with lower pre-training task diversity, that do not generalize o.o.d (Figure 6 c, d). Note that as the number of pre-training tasks ni.a. decreases, the attention map starts to get mosaicked (top panels); and the PCA projections lose their shape (bottom panels). As result, these models lose the ability to perform ICL on modular arithmetic out-of-distribution.
5.2 MLPS Implement Skill II and III
After applying Skill I, the model maps numbers onto circular feature spaces, enabling discrete modular operations such as multiplication, division, addition, and subtraction through subse- quent model components. To analyze this behavior, compute the cosine-similarity between layer / (i.c. tth Transformer block) k-shoo output vectors (standardized) at token position y: hi (I1, Vi, 21, ". Tk, yks 2k, Z. y).
Thii is expected since logzr is mathematically ill-defined and needs special treatment.

---

kh Mk Fk, v) (.,.... , 3k. Mk: 2ks x., .2/ 1 (.... Jee yk FR, x. v) hl (..., UR yk RA, x², ./1
cos(hi (i (a, w), 1 g 22))
depth
depth =
We evaluate this similarity metric across all possible input pairs (I, y) and (r',8'), resulting in a p² x p² matrix. We also repeat the analysis to layer output vectors 2 tokens h-)(). In line with our previous analysis, we use controlled (fixed) in-context examples (x1 y1 21 ... 2.k yh =x.-
kyer2
layer
3
layor
For the d panel of Figure 7 reveals a distinctive 36 kaleidoscopic pattent of high cosine- similarities. The pattern corresponds to imput pairs (xx, y) that share the same ra- tio y mod In the MLP modules, we find highly structured activations as functions of inputs (x. y) (Figure left). Together with our earlier analysis of Skill I, these observations suggest that the MLP layer following the attention module with "circle-of-cireles" head (layer head for ments division operations over (p) The bottom panel of Figure 7 succinetly demonstrates the algorithmic shift from Ratio Matching to Modular Regression as the number of in-context examples increases. At 2-shot, layer 4 features collapse to identical vectors except when the ratio y/z matches one of the ratios (3), /xz or y2/ /2). At higher shots, we see a transition to a task-dependent pattern where features align for input pairs (a, y) for which by mod p match signature of Modular Regression.
model, the top
Figure 8: MLP activations Reo)t wrt inputs (x. v)
model, see Figure 13 (a)) imple-
For d = 2 models, we again observe the kaleidoscopic pattern in sisine-iimilarities (Figure topp right) and structured MI.P activations (Figure X right). However, in this case we only find signatures of Ratio Matching in the cosine similarities (Figure bogtto-rotto as expected.
Thus, we conclude that the d model has scarcely acquired skill III, due to its limited capacity. On the other hand, d = 4 model is very good at combining equations via skill III, explaining its superior performance. For more detailed discussion, see Appendix
Discussion
We have investigated the emergence of in-context learning and skill composition in autoregressive models on a novel algorithmic dataset. The dataset includes large discrete set of modular arithmetic tasks and is specifically designed to force models to learn how to solve variety of tasks. It consists of learning linear modular functions, where the model is expected to identify and perform a modular operation in-context. When the number of tasks is relatively small the models can only generalize in-distribution. Although such models develop ICL capabilities, they simply memorize these task vectors and use them to classify the input vectors. Once the number of training tasks becomes too large, the models transition to a gualitatively different algorithmic approach, where the task vector is determined at inference time.
Finally, we have examined the learnt representations and shown that qualitatively different circuits are formed in different phases. In the o.o.d. generalization phase, we explain the learnt algorithms and highlight an in-context algorithmic shift deeper models.
Limitations We have limited this work to algorithmic datasets. It remains to be investigated what lessons can be translated to realistic language models, and what lessons are specific to the current setting. White-box interpretability analysis of deeper models has proved to be much more difficult than that of shallower models. Consequently, we still do not understand the role of every individual component of the network in the deeper cases.
Patterms cropped due to file size limit, check the GitHub repo for the full kaleidoscopic pattern.
10

---

Acknowledgments
T.H. thanks Yue Xu and Dayal Singh Kalra for helpful discussions. A.G.'s work at the University of Maryland was supported in part by NSF CAREER Award DMR-2045181, Sloan Foundation. The authors acknowledge the University of Maryland supercomputing resources (http://hpcc. und edu) made available for conducting the research reported in this paper.
References
[1] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand. and Suvrit Sra. Transformers learn to implement precon- ditioned gradient descent for in-context learning. In A. Oh, T. Neumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 45614- 14- 45650. Curran Associates, Inc., 2023,
[2] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linearmodels, 2023. URL SRRLltttss//ppeeeeviiee http:/pprrreiiiieeeerroerr Id=0gOX4HBy84I
[3] Sanjeev Arora and Anirudh Goyal. theory for emergence of complex skills language models. arXiv preprint or Xiv: 15936, 2023.
Enric Boix-Adserå, Omid Saremi, Emmanuel Abbe, Samy Bengio, Etai Littwin and Joshua M. Susskind. When can transformers reason with abstract symbols? In The Twelfth International Conference on Learning Representations, 2024. URL http::::/ppppeer///////////
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Ncelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan. Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877 1901 Curran As- sociates, Inc.. 2020. URL https: p/poceeddnns neur wttpsscccppppprppppppppppp www.fflllllpppep.eeeeeeee file/ .dff
1457c0d6bf 00w.6668800009777773373861144464
[6] Darshil Doshi, Aritra Das, Tianyu He, and Andrey Gromov. To grok or not to grok: Disentangling general- ization and memorization corrupted algorithmic datasets. In The Twelfth International Conference on Learning Representations, 2024. URL
htw.:pppwrevv..eeeettttrrrrr/aea E5555B E5v5MB7
[7] Darshil Doshi, Tianyu Hc, Aritra Das, and Andrey Gromov. Grokking modular polynomials, 2024. URL https://arrxt org/ aba. 2406. 03495
[8] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma. Dawn Drain, Deep Ganguli, Zac Hatdield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse. Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan. Sam McCandlish, and Chris Olah. mathe- matical framework for transformer circuits. Transformer Circuits Thread, 2021. https//ttransfommrr circuits.pub/2022 1 /fffa//workiiddx...eee
[9] Deep Ganguli, Danny Hernandez, Liane Lovitt, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova Dassarma, Dawn Drain, Nelson Elhage, et al. Predictability and surprise in large generative models. In Proceedings of the 2022 ACM Conference Fairness, Accountability, and Transparency, pages 1747-1764, 2022.
[10] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers lear in-context? case study of simple function classes. 2023.
[11] Andrey Gromov. Grokking modular arithmetic, 2023. URL https://arxiv..ott www./ogg///aaaaeeeeeeeeeeee http://b//300...2220 92677
[12] Jiuxiang Gu, Chenyang Li, Yingyu Liang, Zhenmei Shi, Zhao Song, and Tianyi Zhou. Fourier circuits in neural networks: Unlocking the potential large language models mathematical reasoning and modular arithmetic, 2024.
[13] Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers leam in-context beyond simple functions? a case study on learning with representations. In The Tweifth International Conference on Learning Representations, 2024, URL htthttts:://////tt htttps://openrew/ net/ net/forumml farrummmddee -iddiddddwwwwaw.....

---

[14] Roee Hendel, Mor Geva, and Amir Globerson. In-contexi learning creates task vectors, 2023.
[15] Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. Gereral-purpose in-context learning by meta-learning transformers, 2024.
[16] Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning, 2024.
[17] Sheng Liu, Haotian Ye, Lei Xing, and James Zou. In-context vectors: Making context learning more effective and controllable through latent space steering, 2024.
[18] 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.
[19] llya Loshchilov and Frank Hutter. Decoupled weight docay regularization. In International Conference on Learning Representations, 2019. URKLLtttpa:/oopereeiiee wwtp:/eeeeetttooooteeeeee netr/forwm/idd-BkkkR....r
[20] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacoh Steinhardt. Progress measures for grokking via mechanistic interpretabiliity. In The Eleventh International Conference on Learning Representations, 2023. URL /ttcpps://peeneeviee..ee net. fo
[21] Eshaan Nichani, Alex Damian, and Jason D. Lec. How transformers learn causal structure with gradient descent, 2024.
[22] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield- Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022.
[23] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, buggh-pefformnce deep learning library, 2019.
[24] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.0211 2022.
[25] Allan Raventos, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence non-bayesian in-context learning for regression. Thirty-seventh Conference on Neurul Information Processing Systems, 2023. URL Bttl:ttttt///ppreriiiee
[26] Aaditya K Singh, Stephanie C.Y. Chan, Ted Moskovitz. Grant, Andrew M Saxe, and Felix Hill The transient nature of emergent in-context learning in transformers. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL htttttt:t//ppreeeee
[27] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.
[28] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2023.
[29] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint 220iv 07682, 2022.
[30] Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks. In Thirty-seventh seventh Conference Neural Information Processing Systems, 2023, URL (ttLittts////peereeee http:///ppppoooeeeeeeeeee///oddd www.://ooornnddd/mmomeoooooo
12

---

A Experimental Details
A.1 Model and Training Hyperparameters
Architecture We used GPT-like architectures with Rotary Positional Embedding (0 10, 000) and ReLU activations. We fix the number of heads H 4, embedding dimension dembed = 512 and MLP widening factor 4 throughout every model. We use d 2, 4, 6 throughout the paper. Embedding layers and the output layer are tied via weight tying.
Initialization All linear layers and embedding layer weights are sampled from Gaussian distribution (0 0.022) at initialization, with the exception that the last linear layer each MLP is sampled from (00, 0.02² /2d). No bias is used in any layer.
Optimization and Schedule We trained most models using Adam W optimizer with learning rate
? 1.5 10 weight decay 2.0, 31 0.9, B2 0.98. e 10 batch size B 1024, = 32 in-context examples for 200k steps, together with 5% linear warmup starting from 0.01n and cosine annealing towards the end to 0.1 Weight decay is not applied to LayerNorm layers.
Hyperparameter Choice For d = 2 models we scanned learning rates 7) E (7.5 x 10-3 1.5 x 10-1 10-1,3x x 10-1, 6 10-1) and weight decay values {0.5, 1.0, 2.0, 5.0). Then we transfer our hyperparameters other depths. Benefiting from the extra scaling in the initialization of the last linear in MLP, we find that the hyperparameters perform well for other depths. For larger p values, we lowering down the learning rate to 10
A.2 Further Details of Each Plot in the Main Text
Figure (b) Phase diagram constructed using data shown Figure 27 threshold for defining
each phase is set to 75% for the corresponding set; (c) accuracy number of shots eurve for d 7l..d. 28, p 29 and a = 0.7. (d, e) model with p 29, ni.d. = 2 and a 0.6.
Figure 3 (a) Selected the best out of three random seeds with early stopping. (b, c) averaged over three random seeds with standard error labeled. All o.o.d. data are measured every 1, 000 step for 16 randomly sampled sequences along the pre-training. (d, e) Used checkpoint at the end of pre-training, averaged over 128 random sequences sampled from So.o.d.
Figure 4 (a. b) Both phase diagrams are the best selected from three random seeds with early stopping.
The loss number of shots curves are plotted with 0.6.
Figure 5 We use a d = 2 model trained with a = 0.6 and n4.d. = 29, and a d = model trained with 0.8 and mi.d. 2 with batch size B 1536.
Figure (a. b) Attention heads extracted form = 2 model trained with ni.a. = 29 and a = 0.6. (c. d) Both models are trained with 0.6.
B Details of resources used
To generate each phase diagram and training curve, we used 90 GPU days on NVIDIA A100 40GB, with automatic mixed precision (BFloat16) and Flash Attention implemented in the PyTorch library (23). During the exploring stage, we also used around another 90 GPU days. Most inferences were running on 1/7-NVIDIA A A100 40GB GPU, of which the cost is negligible compared to pre-training cost.
C Table of log map for p
29
0
2
3
4
5
6
7
8
9
10
11
12
13
14
loga7 (n)
28
15
19
2
22
6
12
17
10
9
11
21
18
27
15
16
17
18
19
20
21
22
23
24
25
26
27
28
logur (n)
13
25
23
24
26
20
16
1
14
13

---

D Task Selection and Sequence Design
Task Selection During the initial exploration phase, we observed that it was challenging for the model to learn multiple modular arithmetic tasks simultaneously. Typically, the loss would stagnate plateau indefinitely.
We hypothesize that when the model is trained on multiple modular arithmetic tasks, the strong bias inherent in each individual task may interfere significantly. When tasks selected randomly, the resultant noise appears to inhibit the learning capabilities of the model, preventing from acquiring any meaningful patterns rules from the data.
Ultimately, we adopted more structured approach to task selection as illustrated in Figure Figure 2 based the following rationale: If the model is to learn specific task vector (a', b²). it would presumably be more straightforward one the two components - either or b' - has already learned by the model. This would leave the model with only the other component to decipher, which we assume is a comparatively simpler task. Thus, we decided to employ the rectangular rule for sampling task vectors, which we believe strategically reduces the complexity of the learning process by partially leveraging previously acquired knowledge.
In Figure g we show an ablation plot of phase diagrams with randomly selected task vector collection for pre-training. We still use the special sequence design.
i.d. train
1d. test
0.0. tran
0.0.d. test
pretraiming inputs fraction (o)
0.1
0.5
6
0.4
40
0.
20
2a 2
2 pre-treining (0...))
2' 2
2a
28 2
22
2 pre-traiingg
2
27
2 (
2
2
27
2
2
pre-traiming lasks (....)
pre-traning (h...)
Figure 9: ==0 6 phase diagram with task vectors randomly selected, selected the best results from two random seeds. Learning rate 77 10
Sequence Design Following similar spirit, we use balanced batch where sequences generated from all task vectors appear exactly the same number times during the training. We further align the examples across sequences generated by different task vectors, which believe reduces the chance of the model getting confused by the same input (z, y) appearing at different positions within the batch. Without this design, could not make the model train.
Here, we also show the training curve per task in Figure 10, trained with all of our tricks. We see that the model first learned very few tasks and then eventually found its way out. For training without the task selection and sequence design, the loss typically plateau around 3.0.
14

---

Tain
3.5
0.5
ACCUURO 26)
1.0
02
0.5
20000
40000
50000
80000
100000
22000
42000
60000
80000
100000
stop
S9OP
3.5
0.5
Mobureey N54
0.6
02
0.
0.0
20000
40000
60000
30000
100000
20000
42000
60000
80000
100000
stop
5900
Figure 10: d = training curve, with loss and accuracy measured per task for the last in context example. Trained with ri.d. 2 and 0.8. There are six tasks learned almost instantly: (5, 26). (21, 23). (28, 27), (3,9), 9), (10, 10), (5,9).
E Additional Interpretability Results
In this section, we show additional results on interpretability. Where we continue our discussion on how the algorithm might be implemented inside the model. Importantly, This includes PCA analysis of embedding (Appendix (.. other attention heads (Appendix(E.)) (..)) attentions and MLPs outputs (Appendix E.3). and finally, the role of MLPs (Appendix E.4). We think it is also crucial to study LayerNorm layers, which we leave for future work.
E.1 PCA over Embeddings
We begin the discussion with the role of embedding layers. By further examining different models, we find highly structured embedding from d = 2 models. models, on the other hand, does not have such structured embedding layers.
First we focus the embedding of d models. shown in Figure I1 clearly, the logarithm of each number is split into even and odd groups, and each group forms one circle, which is suitable embedding for doing modular multiplications. However, one should note that this will not obviate the importance of the head shown in Figure ((a), the model still needs way to distinguish the same number that appears in the different positions in the context.
Curiously, we could not find such a structured embedding for the d = model, as shown in Figure I2 However, as we will show in the next subsection, this non-structural embedding, together with the first layer, prepares a foundation for the essential heads in the latter layer to create similar Peirele-of-cirelee feature as we shown in Figure 6 (a, b).
E.2
Attention Heads
E.2.1
d 4 Model
To continue the story, we analyse the attention heads different models. We first study the d model, where we also find similar head with "cirole-of-circles" (Figure Moreover, two other
15

---

0.0-
0.06
0.0-
144
prhopple component
2 princpicc comparent
principle component
0.02-
0.06
3
10
23
0.00-
20
15
-0.02-
10
1
20
20
0.00-
-0.00
15
16
0.02
-004-
-0.05 0.00 -3.02 aãoo dide 004 0.00 0.00 principle comporent
-0.05-0 0.000000000 0.0z 0.d a.d oite
-0.04 -0.0z oce o.t2
0.04 0.05
o.d
1 principle companent
2nd principle component
(a) seed
(b) seed
Figure 11: PCA over the embedding layers of = 2 models with different random seeds. Each number is annotated by its log27
value.
0.00
0.08
12
12
prinsple component
0.08
3 principle comporent.
grnccile comparennt
0.04
0.04
15F
0.02
130
15
0.02
18
0.02
264
14
0.00
2g
0.00
21
0.00
261
-0.02
20
-004-
-0.44
24
-0.04
0.56
-0.0
10
-0.08
10
-0.08-0'o2 02 ono nite obe oiss obnn 0.10 principle component
-o.0e-o 02 no oiz o.n 0.60 principle component
-0.06-0.ot oloo o olz ab4 nos obn 2nd component
(a) PCA over the embedding layer for top-3 principal components
Figure 12: PCA over the embedding layers of d 4. Each number is annotated by its log27 value.
heads put together are seemingly equivalent to Figure 6(b.. We surmise that any model that solves modular arithmetic in-context requires such heads.
From Figure 13 a), we see that the head still pays attention locally within three token positions. Importantly, also creates a circle-of-circles while performing PCA over concatenated (r, y) features. The difference here that the circle winds twice to go back its origin. We believe that this factor of two differences in the period means that this head is effectively combing the row of the embedding layer of the d 2 model and the role of the head Figure (aa. Overall, we not yet know why the model needs to split the logarithm of numbers into even and odd groups.
From Figure 13 c), find two heads that are structured but different from Figure b The PCA pattern forms circles while annotated within logarithm space, but depends on the choice of the sequences. This hints that those two heads are trying to exchange information with the previous inputs. We leave PCA analysis across different in-context examples to future work.
E.2.2 Model
Next we focus on the model. In Figure 14 we plot similar PCA to the one in Figure 6 but with different concatenated features from the same heads: (x, 2) and (y. z). We again see circle-of- circles, solidifying our argument that this head provides proper spaces for later heads to perform modular operations.
In Figure 15 we dump all the attention patterns we had for = models. The corresopnding PCA over concatenated (r, v) features are shown in Figure (except for layer 1 head 4, for which we plot for (a., =)). The special head we choose here has a similar behavior to the one in Figure (()) with its main focus shifted by position but still within two preceding tokens.
16

---

0.033
-0.034
2nd priciple componant
principle companent
0.02
0.038
2nd priccple
+0007-
0.00
0.000
-0.01
-0.009
-002-
0.02
001 0.00
0.11
0.02
0.3
00000,0010 0000 dos dos dos do7 1st principle component
0.02
0.00
0,02
0.04
obs
1st principle component
1st principle component
(a) layer 2 head 4
(b) layer head 1
(c) layer 3 head 4
Figure 13: PCA over the features of specific heads for log27 value. We plotted PCA over (a. b) (z,, y) and (c) (y, z) concatenated features. Where the circle in (b, c) heavily depends on the sequence and (a) remains unchanged.
model. Each number annotated by its
0.10
princppis componort
0.08
2nd primaple componeert
0.04
0.00
-0.02
-0.5-
-0.4
-0.05
-0.10-
-000--0 -olo noo uba obe nós aim principle component
principee component
2x)
(b) (v/
Figure 14: PCA over the same head in Figure 6(a,, with feature concatenated differently. We annotate numbers with their corresponding logarithm.
In Figure 17 extend our PCA analysis to different task vectors and up top-4 components, where the spatial structure of the features is better depicted.
We have reasons to believe that the non-circle heads are not essential to the models' performance. However, leave careful exploration of this front for future work.
E.3
PCA over Attention and MLP Outputs
In Figure 18 we plot d = 2 PCA analysis up to top-4 components for attention and MLP outputs.
17

---

(a) layer I, head
(b) layer I, head 2
(c) layer 1, head 3
(d) layer 1, head 4
(e) layer 2, head
(I) layer head
(g) layer 2, head
(h) layer 2, head
Figure 15: All attention heads in d 2 model.
0:24
20d paniee comparentt
2rd permal compaaree
201 MPV- comporent
A64
A82
pr DA componeei
pwpppee compenent
pr80 Companest
(a) layer 1, head
(b) layer 1, head
(c) layer 1, head
(d) layer 1, head 4, (zz, 2)
0112
2:4 www.0000 01000000
1015
0228
4.5
-0330
olo
219
<10
Joi sàt ols
410 olo n/c 8/6 0.06 8.6 cll
.005
3/5
(c) layer 2, head
(f) layer 2, head
gg)
layer 2, head 3
(h) layer 2, head
Figure 16: PCA for concatenated (x, yax 2) features for all heads in d 2 model.
E.4 Cosine-similarity analysis for MLPs
As mentioned in Section 5.2 Skills II and IIII are likely implemented within the MLP layer. While we provided extensive details on how the signature of Skill was observed in MLP, we could not find clear signals similar to those in Nanda al. [00 Gromov [1] from MLP layers Skill III.
E,4.1
d = 4 model
Here, we show an extended version of Figure for all layers from the d model. plot for 2-shot (Figure 20) and 16-shot (Figure 21, the task vector used is (a, (2, 6).
18

---

layer head
layer head
Figure 17: More PCA analysis similar to Figure 6 Top-2 components of PCA contribute ~ 40% variance of the layer 1 head and 20% variance of the layer 2 head, while the top-4 components contribute ~ 60% and ~ 40% accordingly.
1yw
MP
Figure 18: PCA analysis of Attention and MLP outputs. Attention modules contribute ~ 50% and 30% to top-4 components of PCA of layer and layer 2. MI.P modules contribute 10% and
12% accordingly.
24 21 5) 6
Layor
Layer2
Layor
Figure 19: MLP (hidden) activations as function of inputs (x,y) for d 2, models.
E.4.2 d d= model
Here, show extended version of Figure 7 for all layers from the d = 2 model. We plot for 2-shot (Figure 22) and 10-shot (Figure 23 the task vector used is (a, b) = (2, 6).
19

---

(a) y, layer
(b) y. layer
(c) layer
(d) y. layer 4
(c) layer
(f) layer
(g) layer
(h) layer
Figure 20: Cosine-similarity of d = 4 model for every block, measured at 2-shot.
(a) y, layer
(b) layer
(c) layer
(d) y. layer
(e) layer
(0 layer
(g) layer
(h) layer
Figure 21: Cosine-similarity of 4 model for every block, measured at 16-shot.
(a) v. layer
(b) layer
(c) layer
(d) z. layer
Figure 22: Cosine-similarity of d 2 model for every block, measured at 2-shot.
20

---

(a) v. layer
(b) y. layer 2
(c) 2, layer
(d) 2. layer 2
Figure 23: Cosine-similarity of d 2 model for every block, measured at 10-shot.
Id. trah
0.0 d. Irain
l.d. tnst
O.O.d. d. tast
26
20
20
33
15
band
20
34
25
18
20
24
26
12
18
27
24
28
20 28
22
(a) single label corruption,
8
20
28
20
05
04
A20
84
02
A00
02
2.78
(b) multiple label corruption, d 4
(c) multiple label corruption, 2
Figure 24: (a) Single label corruption Performance of last token prediction as function of changing position for a single label corruption. The x-axis shows different shots presented to the model, while the y-axis shows the position where the label is corrupted 2j + 2/,. The d 4 model is remarkably robust for long sequences, indicating the use of an algorithm which is not sensitive to the position of preceding examples. Qualitatively similar plots hold for other task vectors well. (b, e) Multiple label corruption, at random locations for d = 4 and d = 2 models respectively. The d = model more resilient to label corruption than the d model, implying that the algorithm the latter employs is imperfect due to limited capacity.
E,5
Label Noise
To gain insight into how the model combines the in-context examples, introduce label-coruppiin in the in-context examples. In particular, we note the effect of (i) amount and (ii) position of label corruption the model's 's performance. When we corrupt single in-context example for d 4 model, the model performance remains unaffected for longer sequences. This hints at that weighted average of the in-context inputs being used in model prediction. The d = 2 model, however, did not show such resilience.
Next, we corrupt multiple in-context examples in random locations. We study the effect on model per- formance as the amount of corrupted labels increases. While the = model is easily overwhelmed, the d = 4 model is able to offer strong resistance even at 40% label corruption, for long sequences. This behavior remains invariant with the change task vector for the particular sequence, indicating the universality of the underlying algorithm necessary for o.o.d. generalization.
21

---

Additional Training Curves
We plot some selected training curves for d (Figure 25 and (Figure 26 from Figure phase diagrams. We see that even for d = 4, ICL can be transient. With increased or nl.d.. the transient nature goes away.
100
7.0
6.0
6.0
1.0
0.0
(a) d 1, 0.7
100
7.0
2.0
1.0
0.
step
step
(b) d = 4. 0.6
Figure 25: Training curves for 4, averaged over three random seeds.
aICCLFoc (%)
75
4.
9.0
2.0
50k
100k step
150k
200k
50K
100k step
150k
200k
(a) d 2, 0.7
9.0
acauraoy (%)
75
2.6
50k
100k step
150K
200k
30K
100k stop
150k
200k
(b) = 2,
Figure 26: Training curves for 2, averaged over three random seeds.
22

---

G Additional Phase Diagrams
In Figure plotted detailed/extended versions of the phase diagrams shown in Figure Figure
The four phases story we have shown in Figure T still hold for other depths.
b.o.d.
100
0.8
fraciio
0.1
60
60
0
40
6.
-20
6.
2a
2
2
29
2
2
2
23
2
2
2
2
2 (n,
2
2
2
29
2
2
2
2
24 2
2
2
2
2
pre-training
(n.
treaitraiiinng
pre-training
pre pre-training
(m)
(a) d = 6, last shot
1.d. train
id. test
o.o.d. train
0.0. d. test
fraction
6.
80
60
05
6A
40
63
20
62
21
2
2
27
2
2
2 pro-training tasks
# of pro-training tasks (n.
# of pro-training tasks (n.
# pre-training tasks (m.
(b) d 4
1.d. train
Idd.
0.o.d. Irein
0.o.d. Test
100
0.6
fracion
6.
6.6
60
40
20
2
24
2
2
2
2
2
2
2
2
2
2
2
2
2a
2
2 24
23
2a
2
2
# of pre-training tasks (n.
of pre-training tasks (7)
#
pre-training tasks (0.a.)
#f pre-training lasks (1)
(c) =
Figure 27: Phase diagrams on all four for d 4 and d 2.
23

---

H Different Choice of
In this section, check the effect of varying task difficulties, i.c. the value of p. In Figure 28 plotted o.o.d. generalization accuracy. Clearly as the task gets harder, the model needs to see more tasks to generalize -itto--distribution.
4.0
100
3.5
3.0
accuracy (%)
75
2.5
50
2.0
1.5
25
1.0
0.5
2²
2"
2'
2"
2"
# of pre-training tasks (M,
of pre-training tasks (ni
Figure 28: We compare the best performance for d = 6 models with different ni.d. values, averaged over three seeds. We use learning rate r 10-1 4 for p and p 47, while keeping other hyperparameters the same as p 29.
24

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our main claim are summarized in Figure Section and Section 5 offfe detailed explainations.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We include the limitations of our work in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
25

---

Justification:
This is not theoretical paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We explained our settings in Section S and hyperparameters in Appendix A
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
26

---

Answer: [Yes]
Justification:
GitHub link offered.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/
wwt.Subbmiiiioooiiioooooooio for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Details are summarized in Appendix A with explanations on data split in Section We also provide code supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We show the standard error in most training curves, with an average of over three random seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
27

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We have a Appendix B on this.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
We have read and understood the code of ethics; and have done our best to conform.
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification; Our work involves small models and toy datasets. does no impact the society at large, beyond improving our understanding of certain aspects of deep learning.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
28

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Since train small models toy-datasets, work poses no risk of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We are not shipping our code with any source code or binary files from any other existing libraries, there are no concerns over getting permission or including license. We did cite open-sourced libraries, e.g. PyTorch, in our paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
29

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: will release our code base with included readme files. We do not ship any source code binary files from any other existing libraries.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
30


Reviews:
[{"review_id": "m1s5JPqAc9", "reviewer": "Reviewer_CrUb", "summary": "This paper develops novel insights into in-context learning and how it works in Transformers. To this end, the authors propose a generalization of the modular arithmetic task explored in several prior works on grokking. Unlike those works, the structure of the defined task is more rich, enabling an analysis of both in-distribution generalization (standard test evaluation) and out-of-distribution generalization (which is itself broken down into two variants).", "questions": "A few questions below that I would like to see answered. \n\n- **PCA variance.** Given this is a rather rich geometry in 2-D, I'm slightly surprised to see PCA captured it. Did you have to do some preprocessing? How much variance is explained by the two projected components? If there are other components that are not shown but have a large variance, what do those components encode---can you try 3D plots?\n\n- **What does the MLP do?** Given the mechinterp focused on attention solely, it is unclear what role MLPs played. Two experiments to try here are: (i) train attention only models to see if MLPs are even necessary, and (ii) perform the PCA analysis to uncover representations' geometry at the level of attentions and MLPs at each block in the model. Experiment (i) may require retraining models, so I understand if the authors are unable to conduct it, but my expectation will be that you will see that model \"internalizes\" task vectors and records them in MLPs. Attention only models can solve the task, but I expect the representations' geometry will be quite different. For experiment (ii) however, I expect that's easy to run and is merely repeating the plotting script on intermediate representations as a forward pass occurs through the model. If the geometry is primarily formed at attention layers, we'll see that in this experiment; vice versa, if it forms via MLPs, we'll see it explicitly.", "limitations": "Limitations are fairly discussed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is fairly well written and clear. Going beyond the standard linear regression task to study ICL was great to see as well.\n\nThe main selling point for me are the empirics though---I really like the results! The visualization of how the model represents concepts relevant to this paper's setup is quite beautiful: the circle of circles was fascinating to look at and, arguably, not something I expected. In retrospect, I can rationalize this as making sense---we get circular embeddings in grokking, so circle of circles is the logical geometrical extension here. Results on scaling are interesting in their own right as well.", "weaknesses": "I do not have any major apprehensions, except for the related work, which I think is relatively sparse.\n\n- **Related Work.** At this point, the topic this paper is focused on has a rather rich literature and I think a more detailed related work is warranted (perhaps in the appendix if space is an issue). For example, the results by Kirsch et al. (which is cited) are very similar to what authors show, especially results on scaling effects. The main different is width scaling in that paper and no geometric analysis, but nonetheless the relationship warranted more emphasis and discussion. Similarly, several recent works have explored OOD generalization of toy ICL tasks defined in prior works (e.g., see Ahuja and Lopez-Paz [1] for work on linear regression tasks and Ramesh et al. [2] for group arithmetic tasks). Regarding grokking, there are several works exploring the phase transition-y nature of this task. For example, see Kumar et al. [3]. The transient nature of ICL also has negative results (see Reddy [4]), which are worth discussion since they are the primary conclusion in depth scaling as I see it.\n\n[1] https://arxiv.org/abs/2305.16704\n[2] https://arxiv.org/abs/2311.12997 \n[3] https://arxiv.org/abs/2310.06110\n[4] https://openreview.net/forum?id=aN4Jf6Cx69", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the encouraging feedback and helpful suggestions.\n\n## Weakness\nWe thank the reviewer for pointing out the highly relevant references. We will add the citations and utilize the additional page allowance in the final version to discuss their relation to our work.\n\n## Questions\n\n**PCA Analysis**: We conducted PCA without any preprocessing. As per the reviewer's suggestion, we expanded upon Figure 6 from the original manuscript by plotting higher order PCA components. The results are presented in Figure G.2 of the attached PDF and discussed in the Global Rebuttal. We analysed the top-4 components, and found highly structured features. The top-4 PCA components account for a significant portion of the PCA variance. (Note that we present 2d slices instead of 3d plots of PCA because sparse 3d plots shown in 2d are difficult to comprehend.)\n\nFurthermore, following the suggestion of the reviewer rwBm, we plot similar features with a different task vector $(a,b)$ and a different number of shots. These results serve as evidence for the claims made in Figure 6 caption. Specifically, the PCA constructed from the first layer's head remains unchanged across different task vectors (up to a negative sign along certain directions). In contrast, the PCA derived from the second layer's head changes with the choice of task vector.\n\n**PCA of Attention outputs**: As per the suggestion of the reviewer, we present the PCA analysis of Attention outputs (as opposed to individual heads) in the top row of Figure G.3 of the attached PDF. We find highly structured top-4 PCA patterns in Layer 1, which also account for a significant fraction of PCA variance. Layer 2 exhibits less structured organization and the contribution of the top-4 PCA components is diminished.\n\n**Analysis of MLP features**: In the attached PDF, we have added two main results concering MLPs. \n\n1.  We extended our analysis to include PCA Multi-Layer Perceptron (MLP) features, as shown in the bottom row of Figure G.3 of the attached PDF. The results demonstrate:\n\n- Layer 1: Highly structured patterns are evident in MLP features. The top-4 components contribute substantially to the overall PCA variance.\n\n- Layer 2: Features exhibit less structured organization, and the significance of the top-4 components is diminished compared to Layer 1.\n\n2. Additionally, in Figure G.1 of the attached PDF, we have shown the post-ReLU neuronal activations from various layers as functions of $x,y$. We find highly structured activation patterns across layers, especially in Layer 3 of $d=4$ model. For a detailed account of the MLP results, please refer to the Global Rebuttal.\n\nWe will discuss these new results in the additional page allowance of the final version. We believe that in a future work, our analysis of the attention heads as well as MLP activations can be tied together to infer an end-to-end algorithm for our setting.", "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": "RPlLeNuaca", "reviewer": "Reviewer_Jerg", "summary": "This paper studies the emergence of in context learning and skill composition in autoregressive models. They create an algorithmic dataset to probe how autoregressive models use tasks learned during training to solve new tasks. They find that more training tasks lead to a generalizing / algorithmic approach instead of memorization.", "questions": "Since multiplication can be viewed as repeated addition, isn’t skill 2 an extension of skill 3 (or can even be viewed as skill 3 composed with itself multiple times)? Is hierarchy of skills important here?", "limitations": "As acknowledged by the authors, this work is limited to particular algorithmic datasets.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This work introduces a new algorithmic dataset (with modular arithmetic tasks) that force models to learn a variety of tasks. The work finds that when the number of tasks goes from small to large, the model transitions from memorization to generalization.\n- This work has many interesting experiments. I found Section 5.2 (Attention Heads Implement Essential Skills) pretty interesting.", "weaknesses": "- The definition of task diversity is not well defined. Is the number of pretraining tasks truly indicative of task diversity? I think the paper could benefit from some justification of this assumption.\n- The paper claims that for larger models, early stopping is necessary (line 52). While I appreciate that the authors used GPT-like architectures to reflect realistic settings, the architectures in the experiments are not that large. Even amongst popular open source models, the smallest are usually around 7B parameters.\n- Many works in the continual learning and meta learning literature suggest that training on multiple tasks at once leads to better generalization. Perhaps it is worth including brief discussion on the connections between this point and the model’s ability to generalize ood which is predominantly determined by the number of pre-training tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments.\n\n## Weaknesses\n\n**Task Diversity**: \nOur definition of task diversity follows the existing works on in-context learning with linear regression, with a key difference: \nsince our tasks are defined over a finite field, the total number of possible tasks (labeled by the task vectors $(a, b)$) is _finite_ and equals $p^2$. This differs from commonly discussed cases of linear regression, where the set of tasks is infinite (and, in fact, _continuous_). Consequently, the number of pre-training tasks (as a fraction of the total tasks $p^2$) is a natural measure of task diversity. There is a subtlety in that tasks may not be completely independent from each other, and a true definition of task diversity should include a reference to independence. However, it is not clear to what extent the model leverages this possible redundancy -- and similar point is also omitted in the works on linear regression. We decided to not go down that rabbit hole and defined task diversity in a naive way.\n\nThat being said, we acknowledge that different ways of sampling the task sequences could also influence o.o.d. generalization. To address this, one might construct a phase diagram with an additional axis representing task sampling. However, this would require an order of magnitude more computations and a detailed multi-page discussion, making it impractical for our current study.\n\n\n**Early Stopping and Larger Models**:\nWe appreciate the reviewer raising this point. \nFirst, we would like to clarify that the \"larger model\" mentioned in line 52 refers to the comparison between the $d=6$ model and the $d=2,4$ models used throughout the paper. \nNotably, these settings are sufficient to demonstrate our point, as the model's scale should be measured relative to the dataset size. The SoTA LLMs are pretrained on corpora much larger and diverse than the arithmetic tasks that we study in this work. We agree that larger-scale experiments would be necessary to transfer the insights gained from our study to modern LLMs. However, such experiments are far beyond our current capabilities due to limited GPU resources. \nFinally, the purpose of including details such as early stopping is to aid reproducibility of our results -- we do hope that the community will explore and generalize our setting.\n\n**Relation to meta-learning and continual learning**:\nWe thank the reviewer for pointing out this interesting connection. It is indeed possible that some of the insights from our work finds connections to these areas. In the current version, we have cited one work [1] related to meta-learning. In the camera-ready version, we will include a more elaborate discussion with more references.\nWe welcome suggestions for specific resources about continue learning that the reviewer has in mind relevant to our study.\n\n## Question\n\nThe reviewer raises an interesting point about the hierarchy of skills, and is correct in pointing out that multiplication can indeed be constructed from repeated additions. \nHowever, it is important to think from the model's perspective. We believe that *efficient* implementation of finite field operations by the model requires separate components to perform addition and multiplication. One intuitive way to think about this is that the models do not have sufficient depth to perform arbitrary repeated additions to construct multiplication. Instead, the models build correct representations to implement multiplication.\nConsequently, it is better to think of the numbers on a finite field $\\mathrm{GF}(p)$, with distinct operations of addition and multiplication.\n\n[1] Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz; \"General-purpose in-context learning by meta-learning transformers\"; https://arxiv.org/abs/2212.04458 (2022)", "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": "ZLQwBMiHgE", "reviewer": "Reviewer_PjEW", "summary": "* The authors propose a synthetic sequence learning problem that I would call\n  'in-context modular regression', an elegant generalisation of prior work\n  studying modular addition and in-context linear regression.\n* Using carefully constructed batches the authors are able to train\n  transformer models to perform regression for a subset of tasks (weights)\n  and a subset of inputs.\n* The authors show that under some conditions on the data distribution and\n  model architecture, the transformers not only achieve good performance on\n  tasks and inputs included during training, but they also generalise to new\n  tasks and/or new inputs. The authors document the conditions governing\n  these generalisation capabilities in detail including showing phase plots\n  and observing that in larger models, the generalisation properties are\n  transient (they appear and then disappear across the training process).\n* The authors postulate a breakdown of skills required to correctly perform\n  the task. They effectively isolate and examine the abilities of their\n  models to perform each component task. They also inspect the activations\n  of each head and identify patterns suggestive of partial mechanisms\n  underlying the generalising behaviour of the models.", "questions": "1. Why is the title 'learning to grok'?\n\n    * Is this meant in the sense that the grokking of a modular addition task\n      is occurring in-context? If so, this seems a little inaccurate, since\n      the phenomenon analogous to 'grokking' seems to still be occurring\n      during pre-training.\n    * To be honest this part of the title has puzzled me since I first looked\n      at the paper. Even if my understanding above is wrong and the title has\n      an accurate interpretation, that I have failed to notice it might be\n      one data point suggesting that if you are going for a title that is\n      both short *and* informative, this might not be the right choice.\n\n2. In the figure 1 caption, is it possible to offer a clearer summary of the\n   difference between in-distribution generalisation and out-of-distribution\n   memorisation? On my first read through, treating the figure and caption as\n   an overview of the work's main results, I had trouble distinguishing these\n   two concepts.", "limitations": "The authors transparently acknowledge all of the limitations I was able to\nidentify within the paper itself.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "I thank the authors for submitting their excellent work which stands to have a substantial impact in the science of deep learning.\n\n* The work makes a meaningful contribution to an exceptionally important and\n  interesting topic of the emergence of capabilities and internal mechanisms\n  in deep learning.\n* The setting and experiments neatly isolate and clearly demonstrate several\n  interesting phenomena of emergence of capabilities and shifting in the\n  solutions found by deep networks throughout training, contributing to the\n  field's developing catalogue of examples of these phenomena.\n* Moreover, the proposed synthetic problem is both rich and elegant. I expect\n  this framework will become a fruitful test-best for follow-up work studying\n  emergence phenomena, helping the field to improve our empirical and\n  theoretical understanding of these phenomena.\n* The authors also offer a partial behavioural and mechanistic analysis which\n  is a solid starting point for a more detailed understanding of the learned\n  structures that emerge in this setting.\n* While some elements of the analysis are complex, the authors have done an\n  exceptional job of clearly presenting their findings. I feel careful study\n  of each section and figure in the main text was rewarded since there was no\n  question that occurred to me that was not addressed in the authors' clear\n  descriptions or figures.\n* The authors have acknowledged all of the related work that I am aware of.", "weaknesses": "I have not noticed any weaknesses in the paper that would temper my overall\nrecommendation to accept. However, I note the following weaknesses, some of\nwhich the authors have already acknowledged, and others which they may like\nto take into consideration if they are interested to improve the paper\nfurther.\n\n1. **Delicate training set-up.** The authors explain that training\n   transformers on multiple modular addition tasks crucially relies on\n   following a delicately balanced batch construction methodology.\n   I am left wondering if this batch construction methodology, as a further\n   departure from the standard language modelling setting, has any other\n   implications for the learning process that may affect the generality of\n   the results.\n   Note: This weakness is not decisive because the authors clearly document\n   their training methodology and it's not *that* artificial anyway.\n\n2. **The mechanistic analysis is only partial.** The authors admit that they\n   have not been able to identify an end-to-end mechanistic model of how the\n   trained transformers perform the task. This leaves their posited skill\n   decomposition and partial mechanistic analysis open to the possibility\n   that they are incomplete.\n   Note: I think the contribution the authors have given in terms of the\n   setting, the generalisation phenomena, and the partial skill decomposition\n   and mechanistic analysis are already significant.\n\n3. **Relationship to prior work.** The related work section does a good job\n   of summarising the contributions of prior work in in-context linear\n   regression and modular arithmetic in the context of transformer models.\n   However, I feel that this section could be improved if the authors\n   attempted to offer greater insight into the relationship between these\n   prior works and the present work. For example, the authors have an\n   opportunity here to informally describe the in-context linear regression\n   and the modular addition problem settings that the newly proposed setting\n   generalises.\n\n4. I noticed some minor text errors as follows, which I expect the authors\n   can easily correct.\n\n    * Line 94: The notation $[1, p^2]$ to me suggests a closed continuous\n      interval, whereas you appear to mean $\\lbrace1, \\ldots, p^2\\rbrace$, also in some\n      cases denoted $[p^2]$.\n    * It seems that equation 2 should read $\\ldots = (z_1^t, z_2^2) \\mod p$\n      and the equation on line 203 should read $c_1x + c_2y \\mod p$. That is,\n      $x$ and $y$ should swap places with $z_1^t$ and $z_2^t$. Is this indeed\n      a mistake, or am I missing something?\n    * In figure 6 (top row) there is a typo: \"Qeury\" on the vertical axis.\n    * In line 445 there is a broken link.\n\nI have not studied all appendices in detail.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the encouraging feedback and valuable comments.\n\n## Weaknesses\n\n1. **Delicate training set-up**: \nThe structured selection of tasks (rectangular rule) and balanced batches largely serve the purpose of making the pre-training more stable. Our intuition for using the specific setup was to reduce batch noise, as the training itself is challenging for this task. We strongly believe that scaling up the batch-size and model sizes will alleviate these constraints. We did not explore this avenue due to compute restrictions.\n\nWe will add this clarification to the final version of the paper.\n\n2. **The mechanistic analysis is only partial**: We have included additional results in the Global Rebuttal and the attached PDF, further strengthening our analysis. Notably, this includes highly structured neuronal activation patterns (Figure G.1). However, the end-to-end algorithm still remains an open question.\n\n3. **Relationship to prior work**:\nWe will incorporate the reviewer's suggestion into the camera-ready version of our paper.\n\n4. **Minor Text Errors**:\nWe have addressed all the minor errors pointed out by the reviewer, except one, which is not an error. Specifically, equation (2) and line 203 are not typographical errors. \nThe algorithm we propose differs from the conventional method humans use to solve linear systems of equations, which involves explicitly computing the coefficients $(a,b)$ from the in-context examples. Instead of finding the unknowns $(a,b)$, the model finds the right way to \"linearly combine\" $(x_1, y_1)$ and $(x_2, y_2)$ to equal $(x,y)$. The re-scaling factors $c_1, c_2$ in this linear combination can then be used to get the correct answer: $z = c_1 z_1^t + c_2 z_2^t$. To emphasize this point, we have modified Equation (2), which now reads:\n$$c_1 (x_1, \\\\; y_1 ) + c_2 (x_2, \\\\; y_2 ) = (x, \\\\; y) \\\\; \\mathrm{mod} \\\\; p\n\\qquad \\xrightarrow{\\text{find} \\\\; c_1, c_2} \\qquad \nz = c_1 z_1^t + c_2 z_2^t \\\\;\\mathrm{mod} \\\\; p \\\\; \\qquad\\qquad (2)$$\n\n    Here is an intuitive way to think about the algorithm: Instead of directly solving for the unknowns $(a, b)$, the model treats the in-context examples $(x_1, y_1, z_1)$ and $(x_2, y_2, z_2)$ as vectors. The task then becomes to fill in the last component of a new vector $(x, y, ?)$. By aligning the first two components using coefficients $c_1$ and $c_2$: $c_1 (x_1, y_1) + c_2 (x_2, y_2) = (x, y)$ mod $p$, the model naturally finds out the correct way of aligning $z_1$ and $z_2$ with $?$, given the fact that they are constructed with a same underlying linear relation $a x + b y$ mod $p$.\n\n    We will add a version of this explanation to the final version of the paper to enhance clarity.\n\n## Questions\n\n1. On the origin of the title: The authors have always been of the opinion that the central theme in grokking is the formation of highly structured representations at the end of training. In fact, these representations can be viewed as a 1st order phase transition (in the proper statistical mechanics sense as explained in https://arxiv.org/abs/2310.03789). \nIf we take this perspective then in the present work grokking happens in _two_ qualitatively different ways: (i) _as optimization time passes_ the model learns to solve the task in-distribution (and sometimes o.o.d.), which requires highly structured representations, and (ii) _as the number of in-context examples increases during inference_ the model performance steadily improves; completely solving the arithmetic problem with enough in-context examples. It is crucial that the predictions are conditioned on the sequence of in-context examples and this conditioning is _emergent_. A more accurate name could be \"grokking grokking\", but we decided to opt for a milder version.\n\n2. We thank the reviewer for pointing out the cumbersome phrasing. We have edited the part of the caption in Figure 1 that distinguishes various phases to make it clearer. The part of the caption now reads: \n\n    > **(b)** Phase diagram for a six-layer model. We find four different phases. (1) in-distribution memorization: The model *only* performs well on tasks $(a,b)$ *and* examples $(x,y)$ from the training set -- it does not generalize on unseen examples or tasks. (2) in-distribution generalization: model generalizes on unseen examples $(x,y)$ but not on unseen tasks $(a,b)$. (3) out-of-distribution memorization: model generalizes on unseen tasks $(a,b)$ but only for examples $(x,y)$ it has seen during training. (4) out-of-distribution generalization: model generalizes on unseen tasks $(a,b)$ for seen as well as unseen examples $(x,y)$. We focus on investigating phase (4) in more detail.\n\n    Additionally, we will add a table clarifying the performance on the sets $S_{train}^{i.d.}, S_{test}^{i.d.}, S_{train}^{o.o.d.}, S_{test}^{o.o.d.}$ in the four different phases, on page 4 of the main text. We hope that this will help avoiding any possible confusion in the definition of the phases.", "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": "PdJabwbJeE", "reviewer": "Reviewer_rwBm", "summary": "The paper studies the emergence of the in-context ability of the GPT-style transformer model trained using autoregressive loss and arithmetic modular datasets. It analyzes the influence of the number of tasks, number of in-context examples, model capacity, etc., on the ICL capability of an appropriately trained model (i.e., using early stopping). It also provides a persuasive “task decomposition hypothesis”, which is well supported by the ablation study and various experiments. The white-box analysis on the attention heads provides convincing evidence of the proposed explanation. Although there is a gap between the grokking settings (i.e., small model and toy dataset) and practical systems, the paper does a good job of explaining many important trends and concepts related to the emergence of compositional in-context ability. I enjoy reading this paper and suggest an acceptance.", "questions": "- The paper claims in line 147 that “As the o.o.d. performance increases, the pre-training performance simultaneously degrades “. However, it is hard to read this information from Figure 3-a panel 1. Maybe a different color mapping or adding numbers on these patches would be helpful.\n- Equation 2 is a bit hard to understand. How does it correlate to $z = ax+by$ ? (Although, from the latter explanations, I know the model relies on $c_1z_1^t + c_2z_2^t$ to get $z$, but it might be helpful to claim how it is derived.)\n- Better to define $GF(p)$, i.e., the Galois field, before using it.\n- Are the results in Figure 6 coming from $d=2$ or $d=4$? I can find the figure for all 8 attention heads for $d=2$ in the appendix, what about the $d=4$ case? It might be helpful to see if the pattern in later layers (i.e., attention focusing on different $z_i$) exists in shallow layers, and vice versa.\n- In line 264, the paper claims that the pattern depends on $(a,b)$, but it is hard to read that from Figure 6b.\n- As also mentioned in the strength part, is it possible to find some specific value in the weight space (e.g., attention weights, readout layers, etc.) that is highly correlated to $c_1, c_2$? If so, the hypothesis that the model first learns skill 2 (scale each example) and then skill 3 (weighted combine different examples) would be further verified.\n- The OOD settings studied in grokking or emergent ability setting are quite related to the compositional generalization and systematic generalization. It would be helpful to discuss them in the related works, here are some of them:\n    \n    [1]  Schott, Lukas, et al. \"Visual representation learning does not generalize strongly within the same domain.\" ICLR 2022\n    \n    [2] Xu, Zhenlin, Marc Niethammer, and Colin A. Raffel. \"Compositional generalization in unsupervised compositional representation learning: A study on disentanglement and emergent language.\" NeurIPS 2022\n    \n    [3] Ren, Yi, et al. \"Improving compositional generalization using iterated learning and simplicial embeddings.\" NeurIPS 2023", "limitations": "Discussions on how the findings help the practical system.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is easy to follow. Good presentation!\n- The experiments are well-designed, providing compelling support for the claims.\n    - The results in Figure 5 are cool.\n    - The skill decomposition discussed in section 5 is great. The clear pattern in attention heads verifies it very well. (The hypotheses could be further verified if the author can link the values of $c_1, c_2$ to some weights in the network, see the question part.)", "weaknesses": "- The emergent ability (or grokking) usually refers to a phenomenon in the model “got stuck” in a non-generalization region and suddenly gained the generalization ability. Hence some discussion about the learning dynamics, i.e., how the accuracy, loss, representation, ability, attention pattern, etc., gradually evolve during training would make the paper stronger.\n- The task and batch sample selection in this paper have many constraints (e.g., the rectangular rule, the balanced number of samples in each batch, etc.). However, the practical systems usually cannot strictly satisfy all these assumptions. Hence a more detailed analysis of how these assumptions influence the generalization ability would provide more insights to practical systems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the encouraging feedback and incisive questions.\n\n## Weaknesses\n\n**Emergent abilities / Grokking**:\nThe loss and accuracy curves are already presented in Figure 3 of the current version of the paper. We agree that the gradual emergence of useful representations as a function of training time are useful results to showcase. In the final version of the paper, we will also include the feature analysis (similar to Figures 5,6) for intermediate checkpoints during training -- as suggested by the reviewer. (Note that we have not included these plots in the PDF of the Global Rebuttal due to space constraints.)\n\nWe also note that we exercise a broader take on emergent ability than only learning dynamics. Emergent behaviours [1] are characterized by qualitative changes in model capabilities upon scaling-up (i) model size, (ii) dataset size, (ii) training duration, (iii) task-diversity etc. While the transition with training duration is an interesting aspect of the literature on Grokking, it is only one part of the emergent phenomena in deep learning. Moreover, careful initialization and optimization have been known to mitigate such effects [2]. However, the transition in capabilities with respect to dataset and model sizes are known to be more robust. \n\n**Pre-training task and batch selection**: \nThe structured selection of tasks (rectangular rule) and balanced batches largely serve the purpose of making the pre-training more stable. Our intuition for using the specific setup was to reduce batch noise, as the training itself is challenging for this task. We strongly believe that scaling up the batch-size and model sizes will alleviate these constraints. We did not explore this avenue due to compute restrictions. \n\n## Questions\n\n**o.o.d. vs pre-training performance**:\nThe performance trade-off between o.o.d. and pre-training is more clear in the $d=2$ and $d=4$ models, shown in Figure 4. We will update line 147 and add a reference to the phase diagrams in Figure 4.\n\n**Equation 2**:\nFollowing the reviewer's suggestion, we will modify Equation (2) to read as follows: \n$$c_1 (x_1, \\\\; y_1 ) + c_2 (x_2, \\\\; y_2 ) = (x, \\\\; y) \\\\; \\mathrm{mod} \\\\; p\n\\qquad \\xrightarrow{\\text{find} \\\\; c_1, c_2} \\qquad \nz = c_1 z_1^t + c_2 z_2^t \\\\; \\mathrm{mod} \\\\; p \\\\; \\qquad\\qquad (2)$$\n\nFor completion, we provide an explanation of the algorithm here: The model finds the right way to \"linearly combine\" $(x_1, y_1)$ and $(x_2, y_2)$ to equal $(x,y)$. The re-scaling factors $c_1, c_2$ in this linear combination can then be used to get the correct answer: $z = c_1 z_1^t + c_2 z_2^t$. Here is an intuitive way to think about the algorithm: Instead of directly solving for the unknowns $(a, b)$, the model treats the in-context examples $(x_1, y_1, z_1)$ and $(x_2, y_2, z_2)$ as vectors. The task then becomes to fill in the last component of a new vector $(x, y, ?)$. By aligning the first two components using coefficients $c_1$ and $c_2$: $c_1 (x_1, y_1) + c_2 (x_2, y_2) = (x, y)$ mod $p$, the model naturally finds out the correct way of aligning $z_1$ and $z_2$ with $?$, given the fact that they are constructed with a same underlying linear relation $a x + b y$ mod $p$.\n\nWe will add a version of this explanation to the final version of the paper to enhance clarity.\n\n**Galois Field**:\nWe thank the reviewer for pointing out the missing definition of Galois Field. We will include that in the final version.\n\n**Figure 6 and attention heads**:\nThe results in Figure 6 are for $d=2$. We will specify that in the caption as well as the main text in the final version.\n\nIn the Appendix of the final version of the paper, we will extend Figure 11 to include all the attention heads from the $d=4$ model. Unfortunately we cannot show them in the Global Rebuttal PDF due to space constraints.\n\n**Task dependence of Figure 6(b)**:\nIn Figure G.2 of the attached PDF (Global Rebuttal) we present the PCA of attention heads for multiple tasks. Taking a close look at the first column in the layer 2, head 2 case, we see that the attention pattern is different for the two different tasks. \n% (Note that the layer 2, head 2 plots may seem denser than Figure 6(b). This is because we have included *all* the points here -- in Figure 6(b) we had only included points with even $x$ values, to keep the figure clean and interpretable.)\n\n**Re-scaling coefficients $\\mathbf{c_1, c_2}$**: \nWe tried using linear probing to extract information of $c_1, c_2$ from the residual stream, but the result is inconclusive. Please see the Global Rebuttal for more details.\n\n**Compositional and systematic generalization**:\nWe thank the reviewer for pointing out the relevant references. We will include their relation to our work in the final version of the paper.\n\n[1] Wei et al.; \"Emergent Abilities of Large Language Models\"; arXiv:2206.07682 (2022)\n\n[2] Kumar et al.; \"Grokking as the transition from lazy to rich training dynamics\"; ICLR 2024", "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:
# Global Rebuttal

We included three new figures in the attached one-page PDF. These new results address questions raised by one or more of the reviewers. Especially, the results about MLP layers are relevant to multiple reviews.

## Figure G.1
We examined how individual neurons (post-ReLU) are activated for different inputs $(x,y)$. We discovered that each neuron only gets activated for highly specific inputs. This can be interpreted as a further skill composition at the neuron level, although the exact role of each neuron remains to be discovered.

Notably, in Layer 3 of the $d=4$ model we find that neuronal activations follow the re-scaling relation $x = k y$. The layer contains all such re-scalings, forming a complete basis. Layer 2 of the $d=4$ model show a periodic pattern wrt $(x,y)$, while Layer 1 neurons only get activated only for specific $x$ values.

Neurons in the $d=2$ model appear to be superposed/compressed versions of those found in the $d=4$ model. This is likely due to $d=2$ model not having enough layers. We observe that the neurons from Layer 1 of the $d=2$ model contains patters similar to Layer 1 and 2 of the $d=4$ model. Neurons from Layer 2 of the $d=2$ model appear to be superpositions of various re-scaling patterns from Layer 3 of $d=4$ model.

## Figure G.2
We expanded upon the PCA plots of the $d=2$ model presented in Figure 6 of the original manuscript. In this extension, we included a different task vector $(a, b)$ and plotted the results using a different shot (16-shot).

We see that The top-3,4 components of PCA of layer 1, head 3 also forms circle of circles, albeit with a different pattern from that of top-1,2 components. In this case, we find pairs of coinciding circles, where the $x$ corresponding to the coinciding circles differ by $(p-1)/2 = 14$.

In layer 2, head 2 we can see that the PCA pattern changes for different tasks. This is in contrast to the layer 1 PCA patterns, which remain unchanged. (as claimed in the main text)

## Figure G.3
We performed PCA on both (i) Attention output (as opposed to individual heads) and (ii) MLP output of the $d=2$ model. In the Attention output of Layer 1, we observe a circle-of-circle structure. The other components also exhibit some structure -- notably, the top-3,4 components in layer 1 form 4 clusters corresponding to even/odd $X$ and $y$ values.

## Additional Experiments

In addition to these results, we also ran a few more experiments. Due to space constraints, we could not include them in the one-page PDF. We describe these experiments and results in words here -- we will include them in the camera-ready version of the paper.

1. Linear probing of $c_1$ and $c_2$: We extracted the feature from the residual stream after each transformer block; attached a new linear layer and fine-tuned it to predict the correct re-scaling coefficients. 

    We simplified the experiment to the $1$-shot case, where the sequences are simply $(x_1, y_1, z_1, x, y, ?)$. The fine-tuned linear layer was used to predict the correct coefficients $c_1$ such that $x_1 \cdot c_1 = x$ mod $p$ (alternatively $y_1 \cdot c_1 = y$ can be used). We observed $15 \\% - 20 \\%$ accuracy across all the layers for both $d=2$ and $d=4$ models. Despite the accuracy being above random guessing ($3\%$), we believe this result to be inconclusive.

2. Linear probing of (a,b): We also ran similar experiments with full sequences ($1$-shot to $31$-shots) and tried to predict the task vector $(a, b)$ along the sequence. We found random performance on the o.o.d. generalization set, suggesting that the model does not explicitly compute the task vector. Note that this is in agreement with our proposed algorithm.

3. Pruning / Activation Patching [1]: In this experiment, we replaced the output of each attention head with its averaged output over pre-training sequences. The average was taken over all the pre-training tasks, and $512$ sequences from each task. We found that:
    
    - For both $d=2$ and $ d=4$ models, pruning the circle-of-circle head immediately brings the model to random guessing. This can be understood as an average over sequences collapsing the circle down to a point, which destroys the feature completely.

    - For the $d=2$ model, pruning any other head causes some performance drop. This is to be expected since the model does not have enough capacity even before patching.

    - The $d=4$ model is significantly more robust to pruning. We can patch all heads except for (i) the three shown in Figure 11 of the manuscript and (ii) the heads in the last layer, with almost no impact on the performance (less than a $5\\%$ performance drop).

[1] Fred Zhang, Neel Nanda; "Towards Best Practices of Activation Patching in Language Models: Metrics and Methods"; https://arxiv.org/abs/2309.16042

Meta-review: This paper studies the emergence of in-context learning and skill composition in Transformers. All reviewers were enthusiastic about the work and highlighted its novelty and generalization of prior work studying modular addition and in-context linear regression. The paper was praised for its impact on deep learning and its meaningful contribution to understanding emerging capabilities. Reviewers found the proposed synthetic problem elegant and a candidate for test-best for follow-up work. The experiments' design was appreciated because they underpinned the paper's claims and contributed to cataloging the emergent phenomena. Finally, the clarity of exposition and visualization was acknowledged, particularly due to some complex analysis elements. This translated to the reviewers' scores all exceeding 7 and a unanimous recommendation of acceptance. 

This is by far the highest-rated paper from my batch (12 items), and due to such enthusiastic reviews, I recommend it to be considered for oral. 

EDIT after SAC feedback: I highlighted the things that the reviewers praised.
Decision: Accept (oral)