SlowGuess/NeurIPS_2023_OCR huggingface dataset:
6 ICL transience may be caused by competition with IWL circuits
We performed initial investigations to understand the origins of ICL transience: why does ICL fade when the IWL solution emerges? Do the ICL and IWL circuits compete with each other for resources, for example, in the transformer's residual stream [16]? If so, we would expect wider transformers with larger embedding size (and therefore greater residual stream capacity) to be less afflicted by ICL transience. In Figure 10a, we see that this is the case: Greater model embedding sizes lead to higher peak heights and gentler decay slopes. However, one confound to consider is that larger embedding may just independently enhance ICL abilities (as they enhance IWL abilities, Figure 10b), rather than reduce competition between circuits.
To better understand how scaling the embedding dimension enhances ICL and IWL, we trained models on data that could be solved via ICL only or IWL only, rather than being solvable by both strategies. For ICL-only training, we used the same type of sequences used in our ICL evaluator, where the network must use ICL to solve the query-prediction task. For IWL-only training, we used the same type of sequences used in our IWL evaluator, where the network must use IWL to solve the query-prediction task. Figure 10d shows that larger embedding dimensions enhance the model's ability to perform IWL, but do not consistently help ICL (Figure 10c). Thus, these results indicate that the increased peak height and shallower decay slope in Figure 10a are likely due to decreased competitive interactions with IWL circuits, rather than an enhancement of ICL capabilities.
We arrived at convergent evidence by building on the regularization experiments from Section 5. We applied weight decay selectively to either the ResNet embedding layers, the MLP layers, or the self-attention layers. When we applied weight decay to only the self-attention layers, ICL was still transient. When we applied weight decay to the MLP or ResNet layers, however, ICL transience was mitigated (Appendix Figure 16). These results can be interpreted in light of prior work indicating that in-weights information is stored in MLP layers [31, 32]. By selectively penalizing this behavior, we enable ICL to persist, thus providing convergent evidence that – when no weight decay is applied – ICL fades due to competition with IWL circuits.
This work builds on several papers that seek to understand in-context learning by analyzing small networks and/or controlled data sources. Xie et al. [7] consider an implicit definition of ICL (defined in terms of the model's loss) and show that the effect can be framed as a form of Bayesian inference. Chan et al. [10] show via experimentation with controlled synthetic datasets that, e.g., the bursty and
(c) In-context learning, when training on data that (d) In-weights learning, when training on data that only supports in-context learning. supports in-weights learning.
Figure 10: (a-b) We find a trend that wider models (larger embedding sizes) have higher ICL peak heights and more gradual decay slopes. (c-d) However, when we trained on data that allowed us to specifically elicit each type of learning, we found that embedding size was only consistently helpful for enhancing IWL, indicating that mitigation of ICL transience is likely due to decreased interaction between the two circuits, as opposed to an enhancement of ICL capabilities.
Zipfian nature of text data could play an important role in the emergence of ICL. Many other works consider how ICL can be used to learn certain function classes using synthetic data, and how ICL may be implementing a form of gradient descent to accomplish this [33-35].
As with the aforementioned work, here we do not experiment directly with billion-parameter transformers trained on web-scale corpora. Nonetheless, our work should be of relevance to studies that aim to derive ICL in increasingly small and compact text-based language models. By training language models of many different sizes on different amounts of data, Hoffmann et al. [13] identified a relationship between the amount of training and model size that appeared compute-optimal with respect to the training loss. This correspondence motivated Chinchilla, a 70B-parameter language model whose downstream performance (on tasks that require ICL) was comparable to far larger models. More recently, Touvron et al. [15] took this further, "overtraining" smaller models to achieve lower training loss according to a given, inference-time, compute budget. Perez et al. [36] point out that in such work, tasks that require ICL are typically included in the models' validation set in order to optimise hyperparameters directly for ICL. Our results indicate that, without appropriate validation and early stopping, training for a long time may eventually lead to ICL gradually disappearing, perhaps even without recognition from model developers. A large number of "classes" or Zipfian data distributions (both properties of natural language [37]) can mitigate this by delaying the disappearance. Moreover, regularization may completely eliminate this risk. Indeed, Hoffmann et al. [13] find that the use of optimizers with weight decay [38] leads to models that perform better, but this benefit only manifests at later stages of training (perhaps when ICL would otherwise decay).
In-context learning is often transient. The key insight of this work is the empirical finding that, if in-context learning (ICL) emerges in a transformer network, it may not necessarily persist as the model continues to be trained. That is, while ICL is often emergent, it may be in many cases also be transient (Section 3). In fact, since the loss of ICL is gradual, its creeping disappearance may not even be known to model developers. This emphasizes the importance of evaluating appropriate validation metrics over the course of training [e.g., 39]. Furthermore, these results suggest that it may be desirable to stop training some transformer models before a compute budget is reached—to
preserve ICL—even if training losses are still decreasing. This would need to be ascertained on a case-by-case basis. Other methods for preventing the disappearance of ICL, like regularization, may also be critical for delivering flexible, in-context learners.
What are strategies for mitigating ICL transience? In Section 4, we found that data distributional properties (increasing the number of classes, adding a Zipfian skew) can strengthen ICL and postpone the transience of ICL, and so can increasing embedding size (and least in the ranges we tested). The fact that most transformers are trained on language, which naturally has both of these data-distributional properties [37], may explain why ICL transience hasn't been documented in LLMs yet. However, as the trend of overtraining continues, there's a risk of losing these impressive ICL capabilities. We found that the most robust method to make ICL persistent was L2 regularization (Section 5).
Why does ICL emerge at all, if it is not asymptotically preferred? Since GPT-3 [6], the emergence of ICL has become an expected phenomenon in large-scale transformer models. Our experiments underline how curious it is that ICL emerges in such models given that it does not seem to be asymptotically preferred. Possible explanations may concern optimization quirks (similar to the slingshot mechanism proposed as an explanation for grokking [29]) or perhaps lottery tickets contained in common initialization methods [28, 40] that incentivize ICL (when trained with relevant data-distributional properties). Future work could explore these possible explanations by considering the effect of alternate optimizers or initialization schemes on the transience of ICL.
Given that ICL does emerge, why does it fade in favor of IWL? On our training data, either ICL or IWL can reach good performance. While we do not know for certain why ICL emerges in a network, the question remains of why it fades if it is a viable strategy to solve the training task. Our results in Section 6 suggest that part of the explanation may lie in a competition between IWL and ICL circuits for resources in the transformer residual stream [16]. Under the assumption that IWL is asymptotically preferred, this competition would explain why ICL fades after emerging.
Why is IWL asymptotically preferred over ICL, when both solve the task? Induction heads [9] are a possible circuit that may be responsible for ICL, and comprise of two interacting components: finding a match in-context, then copying some token forward. The match operation is fundamentally limited by the softness of the attention operation [41-43]. On the other hand, IWL relies on learning the exemplar-label mapping in-weights, which we show is feasible in Figure 10d. We postulate—and hope to investigate in future work—the possibility that although ICL and IWL can achieve perfect accuracy on the task, this "imperfect match" to prior context asymptotically incentivizes solutions (when training with cross-entropy loss, as is standard) that do not rely as much on context and instead learn in-weights.
How do we get ICL and IWL to co-exist asymptotically in the same model? Large language models clearly show a co-existence of ICL and IWL. Chan et al. [10] found that training on Zipfian data allows the network to perform ICL on rare classes, but learn common classes in-weights. Regularization was the best strategy we found to make ICL persistent, but it leads to little-or-no IWL. Given our results, the only effective way currently to get co-existence of ICL and IWL appears to be stopping training at the right time. Future work could investigate other factors that may restore co-existent ICL and IWL, even asymptotically.
Limitations and future directions. Perhaps the most significant limitation is that our experiments were not performed on large language models, one of the most important applications of transformers. Our training resources limited our ability to perform such experiments in the current work, but we believe that this is a necessary future step, to understand the full import of the current results. As a first step toward this, we demonstrated ICL transience on datasets with language model token embeddings (Section 4.3), but our datasets are still restricted to the exemplar-label setup. It would also be interesting to explore the effect of context length: For a given train or inference budget, a system designer could trade off between longer context (possibly enabling better ICL) and a larger model (possibly enabling better IWL). We also note that our definitions of ICL and IWL are relatively specific, and many more important mechanisms are likely at play in transformers—for example, contextual information could be used in conjunction with IWL to narrow down possible outputs via copying [16], which we find some preliminary evidence of in Appendix B. Future work could investigate more mechanistic explanations of ICL transience, the tradeoff between ICL and IWL, and additional recipes for reducing ICL transience.
A.K.S. and T.M. are supported by the Gatsby Charitable Foundation. This work was supported by a Schmidt Science Polymath Award to A.S., and the Sainsbury Wellcome Centre Core Grant from Wellcome (219627/Z/19/Z) and the Gatsby Charitable Foundation (GAT3755). A.S. is a CIFAR Azrieli Global Scholar in the Learning in Machines & Brains program.
The authors thank Kira Dusterwald and DJ Strouse for insightful discussions throughout the course of the project. In addition, the authors thank Andrew Lampinen, Pierre Richemond, Roma Patel, and Murray Shanahan for feedback on the draft.
[1] Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. Advances in neural information processing systems, 29, 2016.
[2] Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta-learning with memory-augmented neural networks. In International conference on machine learning, pages 1842-1850. PMLR, 2016.
[3] Jane X Wang. Meta-learning in natural and artificial intelligence. Current Opinion in Behavioral Sciences, 38:90–95, 2021.
[4] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1126–1135. PMLR, 2017.
[5] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[6] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfbcb4967418bf8ac142f64a-Paper.pdf.
[7] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021.
[8] Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202.12837, 2022.
[9] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html.
[10] Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 18878-18891. Curran Associates, Inc.,
[43] David Chiang and Peter Cholak. Overcoming a theoretical limitation of self-attention, 2022.
[44] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. CoRR, abs/1512.03385, 2015. URL http://arxiv.org/abs/1512.03385.
[45] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. Imagenet large scale visual recognition challenge, 2015.
A Model and training procedure details
All experiments used the same model and training procedure unless stated otherwise. The transformer [5] consisted of 12 layers with an embedding dimension of 64 and 8 heads. The images were embedded by a ResNet [44] with two blocks per group and channels per group (16, 32, 32, 64). The ResNet embedder was learnt jointly and not pre-trained. The integer labels were embedded using a standard embedding layer. The input embeddings were augmented with a standard sinusoidal positional encoding [5] with timescale 30 (as our sequences are all length 17). Experiments were run for up to 5e7 training steps with batch size 32, on 16 TPU v2 or v3 cores. They were trained using Adam [25] (with default parameters of  $\beta_{1} = 0.9$ ,  $\beta_{2} = 0.999$ ) and a learning rate schedule with a linear warmup up to a maximum learning rate of 3e-4 at 4,000 steps, followed by an inverse square root decay. L2 regularization was implemented by adding the squared weights of the model (excluding batch norm parameters) to the loss term. All experiments were run with 2 seeds each. In all figures, (shaded) error bars indicate standard deviation around the mean.
Also, we note that, despite the long training times, all our experiments are still in the  $<1$  epoch regime, due to the combinatorial nature of our dataset. We estimate there are over  $1,600 \times 1,599 \times (1,598 \times 1,597/2) \times (8!/(3! \times 3!)) \times (20^9) = 1.87 \times 10^{27}$  possible training sequences, of which our model sees a maximum of  $32 \times 5 \times 10^7 = 6.4 \times 10^8$ .
Source code can be found at github.com/google-deepmind/emergent_in_context_learning and github.com/aadityasingh/iccl-transience.
B Alternate evaluation metrics
While the primary evaluations used throughout the paper are inspired by the ICL and IWL evaluators of Chan et al. [10], we did consider a few alternate evaluators. Prior work has shown that transformers are prone to simple copying biases [16] and that the importance of few-shot exemplars might be more in their format rather than correctness [8, 11]. To account for copying biases, we introduce the "IWL-copy-available" evaluator, wherein we randomly replace one of the labels with the true label. Note that, during training, the correct label always appears in context, which may incentivize the formation of such naive copying circuits. To perform a stricter test of ICL performance, we thus take inspiration from the "flipped-ICL" evaluation in Wei et al. [11].
For clarity, we illustrate some abbreviated example sequences below:
- Train: X 23 Y 24 X 23 X 23 A 0 Y 24 B 1 Y 24 X?. Note there are 3 instances of X 23, 3 instances of Y 24, and one instance each of A 0, B 1. Exemplar X is always associated to label 23 through training, Y to 24, A to 0, and B to 1.
- IWL eval (in main text): A 0 C 2 F 5 B 1 Y 24 D 3 H 7 X?. Neither the exemplar nor label appears in context.
- IWL-copy-available eval: A 0 C 2 F 5 B 23 Y 24 D 3 H 7 X?. We randomly replace one of the labels with the desired answer label. Note, this task allows the network to benefit from copying, but still provides no information in-context to learn that X should map to 23.
- ICL eval (in main text): X 0 X 0 Y 1 X 0 Y 1 Y 1 X 0 Y 1 X ?. Both X and Y's labels are replaces with 0/1. Note, we also tried using other pairs of labels for the ICL evaluator (e.g., 15/20), but found little variance across this choice.
- flipped-ICL eval: X 24 X 24 Y 23 X 24 Y 23 Y 23 X 24 Y 23 X ?. Note, if the network outputs 24 here, it demonstrates in-context learning, as it must be using the exemplar-label from context. Alternatively, if the network outputs 23, it demonstrates in-weights learning, as its ignoring the exemplar-label mapping provided in context. Accuracy=1 on this evaluator means the network is doing ICL, Accuracy=0 means the network is doing IWL.
Following the results of Section 4.3, which showed that fixed embeddings exhibit transience on shorter timescales, we conducted some experiments using the above evaluators on a dataset of fixed omniglot embeddings by preprocessing all images through a Imagenet-pretrained-and-frozen Resnet18 encoder [44, 45]. Like the main paper, we considered datasets with either 1,600 or 12,800 classes, with 20 exemplars per class. We did not run these experiments as long as those in the main paper, as the purpose was to test the new evaluators.
Figure 11: Results of alternate evaluators on fixed embedding Omniglot data. Note the shorter  $x$ -axis scale. (a) ICL evaluator, same as main text. (b) IWL evaluator, same as main text. (c) flipped-ICL evaluator. Note how signs of transience are visible sooner on this stricter measure. (d) IWL-copy-available evaluator. Similarly, IWL-copy-available starts growing sooner than the main IWL evaluator, offering earlier indications of transience
As seen in Figure 11, the alternate evaluators' asymptotic behaviors largely mirror the evaluators used by Chan et al. [10]. That said, there are some interesting differences. For example, under the flipped-ICL evaluator, models trained on 1,600 classes are never able to fully override the mapping of exemplars to labels seen in training (as evidenced by accuracy remaining below 0.5). However, at the scale of 12,800 classes, models are able to (accuracy goes above 0.5 before decaying). Intriguingly, these results may connect to Wei et al. [11], except in the sense of a more diverse dataset rather than a larger parameter count, leading a model to "do in-context learning differently".
With respect to transience, we do still observe it in both cases. For 1,600 classes, the increase in performance indicates the network may be overriding its exemplar-label mapping on some sequences, but this fledgling ability is short-lived. For 12,800 classes, while the shallow decay slope makes transience on the standard ICL evaluator not visible due to the short timescale (note, ICL is still transient at this # of classes, as shown in Figure 5a), transience is clearly visible on the flipped-ICL eval.
Furthermore, we see that there is indeed an effect of copying in our networks, as seen by the improved performance of models on the IWL-copy-available evaluator as compared to the IWL evaluator.
For a clearer comparison to prior work on synthetic datasets, we opted to use the same evaluators as Chan et al. [10] in this work, but hope the above analysis shows the importance of good evaluators in diagnosing ICL transience.
C Construction of LLaMa-based exemplar-label datasets
In this section, we provide more detail as to the procedure used to construct the fixed embedding datasets used in Section 4.3.
C.1 Constructing datasets
We save the token embedding matrix for each LLaMa v1 model. These have dimensions 32,000  $\times d_{model}$ , where  $d_{model} = 4096,5120,6656,8192$  for the 7B, 13B, 30B, and 70B, LLaMa models, respectively. To go from these raw embeddings to "classes" that we can use in our setup, we perform the following steps:
1. Subselect: Many tokens correspond to individual bytes or characters from non-English languages. We observed that the statistics of these token embeddings are often different than
(b) In-weights learning.
Figure 12: A model initialization seed for which no ICL emerges on LLaMa-derived datasets
the others, so we remove these tokens. Specifically, we use tokens 259 to 29,870, inclusive, for a total of 29,612 tokens.
2. Cluster: We cluster the token embeddings using K-means clustering and cosine distance using the FAISS library [24]. As FAISS does not offer a valid minimum number of points per cluster, we first cluster into 2400 (or 4800) clusters, then pick all clusters that have more than 10 (or 5) points. This gives us at least 1,600 (or 3,200) clusters. We randomly pick the actual 1,600 (or 3,200) clusters.
3. Select from cluster: Many of these clusters contain more than 10 (or 5) points. To maximize in-class variation, we pick the furthest 10 (or 5) points from the cluster centroid. Thus, we have our 1,600 classes of 10 examples, or 3,200 classes of 5 examples.
Note that step 2 is highly dependent on the random seed used to initialize cluster centroids. We run the above procedure for all 4 LLaMa models on 2 random seeds.
C.2 Example clusters
If we map the embeddings back to tokens, we can get a sense of what various clusters may represent. Many clusters cover similar tokens, but some clusters can be more semantic or even multilingual! We show some extracted clusters (for 3,200 classes, 5 embeddings per class, on the 70B LLaMa) below:
_erm | _room | rm | _Room | _rm  
_utter | _extremely | _partially | _entirely | Comple  
hover | _mysqli | gz | RGB | rgb  
async | _predicate | Observer | _deleg | _delegate  
ive | ous | _cable | rable | ible  
FLAG | _flag | _Flag | flag | Flag  
_specified | _provided | _supplies | _supply | _supplied  
ring | _dent | dent | _rent | RENT  
_Derby | _Dit | _pedig | _ort | _Ort  
_Britain | _England | _Madrid | _Spanish | Espagne
C.3 The difficulty of emergent ICL