SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Aaditya K. Singh*
Stephanie C.Y. Chan*
Gatsby Unit & SWC, UCL
Andrew M. Saxe†
Felix Hill†
Transformer neural networks can exhibit a surprising capacity for in-context learning (ICL) despite not being explicitly trained for it. Prior work has provided a deeper understanding of how ICL emerges in transformers, e.g., through the lens of mechanistic interpretability, Bayesian inference, or by examining the distributional properties of training data. However, in each of these cases, ICL is treated largely as a persistent phenomenon; namely, once ICL emerges, it is assumed to persist asymptotically. Here, we show that the emergence of ICL during transformer training is, in fact, often transient. We train transformers on synthetic data designed so that both ICL and in-weights learning (IWL) strategies can lead to correct predictions. We find that ICL first emerges, then disappears and gives way to IWL, all while the training loss decreases, indicating an asymptotic preference for IWL. The transient nature of ICL is observed in transformers across a range of model sizes and datasets, raising the question of how much to "overtrain" transformers when seeking compact, cheaper-to-run models. We find that L2 regularization may offer a path to more persistent ICL that removes the need for early stopping based on ICL-style validation tasks. Finally, we present initial evidence that ICL transience may be caused by competition between ICL and IWL circuits.
In-context learning (ICL) is the ability of a model to use inputs at inference time to adapt its behavior, without weight updates, in order to solve tasks not present during training. The first examples of this behavior in neural networks were observed in architectures specifically designed and trained for few-shot learning, the capacity to learn a desired behavior from only a few examples [1-3]. Training in these cases involved shuffling the labels corresponding to input exemplars on each "episode" so that, to perform well on the training set, the model had to recall exemplar-label mappings from context to make future predictions. At test time, novel exemplar-label mappings were provided, and the network had to use these to classify query exemplars.
Research on ICL changed with the development of the transformer [5]. Brown et al. [6] first observed that a transformer-based language model, GPT-3, trained auto-regressively at sufficient scale, exhibited ICL without any specific effort of the authors to promote it via the training objective or data. Since then, a large body of literature has studied [7-12] or reported [13-15] examples of ICL in transformers trained at scale.
*Co-first authors; direct correspondence to auditya.singh.21@ucl.ac.uk and scychan@google.com
† Co-senior authors
Few-shot learning can also be implemented via methods that apply gradient updates [e.g., 4], but we focus on implementations of ICL that do not require any weight updates.
(a) In-context learning.
(b) In-weights learning.
(c) Train log loss.
Figure 1: In-context learning is transient, shown for our "default" settings: 12 layers, embedding dimension of 64, trained on 1,600 classes, with 20 exemplars per class. All training sequences are bursty (see Figure 2a for details). Chan et al. [10] found these settings to strongly incentivize ICL, but did not observe ICL transience (see Figure 3), as they did not train long enough. (a) ICL evaluator accuracy. (b) IWL evaluator accuracy. We note that, while accuracy on train sequences is  $100\%$ , accuracy on the IWL evaluator is very slowly increasing, as the test sequences are out-of-distribution. See Appendix B for further investigation. (c) Training log loss. Two colors indicate two seeds used for experiments.
These compelling findings have sparked a surge of research into the phenomenon of emergent capabilities in large neural networks. Nonetheless, recent work has shown that ICL is not a guaranteed outcome of training transformers. Chan et al. [10] found that specific properties that are typical of language data, such as burstiness and its highly skewed distribution, play an important role in emergent ICL in transformers. When trained on data without these properties, Chan et al. [10] found that transformers tend to resort to in-weights learning (IWL). In the IWL regime, the transformer relies on information stored in the weights of the model, rather than new information provided in-context.<sup>3</sup> Importantly, ICL and IWL often appear to be in opposition, with ICL emerging most readily when training data has a large number of tokens or classes and when that data is bursty, with items appearing in clusters rather than being uniformly distributed.
Such controlled studies based on known data-generating distributions are critical for better understanding the ICL phenomenon in transformers. At the same time, a complementary body of work studies emergence in massive models trained directly on organic web-scale data, with the takeaway that impressive capabilities such as ICL are more likely to emerge in large models trained on more data [14, 17-19]. However, this reliance on massive models poses serious practical challenges, such as how to innovate rapidly, how to train energy-efficiently and in low-resource settings, and how to run efficiently at deployment. Thus, a growing line of research [20-22] has focused on achieving comparable performance (including emergent ICL) in smaller transformer models.
Overtraining is currently the favored approach for training small yet performant transformers. These small models are trained on more data, possibly with repetition, beyond what is prescribed by scaling laws and compute budget [13, 15]. At its core, overtraining relies on an assumption that is implicit in
<sup>3</sup>We also note that ICL and IWL are not so cleanly distinguishable in many cases: ICL can depend at least partially on information stored in-weights, and behavior that looks like IWL may still fact using context naively, e.g., via simple copying biases [16].
(a) Example sequences and outputs.
Figure 2: An overview of our setup. (a) Example sequences during training, ICL evaluation, and IWL evaluation. Example outputs are colored green when correct and red when incorrect. Note that for train sequences, ICL and IWL strategies both result in the correct answer. On ICL eval sequences, the IWL prediction is incorrect. ICL is required to remap exemplars to the randomized 0 or 1 labels. On IWL eval sequences, there are no matching exemplar-label pairs in context, so IWL is necessary. (b) Model schematic. Training and evaluation focuses on the predicted label for the final exemplar.
(b) Model schematic.
most (if not all) recent explorations of ICL in LLMs: persistence. That is, it is assumed that once a given model has been trained sufficiently for an ICL-dependent capacity to emerge, that capacity will be retained as training progresses, provided the training loss continues to decrease.
Here, we show that the general assumption of persistence is false. We do so by adapting a standard image-based few-shot dataset [23], which allows us to evaluate ICL carefully under controlled conditions. We characterize simple cases where ICL emerges, only to be subsequently lost while the model's loss continues to decrease. In other words, while ICL is well-known to be an emergent phenomenon, we should also think of it as a potentially transient effect (Section 3, Figure 1). We find this transience occurs across a range of model sizes (Section 4.1), dataset sizes (Section 4.2), and dataset types (Section 4.3), though we do identify some properties (Section 4.4) that can postpone transience. In Section 5, we study ways to alleviate this risk, and show that regularization may offer a path to persistent ICL. In Section 6, we present preliminary evidence that ICL transience is caused by competition with IWL circuits.
In general, if networks are trained carelessly for longer and longer, we find that ICL can disappear as readily as it emerges, leaving models without capabilities that we are increasingly expecting to find in modern AI systems.
2 Experimental Setup
In order to study the transience of ICL across a range of model and dataset sizes in a controlled fashion, we use a synthetic data-generation process inspired by Chan et al. [10], which allows us to properly assess in-context versus in-weights learning. Our data generators are primarily built from the Omniglot dataset, a standard benchmark for few-shot learning, which consists of images of handwritten characters from different alphabets [23, MIT License]. The Omniglot dataset consists of 1623 character classes, with 20 exemplars each. In experiments with more than 1,600 classes, we construct additional image exemplars by applying rotations and flips of the original Omniglot images.
Inputs to the transformer model are sequences of exemplars and labels. We use 8 exemplar-label pairs (the "context") followed by one "query" exemplar, for an input sequence length of 17 (Figure 2a). The model is trained to minimize the loss in predicting the label for the query exemplar (Figure 2b).
To explore the generality of our results, we also extend to exemplars derived from language model token embeddings. We use pre-trained token embeddings from the LLaMa family [15] of large language models, which use a vocabulary size of 32,000 tokens. We extract the input embedding matrices from all four open-source LLaMa 1 models. We then subset and cluster [24] these token embeddings to form 3,200 classes of 5 tokens each. For each LLaMa model size, we construct two such datasets, with different clustering seeds. The full procedure is described in Appendix C.
(a) In-context learning.
(b) In-weights learning.
Figure 3: (Reproduced with permission from Chan et al. [10].) It was previously shown that ICL can be transient (purple and red curves) when transformers are trained on data that is only weakly conducive to ICL (e.g., with low levels of burstiness). P(bursty) indicates the fraction of training sequences that are bursty. All curves in this figure train on a dataset with 1,600 classes and 20 exemplars. Note the  $x$ -axis scale: At this number of iterations, there is no sign of ICL transience for the highest level of burstiness, but our experiments are run for a much larger number of steps. We also note that, when P(bursty)  $< 1$ , some of the sequences during training are of the same form as the IWL evaluation sequences. As a result, the IWL evaluator is less out-of-distribution, which is why the accuracies in (b), when P(bursty)  $< 1$ , are so high.
Training sequences. The training task (see top row in Figure 2a) can be solved via ICL or IWL. The context for each training sequence contains 3 exemplar-label pairs from the query class ("bursty" sequences), which prior work has found to incentivize ICL [10]). Bursty sequences allow the network to look for matching exemplars in-context to output a label for the query example, an example of ICL. To avoid repetition biases, the training context also includes 3 exemplar-label pairs from a distinct, distractor class. We keep exemplar-label mappings fixed across sequences, so the network could learn a mapping from exemplars to labels in-weights (without needing to find matching exemplars in-context). This design allows us to observe whether the model prefers ICL or IWL at different points in time, as a function of model size or in response to the properties of the training data.
Evaluation sequences. We follow Chan et al. [10] and construct two kinds of evaluation sequences that measure the extent to which a trained network relies on ICL versus IWL. To evaluate ICL (see second row of Figure 2a), we use sequences with 4 exemplar images from each of two classes, but we set the class labels to either 0 or 1 at random for each sequence (unlike during training, where the exemplar-label mappings are fixed). Accuracy on this evaluator is measured across 0 and 1 as possible outputs (so chance-level accuracy is  $50\%$ ). As these labels were not associated with these exemplars during training, the only way to achieve above-chance accuracy on this evaluator is for the network to refer back to the context and copy the label corresponding to matching exemplars. To evaluate IWL (see third row of Figure 2a), we use sequences where none of the context exemplars come from the same class as the query, but all of the exemplar-label mappings are the same as during training. In this case, ICL is not useful (as there are no matching exemplars in-context), so the network must rely on exemplar-label mappings stored in-weights (IWL) for above-chance accuracy. Other evaluation techniques for ICL and IWL are considered in Appendix B.
2.2 Model details
For our transformer model, the "default" settings consist of 12 layers, with an embedding dimension of 64 and an additive, sinusoidal positional encoding scheme [5]. Each element of the sequence is passed through an embedder, either a standard embedding layer for the label tokens, a ResNet encoder for the exemplar images, or a simple linear layer for exemplars created from LLaMa tokens. (Figure 2b). Both embedders are trained jointly with the network. In our experiments, we vary multiple aspects of this architecture and observe the effect on the transience of ICL. All experiments were run with 2 seeds and used the Adam optimizer [25]. We did not observe qualitative differences in behavior between seeds. We ran most experiments for 5e7 iterations, but truncated some early when signs of transience were clear. More details (e.g., specific hyperparameter choices, and links to code) can be found in Appendix A.
(b) In-weights learning.
Figure 4: (a-b) ICL is transient regardless of model depth, with no clear trend of peak height or peak onset. Decay slopes are roughly similar across model sizes.
3 In-context learning is transient
Figure 1a shows that ICL is transient even when using data-distributional settings previously found by Chan et al. [10] to strongly incentivize ICL. We annotate various aspects of the ICL curve, such as the peak height, peak onset, and decay slope. Prior work [10] stopped training at 5e5 training steps, before peak onset. Training for much longer in an identical setting leads to the disappearance of ICL. Notably, ICL is slowly replaced by IWL (as seen by the steadily increasing accuracy on the IWL evaluator, Figure 1b), all while the training loss continues to decrease (Figure 1c).
Chan et al. [10] found some transient ICL in other data regimes, such as when some small percentage of sequences are non-bursty, where the context does not contain the query (Figure 3). In these settings, the eventual rise of IWL is to be expected, as ICL cannot minimize the loss on the non-bursty sequences. Notably, our results (Figure 1) show surprisingly that ICL transience extends even to cases where ICL is "sufficient" to fully solve the training task (such as when the context always contains a relevant exemplar-label mapping that can be used). When comparing to the results from Chan et al. [10] (Figure 3), we see that higher levels of burstiness lead to a higher peak height and gentler decay slope. When training for fewer iterations, it might appear that ICL is persistent, but our results from training longer in Figure 1 demonstrate that ICL is actually transient.
4 In-context learning is transient across settings
4.1 Effect of model size
We investigate whether the ICL transience effect is modulated by model size, given the current prevalence of extremely large transformer models. We find that there is no consistent effect of model depth on ICL transience (Figure 4a-b). In fact, deeper models can lead to lower peak heights (e.g., the 18-layer model has a lower peak height than the 12-layer model). Notably, the decay slope is similar across different model depths, indicating that the ICL is unlikely to persist simply by scaling model depth.
We also experiment with the width of the model in Section 6, which sheds some light on the causes of ICL transience.
4.2 Scaling dataset size
Next, we investigate the effects of dataset size on ICL transience. There are two primary ways of varying the dataset size: increasing the number of exemplars per class (a form of increasing in-class variation), or increasing the number of classes.
Chan et al. [10] found that in-class variation was crucial for ICL to emerge, and we find a similar trend. When the number of exemplars per class is increased (from 10 to 15 to 20) while keeping the number of classes fixed, we see an increase in ICL peak height (Figure 5a). However, despite higher peak heights, we also see steeper decay slopes, indicating that increasing in-class variation is unlikely to make ICL persistent.
(b) In-weights learning.
Figure 5: In-class variation improves ICL, as previously known and seen here by the higher peak heights, but ICL is nonetheless transient across settings.
(b) In-weights learning.
Figure 6: Effect of dataset size on ICL transience. Increasing the number of classes leads to higher ICL peak height and a more gradual decay slope. Despite the number of exemplar-label pairs being the same in the purple and red curves, we note that the ICL and IWL behavior is very different when we increase the number of classes (purple) versus the in-class variation (red). Also, note that chance level performance for IWL differs slightly: 1/12,800 for the other curve (purple), and 1/1,600 for the 1,600 classes lines (red and orange).
When increasing the number of classes, we again find a similar trend to Chan et al. [10] with peak heights increasing (Figure 6a). In this case, we also find that decay slopes become more gradual. To better illustrate the differences between increasing in-class variation and increasing the number of classes, we do another experiment where the total number of exemplar-label pairs is fixed to  $12$ ,  $800 \times 20 = 1$ ,  $600 \times 160 = 256$ , 000, with one run having more classes and less in-class variation, and another run having fewer classes and more in-class variation. We find that increasing the number of classes is a more effective way to promote ICL (higher peak height, shallower decay slope) than increasing the in-class variation, perhaps due to each class appearing less often during training. This result also indicates that, beyond a certain point, in-class variation actually hurts ICL performance (lower peak height, steeper decay slope, as seen in Figure 6a). To conclude, ICL remains transient as we increase the number of classes, but training for longer and longer becomes necessary to make it go away, given the shallower decay slope.
4.3 Extending to language model token embeddings
To explore the generality of our results, we replaced the image exemplars with exemplars derived from language-model token embeddings (procedure described fully in Appendix C). On all of these
(a) In-context learning.
Figure 7: ICL is transient across a range of input datasets, derived from clusterings of LLaMa 1 token embeddings from the four different sizes. Error bars indicate standard error across two possible datasets derived from each model's embedding matrix. Further results can be found in Appendix C.
(b) In-weights learning.
token-embedding derived datasets, we find that whenever ICL emerges, it is subsequently transient (Figure 7). Furthermore, the peak height is lower than in our other experiments. In terms of timing, the peak onset as well as subsequent transience occurs on a quicker timescale (note the shorter  $x$ -axis scale in Figure 7). These results also further corroborate our findings on dataset size (Section 4.2), as the constructed token embedding-label dataset is smaller ( $3,200 \times 5 = 16,000$  exemplar-label pairs, as opposed to  $1,600 \times 20 = 32,000$  pairs in our Omniglot experiments), so we would expect lower heights and sharper decay slopes. Overall, these experiments are a promising step towards connecting our findings to natural language, and demonstrate that the transience phenomenon is not limited to image-based exemplar-label datasets.
4.4 Changing the data distribution
Until now, we have trained on data for which the marginal distribution over classes was uniform. However, a common property of naturalistic data is a Zipfian (power-law) skew [26]. Zipfian data contains a few "common" classes that appear very often, and many "rare" classes that appear less often than they would if the class distribution were uniform. In Section 4.2, we found that more classes (corresponding to each class appearing less often) led to shallower decay slopes. We explored whether data with a Zipfian distribution could play a similar role in reducing the transience of ICL.
Indeed, in Figure 8, we find this to be the case. We implemented different levels of skew across data classes according to a Zipfian distribution  $p(X = x) \propto x^{-\alpha}$  ( $x$  corresponds to the rank of an input class—e.g.,  $\alpha = 0$  corresponds to a uniform distribution). At Zipf exponent  $\alpha = 1$ , we find a severely delayed ICL peak onset, thereby delaying the subsequent transience, and a shallow decay slope following this delayed onset (similar to that seen with a large number of classes in Figure 6a). Furthermore, corroborating the findings of Chan et al. [10], we find that stopping training early, when using skewed distributions, can lead to a combination of ICL and IWL in the same network (as evidenced by the increased IWL at Zipf exponent  $\alpha = 1$  as compared to Zipf exponent  $\alpha = 0$ , Figure 8b). However, it is important to note that if the Zipfian exponent is too extreme (e.g.,  $\alpha = 2$ ), ICL does not emerge at all, and the network uses IWL to solve the task. We conclude that moderately skewed Zipfian data does not eliminate transience, but mitigates it by necessitating training for an even longer time to observe transience (similar to when there is a large number of classes).
Figure 8: ICL transience as a function of dataset skew, implemented across classes in accordance with a Zipfian distribution. (a) A moderate Zipf exponent of 1 leads to a significant delay of ICL peak onset followed by a shallow decay slope, implying models would need to be trained for even longer to get full transience. However, an extreme level of skew (Zipf exponent 2) leads to the complete loss of ICL altogether. (b) IWL, evaluated on the same skewed distribution over query classes as in the training data (thus high accuracy on Zipf exponent 2 is not unexpected, as the IWL test set is dominated by a small number of classes).
5 Regularization may eliminate ICL transience
If we consider that ICL is a more generally useful solution than IWL, then ICL transience is, in a sense, the opposite of the grokking phenomenon, where a model changes from a less to a more general solution after long periods of training [27]. Regularization (or implicit regularization) may help to drive grokking [28, 29] (and some work even shows an impact of dataset size as well [30], analogous to what we find above). Thus, would we see regularization mitigating the ICL transience phenomenon?
(b) In-weights learning.
Figure 9: L2 regularization potentially eliminates ICL transience. Increasing regularization decreases ICL transience, at least up to a point. Regularization also seems to lead instead to IWL transience. (Note that the  $y$ -axis in Figure 9b is zoomed in so that IWL transience is visible.)
We see in Figure 9 that this is indeed the case, with L2 regularization seeming to even eliminate ICL transience entirely, at least on the timescales that we tested. Without any regularization, we see transient ICL as before. However, as we add increasing levels of L2 regularization, we see that the decay slope of ICL performance goes to 0. On the other hand, we actually observe some transience of in-weights learning (purple and maroon curves, Figure 9b). These results indicate that the ICL circuit is a lower norm solution than the IWL circuit. Thus, regularization may offer a path to persistent ICL. However, with values of regularization that are too high ( $>0.0001$ ), the network is over-regularized and ICL transience returns, but, notably, IWL accuracy never rises above zero either (see Appendix D for more train loss curves and more extreme values of regularization).