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_cc_2024_Conference
Published: 2024-05-15 21:42:34

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.

Main paper: 
 introduction: 
 Large language models (LLMs) can perform simple tasks absent from their training data. This ability is 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 is usually attributed to the following mechanism. First, LLMs learn a large variety of simple tasks and, then, how to compose those skills to form very complex skills [3] . Furthermore, LLMs also exhibit "emergent capabilities" -a sudden emergence of a new complex skill as a function 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 on algorithmic tasks [24] , which results from the model learning very structured representations [18, 11, 20] . Once these representations emerge, the model abruptly learns how to perform the task.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 [1, 25] , we introduce a finite collection of discrete modular arithmetic tasks [24] generalized to the in-context learning setting. Each task corresponds to learning a linear function z = a x + b y mod p over Z p from the examples provided in context of the autoregressive model(AM). In the bi-variate case there are p 2 such functions labeled by the vector (a, b). The main objective of this algorithmic dataset is to probe how AM utilizes the tasks it has learnt during training to solve the new tasks. 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, as it was shown in [1] , a single linear attention layer can solve the regression problem, while adding extra layers and non-linearities slightly modifies 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. 1 . Our main findings are: (i) there are four different phases in the end of pre-training depending on the number of tasks, n i.d. , and number of examples per task, α. (ii) At inference time, there is a generalization transition in the number of few-shot examples, as the 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 a handful of tasks and memorize the training set. To avoid such minima we make sure that every batch contains equal number of tasks (meaning that no tasks is over-or under-represented in 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 2 contains the literature review. In Section 3 we explain our notations and discuss the experimental details. In Section 4 we demonstrate empirically that the out-of-distribution ICL ability emerges as the number of training tasks increases. We also study the effects of model depth and task difficulty. In Section 5 we carefully examine a 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] . 
related works: 
 In-Context Learning (ICL) Brown et al. [5] first demonstrated that large models performance improves substantially when a few examples of the task at hand are provided at inference time, in the prompt. Akyürek et al. [2] , Ahn et al. [1] , 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. [10] demonstrated that transformers can learn several simple classes of functions in context. Kirsch et 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 out-of-distribution ICL ability on linear regression tasks. Lin and Lee [16] 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 to predict the correct results. Boix-Adserà et al. [4] showed theoretically and experimentally that with enough pre-training data, a transformer model can perform abstract reasoning that a MLP cannot do. Guo et al. [13] showed that transformers can use lower layers to memorize and upper layers to perform ICL in a feature regression setting. It was found in [26] that ICL is a 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. [24] 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. [12] 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 et al. [7] showed that MLP and transformer models can solve a specific family of modular polynomial tasks by bijectively mapping them to modular addition tasks.Interpretability Elhage et al. [8] , Olsson et 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. [21] 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: z t i = a t x i + b t y i mod p. We will refer to the coefficients (a t , b t ) ∈ Z 2 p as the task vector. The superscript t ∈ {1, • • • , p 2 } labels the p 2 possible tasks. We will refer to (x i , y i ) ∈ Z 2 p as the input vector, which is labeled by the subscript i ∈ {1, • • • , p 2 }.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 d embed . Each number is tokenized as an independent token. The pre-training is done following a slightly modified next-token prediction setup, with sequences of the form: 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.EQUATIONDetailed discussions on 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 d embed = 512, with n ctx = 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. 
