Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks

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.

Strengths

- The paper is easy to follow. Good presentation! - The experiments are well-designed, providing compelling support for the claims. - The results in Figure 5 are cool. - 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. - 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

- 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

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Discussions on how the findings help the practical system.

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.

Strengths

I thank the authors for submitting their excellent work which stands to have a substantial impact in the science of deep learning. * The work makes a meaningful contribution to an exceptionally important and interesting topic of the emergence of capabilities and internal mechanisms in deep learning. * The setting and experiments neatly isolate and clearly demonstrate several interesting phenomena of emergence of capabilities and shifting in the solutions found by deep networks throughout training, contributing to the field's developing catalogue of examples of these phenomena. * Moreover, the proposed synthetic problem is both rich and elegant. I expect this framework will become a fruitful test-best for follow-up work studying emergence phenomena, helping the field to improve our empirical and theoretical understanding of these phenomena. * The authors also offer a partial behavioural and mechanistic analysis which is a solid starting point for a more detailed understanding of the learned structures that emerge in this setting. * While some elements of the analysis are complex, the authors have done an exceptional job of clearly presenting their findings. I feel careful study of each section and figure in the main text was rewarded since there was no question that occurred to me that was not addressed in the authors' clear descriptions or figures. * 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 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

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.

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors transparently acknowledge all of the limitations I was able to identify within the paper itself.

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.

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

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?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

As acknowledged by the authors, this work is limited to particular algorithmic datasets.

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

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

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.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are fairly discussed.

Thanks for clarifying and for your proposed improvements to an already strong paper

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.