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

In Figure 3(a), we show the accuracy of $d=6$ models vs the number of training tasks $n_{\mathrm{i.d.}}$ and the number of few-shot examples quantified by the fraction of the total number of data points, $\alpha$; on sets $S^{\mathrm{i.d.}}_{\mathrm{train}}$, $S^{\mathrm{i.d.}}_{\mathrm{test}}$, $S^{\mathrm{o.o.d.}}_{\mathrm{train}}$ and $S^{\mathrm{o.o.d.}}_{\mathrm{test}}$. The phase diagram in Figure 1 is constructed by merging the last shot accuracy version of these four diagrams shown in Figure 18(a) of Appendix G.

The ability of the model to generalize in-distribution increases with $\alpha$, as can be seen by comparing the first two panels of Figure 3(a). This behavior is in correspondence with the original work on grokking, where the transition to generalizing solution is driven by the amount of data. Further, we observe that an increase in $n_{\mathrm{i.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 a task-memorizing solution to the one that generalizes out-of-distribution. The model switches to a new algorithmic way of solving the task and the solution is more few-shot-sample-efficient.

Shifting our focus to the last two panels of Figure 3(a), we see that when $n_{\mathrm{i.d.}}\geq 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^{\mathrm{i.d.}}_{\mathrm{train}}$) and the o.o.d. test accuracy (set $S^{\mathrm{o.o.d.}}_{\mathrm{test}}$) as a function of training steps (Figure 3(b, c)); for various values of $n_{\mathrm{i.d.}}$. We observe a clear memorization-to-generalization transition as task diversity increases. Interestingly, for $n_{\mathrm{i.d.}}=2^{8}$, the ICL ability on $S^{\mathrm{o.o.d.}}_{\mathrm{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. [25], albeit with a different setup.

Further evidence supporting the two-circuit competition can be observed in panel (d) of Figure 3. The loss curves show a “monotonic $\rightarrow$ non-monotonic $\rightarrow$ monotonic" transition as the task diversity increases. With a minimal number of pre-training tasks, the model primarily engages in memorization, resulting in 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 $n_{\mathrm{i.d.}}=2^{8}$ 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 a monotonic loss curve. Similar insights can be derived from examining the monotonicity of the accuracy curves in panel (e).

A natural question to ask is if similar phenomena can be observed with different model sizes or task difficulties. Here we present our results with $d=4$ and $d=2$ in Figure 4 and leave the results for other prime $p$ values in Appendix H.

When comparing phase diagrams in Figure 3 with Figure 4, we observe that those phase diagrams across different depths are qualitatively similar, where the o.o.d. generalization only emerges with a large enough number of 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 the pre-training accuracy falls drastically as the 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 a function of context length at 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_{\mathrm{i.d.}}=2^{9}$, there is no loss surge in the $1$-shot case across all three depths. Furthermore, the $d=4$ model with $n_{\mathrm{i.d.}}=2^{8}$ behaves significantly differently from the corresponding one with $d=2$ case, where the model fails to perform ICL for the o.o.d. test set. This is also distinct from the $d=6$ case, where the model tends heavily toward memorization due to its excessive capacity. Instead, the $d=4$ model manages to maintain a better balance between memorization and generalization at the end of pre-training. Consequently, the model has a $1$-shot loss surge followed by a notable drop in ICL loss. This suggests that $d=4$ optimally leverages the available model capacity to facilitate effective learning dynamics for o.o.d. generalization.

We find telling signatures of Equation 2 upon analyzing model predictions with varying numbers of examples in context. Consider the 1-shot scenario, wherein only one demonstration is given to the model. This is equivalent to setting $c_{2}=0$ in Equation 2. In this case, one would expect that the model will only solve the task correctly for examples that obey $(x,y)=c_{1}(x_{1},y_{1})$ for some $c_{1}$. Indeed, this is exactly what we observe in fig. 5 – the model correctly predicts all the targets when the inputs $(x,y)$ are re-scalings of the in-context input $(x_{1},y_{1})$.

For few-shot scenarios, the model does two things. It can predict the targets correctly whenever the query input is proportional to any of the inputs of the previous demonstrations (Figure 5 row 1). Additionally and more importantly, it can correctly solve substantially more than just these cases (Figure 5 row 1,3). This results from the model’s ability to combine in-context examples and use them for prediction. Evidence for this can be seen in the proliferation of red points in row 3 of Figure 5 as more examples are given in-context. The $d=2$ model performs inferior to its $d=4$ counterpart in this latter skill, hinting at an imperfect implementation of the algorithm and the partially correct choice of the set of constants $\{c_{1},c_{2},\cdots\}$.

In the previous section, we found “black-box" evidence suggesting that the model is implementing an algorithm akin to solving equation eq. 2. Now, we turn to “white-box" interpretability, to identify the components within the transformer model that are responsible for the essential skills outlined at the beginning of Section 5. In figure Figure 6(a, b) we analyze the important attention heads in a model that generalizes o.o.d. We compare them with the heads from a model that does not generalize o.o.d in Figure 6(c, d).

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 $(x_{1}\;\;y_{1}\;\;{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}z_{1}}\;\;\cdots\;\;x_{k}\;\;y_{k}\;\;{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}z_{k}}\;\;x\;\;y\;\;{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}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}\times 2d_{\mathrm{embed}}$ batch of features – and perform PCA on this matrix. We project all these $2d_{\mathrm{embed}}$ dimensional features onto the first two principal components. Annotating the pairs $(x,y)$ with $(\log_{27}x,\log_{27}y)$²²2We use 27 as the base of logarithm, which is a primitive root of $\mathrm{GF}(29)$, we find a “Clock-of-clocks" – where clocks of period 28 are themselves arranged in a bigger clock of period 28. Number 0 is located at the center of the clocks³³3This is expected since $\log_{27}0$ is mathematically ill-defined and needs special treatment.. We observe similar clock-of-clocks for concatenated features from $(x,z)$ and $(y,z)$ as well. We refer the reader to Appendix E for further details.

In Figure 6(b), we analyze the head from layer 2 that effectively implements Skill II – it does so by building upon skill 1 and leveraging the features discussed above. The upper panel shows that the highly structured attention map that focuses on the current example as well as the examples preceding the current one. This pattern aligns with the step in our proposed algorithm where the model needs to compare $(x,y)$ pairs across different examples in-context; and selectively re-scale some of them for later combination. Additionally, as the number of in-context demonstrations increases, the attention map indicates that the model begins to focus on multiple examples simultaneously. This suggests that the model might be employing a potential weighted average approach over these examples, as discussed earlier.

By conducting a PCA analysis similar to that in Figure 6(a), we also identify clocks when annotating examples in the $(\log_{27}x,\log_{27}y)$ format. Unlike the previously discussed pattern, the specifics of this "clock-of-clocks" arrangement vary depending on the position and the choice of task vector $(a^{t},b^{t})$. This variability suggests that the head in question is utilizing information about the specific task from the context to appropriately rescale the equations.

We further highlight the importance of the signatures we find in the above paragraphs via comparison with a model that does not generalize o.o.d, in Figure 6(c, d). Note that as the number of pre-training tasks $n_{\mathrm{i.d.}}$ decreases, the attention map similar to the one in panel (a) starts to be mosaicked; and at the same time the PCA projections lose their shape. As a result, those models also lose the ability to perform ICL on modular arithmetic out-of-distribution.

From Figure 5, we infer that the $d=2$ model has only partially 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.