emergence of in-context learning and task composition: 
 In this section, we demonstrate that a transformer model with depth d ≥ 2 can develop ICL and out-of-distribution generalization on 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 out-of-distribution is predominantly determined by the number of pretraining tasks n i.d. .In Figure 3 Shifting our focus to the last two panels of Figure 3 (a), we see that when n i.d. ≥ 256, the model can solve new tasks that were absent in the training set. Notably, there appears to be a 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-training performance simultaneously degrades. This phenomenon indicates a shift in the algorithm implemented by the model. Prior to 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 a more universal approach to solve the task in-context. We emphasize, that the model learns to perform ICL 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 as one of the seen tasks or as unknown. Once it matches the sequence to one of the 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 a method of determining the task vector algorithmically at inference time. In that case model performs equally well on seen and un-seen tasks alike. In fact, the small two-layer model we study has such a low capacity that it entirely skips the in-distribution generalization phase and immediately jumps from pure memorization to out-of-distribution generalization.Next, to further illustrate the effect of task diversity, we plot the pre-training accuracy (set S i. test 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 one analyzed in Singh et al. [26] , albeit with a different setup. test for α = 0.6.Further evidence supporting the two-circuit competition can be observed in panel (d) of Figure 3 . The loss curves show a "monotonic → non-monotonic → monotonic" transition as the task diversity increases.  
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 i s.t. c i (x i , y i ) = (x, y) mod p then z = c i z t i mod p. This algorithm only works when y/x = y i /x i mod p holds for at least one of the in-context examples.Modular Regression: Find {c 1 , c 2 , . . . , c k } s.t. k i=1 c i (x i , y i ) = (x, y) mod p. Then the prediction is z = k i=1 c i z i 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 (x 1 y 1 z 1 x y ?). In this case, Ratio Matching will only solve the task for inputs that obey (x, y) = c 1 (x 1 , y 1 ) mod p for some c 1 . Indeed, this is exactly what we observe in Figure 5 To implement these algorithms, the model needs to perform linear operations over GF(p), which can be broken down into the following essential skills (in addition to copying information, which is readily implemented by attention heads). In the remainder of this section, we show the special attention heads that implement skill I (Section 5.1) and MLPs that implement skills II, III (Section 5.2). We explicitly show the deterioration in this (6, 6) , loop over (x i , y i ) at a specific shot while the previous parts of the sequence are fixed. We annotate each PCA plot with the (log 27 x i , log 27 y i ) mod p pairs. (a) The principal components form a circle of circles where the position of the outer circle is controlled by x i . This pattern remains the same for different task vectors or the shot choices; (b) Only plotted pairs with even log 27 x i , with each log 27 x i circled with different colors. The PCA pattern forms a similar double circle as 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. structures as pre-training task diversity (n i.d. ) decreases. We also elucidate the algorithmic shift in the d = 4 models as the number of in-context examples increases.So far we have presented "black-box" evidence suggesting that the model is implementing the proposed algorithms (Figure 5 ). Now, we turn to "white-box" interpretability, to identify the components within the transformer model that are responsible for the essential skills. In Figure 6(a, b ), we analyze the important attention heads in the d = 2 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 1 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 x i , y i , z t i ; 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 a batched k-shot sequence of the form (x1 y 1 z 1 • • • x k y k z k x y z),where the first k inputs are fixed and the last input (x, y) is scanned over all p 2 possible pairs. We concatenate the resulting features from x and y, resulting in p 2 × 2d embed batch of features -and perform PCA on this matrix. We project all these 2d embed dimensional features onto the first two principal components. Annotating the pairs (x, y) with (log 27 x, log 27 y) 1 , we find a "circle-ofcircles" -where circles of period 28 are arranged in a bigger circle of period 28. Number 0 is located at the center of the circles2 . We observe similar circle-of-circles for concatenated features from (x, z) and (y, z) as well. We refer the reader to Appendix E for further details.EQUATIONIn Figure 6 (b), we 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 (x, 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 (log 27 x, log 27 y) format. However, unlike the previously discussed pattern, the specifics of this "circle-of-circles" arrangement vary depending on the position and the choice of task vector (a t , b t ). This variability suggests that the head in question possesses information about the specific task from the context.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 n i.d. decreases, the attention map starts to get mosaicked (top panels); and the PCA projections lose their shape (bottom panels). As a result, these models lose the ability to perform ICL on modular arithmetic out-of-distribution.After applying Skill I, the model maps numbers onto circular feature spaces, enabling discrete modular operations such as multiplication, division, addition, and subtraction through subsequent model components. To analyze this behavior, we compute the cosine-similarity between layer l (i.e. l th Transformer block) k-shot output vectors (standardized) at token position y: We evaluate this similarity metric across all possible input pairs (x, y) and (x ′ , y ′ ), resulting in a p 2 × p 2 matrix.EQUATIONWe also repeat the analysis to layer output vectors at z tokens h l z (•). In line with our previous analysis, we use controlled (fixed) in-context examples(x 1 y 1 z 1 . . . x k y k z k ).For the d = 4 model, the top panel of Figure 7 reveals a distinctive kaleidoscopic pattern3 of high cosinesimilarities. The pattern corresponds to input pairs (x, y) that share the same ratio y/x mod p. In the MLP modules, we find highly structured activations as functions of inputs (x, y) (Figure 8 left ). Together with our earlier analysis of Skill I, these observations suggest that the MLP layer following the attention module with a "circle-of-circles" head (layer 2, head 4 for d = 4 model, see Figure 13 (a) ) implements division operations over GF(p). The bottom panel of Figure 7 succinctly 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/x matches one of the ratios (y 1 /x 1 or y 2 /x 2 ). At higher shots, we see a transition to a task-dependent pattern where features align for input pairs (x, y) for which a x + b y mod p match -a signature of Modular Regression.For d = 2 models, we again observe the kaleidoscopic pattern in cosine-similarities (Figure 7 topright) and structured MLP activations (Figure 8 right). However, in this case we only find signatures of Ratio Matching in the cosine similarities (Figure 7 bottom-right), as expected.Thus, we conclude that the d = 2 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 a more detailed discussion, see Appendix E. 
discussion: 
 We have investigated the emergence of in-context learning and skill composition in autoregressive models on a novel algorithmic dataset. The dataset includes a large discrete set of modular arithmetic tasks and is specifically designed to force models to learn how to solve a 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 qualitatively 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 in 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.Architecture We used GPT-like architectures with Rotary Positional Embedding (θ = 10, 000) and ReLU activations. We fix the number of heads H = 4, embedding dimension d embed = 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 N (0, 0.02 2 ) at initialization, with the exception that the last linear layer in each MLP is sampled from N (0, 0.02 2 /2d). No bias is used in any layer.We trained most models using AdamW optimizer with learning rate η = 1.5 × 10 -4 , weight decay λ = 2.0, β 1 = 0.9, β 2 = 0.98, ϵ = 10 -8 , batch size B = 1024, n ct = 32 in-context examples for 200k steps, together with a 5% linear warmup starting from 0.01η and a 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.5 × 10 -5 , 1.5 × 10 -4 , 3 × 10 -4 , 6 × 10 -4 } and weight decay values λ ∈ {0.5, 1.0, 2.0, 5.0}. Then we transfer our hyperparameters to 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 -4 . test . 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. 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 A100 40GB GPU, of which the cost is negligible compared to the pre-training cost.C Table of 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 at a 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 are selected randomly, the resultant noise appears to inhibit the learning capabilities of the model, preventing it from acquiring any meaningful patterns or rules from the data.Ultimately, we adopted a more structured approach to task selection as illustrated in Figure Figure 2 , based on the following rationale: If the model is to learn a specific task vector (a t , b t ), it would presumably be more straightforward if one of the two components -either a t or b t -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 9 , we show an ablation plot of phase diagrams with a randomly selected task vector collection for pre-training. We still use the special sequence design. Sequence Design Following a similar spirit, we use a balanced batch where sequences generated from all task vectors appear exactly the same number of times during the training. We further align the examples across sequences generated by different task vectors, which we believe reduces the chance of the model getting confused by the same input (x, y) appearing at different positions within the batch. Without this design, we 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. 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 E.1), other attention heads (Appendix E.2), 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.We begin the discussion with the role of embedding layers. By further examining different models, we find highly structured embedding from d = 2 models. d = 4 models, on the other hand, does not have such structured embedding layers.First we focus the embedding of d = 2 models. As shown in Figure 11 , clearly, the logarithm of each number is split into even and odd groups, and each group forms one circle, which is a suitable embedding for doing modular multiplications. However, one should note that this will not obviate the importance of the head shown in Figure 6 (a), as the model still needs a 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 = 4 model, as shown in Figure 12 .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 a similar "circle-of-circles" feature as we shown in Figure 6 (a, b ).To continue the story, we analyse the attention heads in different models. We first study the d = 4 model, where we also find a similar head with "circle-of-circles" (Figure 6 (a)). Moreover, two other 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, it also creates a circle-of-circles while performing PCA over concatenated (x, y) features. The difference here is that the circle winds twice to go back to 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 in Figure 6(a) . Overall, we do not yet know why the model needs to split the logarithm of numbers into even and odd groups.From Figure 13(b, c ), we find two heads that are structured but different from Figure 6(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.Next we focus on the d = 2 model. In Figure 14 , we plot similar PCA to the one in Figure 6 (a), but with different concatenated features from the same heads: (x, z) and (y, z). We again see circle-ofcircles, 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 d = 2 models. The corresopnding PCA over concatenated (x, y) features are shown in Figure 16 (except for layer 1 head 4, for which we plot for (x, z)). The special head we choose here has a similar behavior to the one in Figure 6 (a), with its main focus shifted by one position but still within two preceding tokens. 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 , we extend our PCA analysis to different task vectors and up to 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, we leave a careful exploration of this front for future work.In Figure 18 , we plot d = 2 PCA analysis up to top-4 components for attention and MLP outputs. As we mentioned in Section 5.2, Skills II and III are likely implemented within the MLP layer. While we provided extensive details on how the signature of Skill II was observed in MLP, we could not find clear signals similar to those in Nanda et al. [20] , Gromov [11] from MLP layers on Skill III.Here, we show an extended version of Figure 7 for all layers from the d = 4 model. We plot for 2-shot (Figure 20 ) and 16-shot (Figure 21 ), the task vector used is (a, b) = (2, 6). (e) z, layer 1 (x = 1, y) (x = 2, y) (x = 3, y) (x = 4, y) (x = 1, y) (x = 2,We plot some selected training curves for d = 4 (Figure 25 ) and d = 2 (Figure 26 ) from Figure 4 phase diagrams. We see that even for d = 4, ICL can be a transient. With increased α or n i.d. , the transient nature goes away. • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It 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 a 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 of range (e.g. negative error rates). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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. It does no impact the society at large, beyond improving our understanding of certain aspects of deep learning. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: Since we train small models on toy-datasets, our work poses no risk of misuse.Guidelines:• The answer NA means that the paper poses no 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 a best faith effort.12. Licenses for existing assets Question: 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?Answer: [Yes] Justification: We are not shipping our code with any source code or binary files from any other existing libraries, so there are no concerns over getting permission or including a 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 or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.We useas the base of logarithm, which is a primitive root of GF(29)This is expected since log 27 0 is mathematically ill-defined and needs special treatment.Patterns are cropped due to file size limit, check the GitHub repo for the full kaleidoscopic pattern. 
 
 Appendix:  


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

Reviews: review: paper_summary: paper_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).
strength_weakness: strength_weakness: I do not have any major apprehensions, except for the related work, which I think is relatively sparse.

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

[1] https://arxiv.org/abs/2305.16704
[2] https://arxiv.org/abs/2311.12997 
[3] https://arxiv.org/abs/2310.06110
[4] https://openreview.net/forum?id=aN4Jf6Cx69
questions: questions: A few questions below that I would like to see answered. 

- **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?

- **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: Limitations are fairly discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
- 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.
questions: 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: limitations: As acknowledged by the authors, this work is limited to particular algorithmic datasets.
paper_summary: paper_summary: * The authors propose a synthetic sequence learning problem that I would call
  'in-context modular regression', an elegant generalisation of prior work
  studying modular addition and in-context linear regression.
* Using carefully constructed batches the authors are able to train
  transformer models to perform regression for a subset of tasks (weights)
  and a subset of inputs.
* The authors show that under some conditions on the data distribution and
  model architecture, the transformers not only achieve good performance on
  tasks and inputs included during training, but they also generalise to new
  tasks and/or new inputs. The authors document the conditions governing
  these generalisation capabilities in detail including showing phase plots
  and observing that in larger models, the generalisation properties are
  transient (they appear and then disappear across the training process).
* The authors postulate a breakdown of skills required to correctly perform
  the task. They effectively isolate and examine the abilities of their
  models to perform each component task. They also inspect the activations
  of each head and identify patterns suggestive of partial mechanisms
  underlying the generalising behaviour of the models.
strength_weakness: strength_weakness: I have not noticed any weaknesses in the paper that would temper my overall
recommendation to accept. However, I note the following weaknesses, some of
which the authors have already acknowledged, and others which they may like
to take into consideration if they are interested to improve the paper
further.

1. **Delicate training set-up.** The authors explain that training
   transformers on multiple modular addition tasks crucially relies on
   following a delicately balanced batch construction methodology.
   I am left wondering if this batch construction methodology, as a further
   departure from the standard language modelling setting, has any other
   implications for the learning process that may affect the generality of
   the results.
   Note: This weakness is not decisive because the authors clearly document
   their training methodology and it's not *that* artificial anyway.

2. **The mechanistic analysis is only partial.** The authors admit that they
   have not been able to identify an end-to-end mechanistic model of how the
   trained transformers perform the task. This leaves their posited skill
   decomposition and partial mechanistic analysis open to the possibility
   that they are incomplete.
   Note: I think the contribution the authors have given in terms of the
   setting, the generalisation phenomena, and the partial skill decomposition
   and mechanistic analysis are already significant.

3. **Relationship to prior work.** The related work section does a good job
   of summarising the contributions of prior work in in-context linear
   regression and modular arithmetic in the context of transformer models.
   However, I feel that this section could be improved if the authors
   attempted to offer greater insight into the relationship between these
   prior works and the present work. For example, the authors have an
   opportunity here to informally describe the in-context linear regression
   and the modular addition problem settings that the newly proposed setting
   generalises.

4. I noticed some minor text errors as follows, which I expect the authors
   can easily correct.

    * Line 94: The notation $[1, p^2]$ to me suggests a closed continuous
      interval, whereas you appear to mean $\lbrace1, \ldots, p^2\rbrace$, also in some
      cases denoted $[p^2]$.
    * It seems that equation 2 should read $\ldots = (z_1^t, z_2^2) \mod p$
      and the equation on line 203 should read $c_1x + c_2y \mod p$. That is,
      $x$ and $y$ should swap places with $z_1^t$ and $z_2^t$. Is this indeed
      a mistake, or am I missing something?
    * In figure 6 (top row) there is a typo: "Qeury" on the vertical axis.
    * In line 445 there is a broken link.

I have not studied all appendices in detail.
questions: questions: 1. Why is the title 'learning to grok'?

    * Is this meant in the sense that the grokking of a modular addition task
      is occurring in-context? If so, this seems a little inaccurate, since
      the phenomenon analogous to 'grokking' seems to still be occurring
      during pre-training.
    * To be honest this part of the title has puzzled me since I first looked
      at the paper. Even if my understanding above is wrong and the title has
      an accurate interpretation, that I have failed to notice it might be
      one data point suggesting that if you are going for a title that is
      both short *and* informative, this might not be the right choice.

2. In the figure 1 caption, is it possible to offer a clearer summary of the
   difference between in-distribution generalisation and out-of-distribution
   memorisation? On my first read through, treating the figure and caption as
   an overview of the work's main results, I had trouble distinguishing these
   two concepts.
limitations: limitations: The authors transparently acknowledge all of the limitations I was able to
identify within the paper itself.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
questions: 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.
- 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.)
- Better to define $GF(p)$, i.e., the Galois field, before using it.
- 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.
- In line 264, the paper claims that the pattern depends on $(a,b)$, but it is hard to read that from Figure 6b.
- 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.
- 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:
    
    [1]  Schott, Lukas, et al. "Visual representation learning does not generalize strongly within the same domain." ICLR 2022
    
    [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
    
    [3] Ren, Yi, et al. "Improving compositional generalization using iterated learning and simplicial embeddings." NeurIPS 2023
limitations: limitations: Discussions on how the findings help the practical system.

Comments: title: Thanks for clarifying and for your proposed improvements to an already strong paper
comment: Thank you for the detailed response. I have increased my score 6 --> 7.
Thank you for clarifying especially my confusion around the proposed vector scaling approach to solving the task. The proposed revisions and the additional experiments will further improve an already strong paper. I maintain my confident recommendation that this paper should be accepted.
Thanks very much for the author's response. The new results are quite interesting. All of my concerns are well resolved. I confirm my evaluation and hope to see its new version.
Thank you for the response. I really like the new results and hope they'll be included in the final paper. I'll maintain my original score.