Data Distributional Properties Drive Emergent In-Context Learning in Transformers
Anonymous Author(s)
AffiliationAddress email
Abstract
Large transformer-based models are able to perform in-context few-shot learning,without being explicitly trained for it.
This observation raises the question: whataspects of the training regime lead to this emergent behavior?
Here, we show thatthis behavior is driven by the distributions of the training data itself.
In-contextlearning emerges when the training data exhibits particular distributional propertiessuch as burstiness (items appear in clusters rather than being uniformly distributedover time) and having large numbers of rarely occurring classes.
In-context learningalso emerges more strongly when item meanings or interpretations are dynamicrather than fixed.
These properties are exemplified by natural language, but arealso inherent to naturalistic data in a wide range of other domains.
They alsodepart significantly from the uniform, i.i.d.
training distributions typically used forstandard supervised learning.
In our initial experiments, we found that in-contextlearning traded off against more conventional weight-based learning, and modelswere unable to achieve both simultaneously.
However, our later experimentsuncovered that the two modes of learning could co-exist in a single model whenit was trained on data following a skewed Zipfian distribution – another commonproperty of naturalistic data, including language.
In further experiments, we foundthat naturalistic data distributions were only able to elicit in-context learning intransformers, and not in recurrent models.
In sum, our findings indicate how thetransformer architecture works together with particular properties of the trainingdata to drive the intriguing emergent in-context learning behaviour of large languagemodels, and how future work might encourage both in-context and in-weightslearning in domains beyond language.
Introduction
Large transformer-based language models show an intriguing ability to perform few-shot learning(Brown et al., 2020).
Such models are able to generalize from a few examples of a new concept onwhich they have not been previously trained.
Earlier work in the context of ‘meta-learning’ showedhow neural networks can perform few-shot learning from a few examples without the need for anyweight updates (Santoro et al., 2016;
Vinyals et al., 2016; Wang et al., 2016)
– this is also referred toas ‘in-context learning’, as the output is conditioned on a context.
To achieve this, the researchersexplicitly designed the training regime to incentivize in-context learning, a process sometimes called’meta-training’.
In the case of transformer language models, however, the capacity for in-contextlearning is emergent .
Neither the model’s transformer architecture nor its learning objective areexplicitly designed with few-shot learning in mind.
[Table] [Table]
[Table]
(cid:199)(cid:214)(cid:143)(cid:200)(cid:232) (cid:251)(cid:250)(cid:253) (cid:134)(cid:186)(cid:181)(cid:210)(cid:143)(cid:231)(cid:210) (cid:199)(cid:214)(cid:143)(cid:200)(cid:232) (cid:134)(cid:186)(cid:181)(cid:210)(cid:143)(cid:231)(cid:210) (cid:272) (cid:248)(cid:249)(cid:248) (cid:248)(cid:248)(cid:250)(cid:253) (cid:256)(cid:247)(cid:254) (cid:251)(cid:252) (cid:248)(cid:247)(cid:247)(cid:255) (cid:252)(cid:254)(cid:256) (cid:255)(cid:249)(cid:248) (cid:247) (cid:248) (cid:247) (cid:247) (cid:248) (cid:248) (cid:248) are fixed across sequences in and sequences occurrences of the classes, can be solved by learning labels across sequences (in-weights learning) or referring back to the context (in-context learning)
Non-bursty sequences were composed of i i d images
(c) To evaluate few-shot in-context learning, the model is presented with a standard few-shot sequence.
The holdout image classes were never encountered in training, and are randomly assigned to labels {0,1}.
Thus the model must use the context to predict the query label.
(d) To evaluate in-weights learning, the model is presented with sequences where the labels are the same as in training.
However, the query class does not appear in the context.
Thus, the model must used information stored in weights to predict the query label.
In the example sequences, we add colors and use only Latin characters for visualization purposes.
Here, we consider the question of how transformer language models are able to acquire this impressiveability, without it being explicitly targeted by the training setup or learning objective.
The emergenceof in-context learning in language models was observed as recurrent models were supplanted bytransformers, e.g. in GPT3.
Was the novel architecture the critical factor behind this emergence?
In this work we explore this possibility, as well as a second: that a capacity for in-context learningdepends on the distributional qualities of the training data .
This hypothesis was inspired by the observation that many natural data sources – including naturallanguage – differ from typical supervised datasets due to a few notable features.
For example,natural data is temporally ‘bursty”.
That is, a given entity (word, person, object, etc) may have adistribution that is not uniform across time, instead tending to appear in clusters (Altmann et al.
,2009; Alvarez-Lacalle et al., 2006; Sarkar et al., 2005; Serrano et al., 2009).
Natural data also oftenhas the property that the marginal distribution across entities is highly skewed, following a Zipfian(power law) distribution with a long tail of infrequent items (Piantadosi, 2014; Smith et al., 2018;
Zipf, 1949).
Finally, the ‘meaning’ of entities in natural data (such as words in natural language) isoften dynamic rather than fixed.
That is, a single entity can have multiple possible interpretations(polysemy and homonymy, in language) and multiple entities can map to the same interpretation(synonymy, in language), usually in a context-dependent way.
The combination of these propertiesmay result in training data that occupies some middle-ground between the data used in canonicalsupervised learning and that used for few-shot meta-training.
In particular, standard supervised training typically consists of item classes that recur with uniformregularity, and with item-label mappings that are fixed throughout training – these properties allowa model to gradually learn over time, by encoding information into its weights, e.g. via gradientdescent.
By contrast, few-shot meta-training generally involves training a model directly on speciallycrafted sequences of data where item classes only recur and/or item-label mappings are only fixedwithin episodes – they do not recur and are not fixed across episodes (Santoro et al., 2016;
Vinyals
[Figure]
holdout classes
[Figure]
et al., 2016).
Naturalistic data, such as language or first-person experience, has characteristics ofboth of these data types.
As in supervised training, items (words) do recur, and the relationshipbetween an entity and its interpretation (or meaning) is fixed, to some degree at least.
At the sametime, the skewed and long-tailed distribution of natural data means that some entities recur veryfrequently while a large number recur much more rarely.
Importantly, however, these rare items areoften bursty, making them disproportionately likely to occur multiple times within a given contextwindow, somewhat like a sequence of ’meta-training’ data.
We can also see the dynamic relationshipbetween entities and their interpretation (epitomized by synonyms, homonyms, and polysemy, in thecase of language) as weaker versions of the completely dynamic item-label mappings that are used infew-shot meta-training, where the mappings are randomly permuted on every episode.
In this paper, we experimentally manipulated the distributional properties of the training data andmeasured the effects on in-context few-shot learning.
We performed our experiments over datasequences sampled from a standard image-based few-shot dataset (the Omniglot dataset; Lake et al.,2019).
At training, we fed each model (such as a transformer or recurrent network) with inputsequences of Omniglot images and labels, varying the natural data-inspired distributional properties ofchoice.
At evaluation, we assessed whether these properties gave rise to in-context learning abilities.
Our results showed that, indeed, in-context learning emerges in a transformer model only whentrained on data that includes both burstiness and a large enough set of rarely occurring classes.
Wealso tested two instantiations of the kinds of dynamic item interpretation observed in natural data –having many labels per item as well as within-class variation.
We found that both interventions onthe training data could bias the model more strongly towards in-context learning.
The models wetested typically exhibited a tradeoff between rapid in-context learning vs. relying on information thatwas stored through slow, gradient-based updates (‘in-weights” learning).
However, we found thatmodels could simultaneously exhibit both in-context learning and in-weights learning when trainedon a skewed marginal distribution over classes (akin to the Zipfian distribution of natural data).
At the same time, architecture is also important.
Unlike transformers, recurrent models like LSTMsand RNNs (matched on number of parameters) were unable to exhibit in-context learning whentrained on the same data distribution.
It is important to note, however, that transformer models trainedon the wrong data distributions still did fail to exhibit in-context learning.
Thus, attention is not allyou need – architecture and data are both key to the emergence of in-context learning.
Experimental Design
2.1 The training data
To investigate the factors that lead to in-context few-shot learning, we created training and evaluationsequences using the Omniglot dataset (Lake et al., 2019, MIT License), a standard image-labeldataset for few-shot learning.
Omniglot consists of 1623 different character classes from variousinternational alphabets, with each class containing 20 handwritten examples.
Using the Omniglot
[Figure]
context learning
[Figure]
holdout classes
[Figure]
dataset allowed us to apply evaluation procedures that are standard in the study of few-shot learning.
The few-shot challenge is to classify an example of a character class that was never seen in training,based only on a few examples of that class and some alternate classes.
The training data consisted of sequences of images and labels (Fig 1b).
The first 16 elements ofeach sequence comprised the ‘context’, and consisted of 8 image-label pairs (where each image wasalways followed immediately by its corresponding label).
The final element was the ‘query’ image,and the aim of the model was to predict the correct label for the query.
Images were allowed to recur throughout training, and the integer label for each image class wasunique and fixed across training, as in typical supervised datasets.
We emphasize that this is a majordeparture from conventional few-shot training, where item-label mappings are completely novel oneach episode, or the items themselves are novel on each episode.
In our standard experiments, we trained the model on a mixture of ‘bursty’ and ‘non-bursty’ sequences.
In the bursty sequences, the query class appeared 3 times in the context.
To prevent the model fromsimply outputting the most common label in the sequence, a second image-label pair also appeared 3times in the context.
For the non-bursty sequences, the image-label pairs were drawn randomly anduniformly from the full Omniglot set.
2.2 The model
[Figure] Each element of a sequence was first passed through an embedder (a standard embedding layer forthe integer labels, and a ResNet for the images;
He et al., 2015).
These embedded tokens were passedinto a causal transformer model (Fig 1a) (Vaswani et al., 2017).
Unless stated otherwise, we used atransformer with 12 layers and embedding size 64.
The model was trained on a softmax cross-entropyloss on the prediction for the final (query) image.
2.3 The evaluation data
We evaluated trained models on two types of sequences, to measure (1) in-context learning and (2)in-weights learning.
As in the training sequences, the evaluation sequences also consisted of 8 pairsof ‘context’ image and label tokens, followed by a single ‘query’ image token.
To measure a trained model’s ability for in-context few-shot learning, we used a standard few-shotsetup.
The context consisted of a random ordering of 2 different image classes with 4 examples each,and the query was randomly selected from one of the two image classes (a ‘4-shot 2-way’ problem,in few-shot nomenclature).
Unlike in training, where the labels were fixed across all sequences, thelabels for these two image classes were randomly re-assigned for each sequence.
One image classwas assigned to 0, and the other to 1 (Fig 1c).
Because the labels were randomly re-assigned for eachsequence, the model must use the context in the current sequence in order to make a label predictionfor the query image (a 2-way classification problem).
Unless stated otherwise, in-context learningwas always evaluated on holdout image classes that were never seen in training.
Although the model’s output layer accounts for all possible labels in the dataset, few-shot accuracy iscomputed by considering the model outputs only for the two labels seen in the few-shot sequence(0 and 1), with chance at 1 / 2 .
This ensures that performance above chance cannot be due to e.g. randomly selecting one of the labels from the context.
Note also that the model was evaluated for incontext learning on novel image classes, but not novel labels (see the appendix for further discussion).
To measure in-weights learning of trained classes in a model, evaluation sequences consisted ofimage classes that were selected uniformly without replacement, with the same labels that were usedin training (Fig 1d).
Because the image classes were forced to be unique within each sequence, thequery had no support in the context.
Thus, the only way for a model to correctly predict the label wasto rely on information stored in the model weights.
For this problem, where the correct query labelcould be any of the labels seen in training, chance was usually 1 / 1600 .
Results
3.1 What kinds of training data promote in-context learning?
Burstiness.
In our first experiments, we vary levels of burstiness in the training data by varying theproportion of bursty vs non-bursty sequences in the training data (as described in Section 2.1).
Theseexperiments replicate the finding that transformers can acquire in-context few-shot learning evenwithout explicit meta-training.
They further show that, as hypothesized, the model displays betterin-context learning with more burstiness in training (Fig 2a).
We also see that in-context learningtrades off against in-weights learning – greater burstiness simultaneously leads to lower weight-basedlearning (Fig 2b).
Interestingly, the models can in some cases lose an initial bias towards in-contextlearning, moving towards in-weights learning over the course of training.
A large number of rarely occurring classes.
Our second set of experiments show that in-contextlearning performance depends on the number of training classes (keeping the level of burstiness fixedat p ( bursty )
= 0 . 9 ).
As we increase the number of classes from 100 to 1600 (and correspondinglydecrease the frequency of each class), we see improvement of in-context learning (Fig 3a).
Asbefore, we also see an accompanying decrease in in-weights learning (Fig 3b).
This accords with ourhypothesis about the importance of having a long tail in the distribution, or a large vocabulary.
Notethat the bias against in-weights learning cannot be explained by the number of exposures to eachclass – even controlling for the number of exposures, the model trained with 1600 classes is muchslower to achieve similar levels of in-weights learning.
Importantly, we need both burstiness and alarge number of classes for in-context learning to emerge.
In order to further increase the number ofclasses beyond the 1623 available in the original Omniglot dataset, we rotated ( 0 ◦ , 90 ◦ , 180 ◦ , 270 ◦ )and flipped (left-right) the images, obtaining 8 × more image classes.
We ensured that the holdoutset did not include transformed versions of train images.
Training on these 12800 classes furtherimproved in-context learning (and reduced in-weights learning)
(Fig 3).
However, some images in
Omniglot have rotational or mirror symmetries, so that the models trained on 12800 classes mayadditionally be pushed towards in-context learning by a label-multiplicity effect, described next.
[Figure]
Multiplicity of labels.
Our third set of experiments explored the effect of dynamicmeanings, with training distributions whereimages did not have completely fixed labels.
Each image class was assigned tomultiple possible labels and, in the data sequences, the label shown after each imagewas randomly selected among the possiblelabels.
If a class appeared more than oncein the same sequence, the label was consistent for all presentations within that sequence (this is commonly the case in natural data such as language, too; Gale et al.,1992).
In Fig 4, we see that increasing the‘label multiplicity’ (the number of labelsper class) also increases in-context learning.
Again, burstiness was fixed for these experiments atp ( bursty ) = 0 .
Within-class variation.
We then explored another source of dynamic variation of meaning – theamount of variation within image classes themselves. In the lowest-variation condition, each image
[Figure]
[Figure]
[Figure]
the full set of 20 Omniglot exemplars per class.
For the remaining experiments, each consisted of only a single Omniglot exemplar image, with varying levels of Gaussian pixel noise.
class consists of only a single image, i.e. the images for a given class were always identical.
Inthe medium-variation conditions, we added Gaussian pixel noise to the images (resampled for eachpresentation).
In the high-varation condition, we used the full Omniglot classes (each class consistsof 20 different images drawn by 20 different people).
To our surprise, we found that greater withinclass variation leads to greater in-context learning (Fig 5).
In other words, making the generalizationproblem harder actually made in-context learning emerge more strongly – it preferentially hamperedin-weights learning more than it hampered in-context learning.
Across all the above experiments, we also evaluated in-context learning on training classes (ratherthan holdout classes).
Evaluations looked similar in all cases, with only slightly higher performance.
3.2 What kinds of training data enable in-context learning and in-weights learning to co-exist
in the same model?
In the previous section, we saw a consistent tradeoff between in-context learning and in-weightslearning – no models could maintain both.
However, it is useful for a model to have both capabilities– to remember information about classes that will re-appear in evaluation, while also being able toperform rapid in-context learning on new classes that appear only in holdout.
Large language modelscertainly do have both of these capabilities.
How might we achieve this?
For all prior experiments, the training data were marginally distributed uniformly over classes, even ifthe data were non-uniform in other ways.
I.e., each class was equally likely to appear, marginalizingacross the dataset.
We postulated that we might achieve both types of learning in the same model byinstead training on marginally- skewed distributions.
In this case, some classes appear very commonly,while most classes appear very rarely.
Many natural phenomena such as word distributions take this [Figure] [Figure] high
level in the
same model
Zipf
exponent = 1, for this
particular training regime).
Coincidentally
, a
Zipf exponent
of
corresponds
approximately
to the
skew
in many
natural languages.
(e) Rare items from training are never memorized (performance is at chance for all Zipf exponents).
form, and are classically described as a Zipfian (power law) distribution (Zipf, 1949):
[Equation]
Here, X is the rank of the class (e.g. 1 for the most common class), and the exponent α ∈ [0 , ∞ )determines the degree of skew.
Fig 6a shows some examples of Zipfian distributions with variousexponents.
Fig 6b shows an example of token distributions in English (from the Brown corpus; Francisand Kucera, 1979).
This type of skewed distribution could allow a model to learn common classes inits weights, while the long tail of rare classes simultaneously induces an ability for in-context learning.
To test this hypothesis, we trained on Zipfian distributions, varying the Zipf exponent and hence thedegree of skew.
We used the same training sequences as before, with 12800 classes and p ( bursty ) =0 .
Our results are shown in Figs 6c-e. We evaluate in-weights learning separately on common classes(the 10 classes seen most often in training) and on rare classes (the remaining classes).
When there isno skew, all classes are relatively rare, and we see high levels of in-context learning but no in-weightslearning.
Increasing the skew leads to the loss of in-context learning and increased in-weights learningof common classes.
In between the two extremes, we observe a sweet spot at Zipf exponent =
1,where the model maintains high levels of both in-context learning and in-weights learning of commonclasses.
Intriguingly, natural languages are best described by a Zipfian distribution with an exponent ofapproximately
(Piantadosi, 2014).
Note though that the sweet spot for simultaneously maintainingin-weights and in-context learning in transformers may differ, depending on the training regime.
[Figure]
3.3 But architecture does matter too.
To investigate whether these results are specific to transformer models, we performed similar experiments using recurrent sequence models.
For these models, we simply replaced the transformer witheither a vanilla recurrent neural network (RNN;
David E. Rumelhart et al., 1985) or a long short-termmemory network (LSTM; Hochreiter and Schmidhuber, 1997).
We used the same training sequencesas before, with 1600 classes and p ( bursty ) = 0 .
We also used the same image and label encoders,and cross-entropy classification loss.
The recurrent models were matched to the transformer for depth,number of parameters, and hidden layer size.
We performed a comprehensive hyperparameter searchfor all models (see Appendix for details).
In these experiments, we see that the recurrent models are never able to achieve in-context learning,despite the parity in training setup (Fig 7).
Interestingly, the transformer actually outperforms therecurrent models on in-weights learning as well (see Fig 8 in the Appendix), indicating that we cannotexplain these results by proposing that recurrent models are simply more biased towards in-weightslearning than transformers.
Discussion
In summary, we find that both data and architectures contribute significantly to the emergence ofin-context learning in transformers.
Data properties that promote in-context learning.
We identify several features of training datathat can promote in-context learning – burstiness, number and rarity of training classes, and dynamicmeaning (as instantiated by multiple labels per class or within-class variation).
These data propertiesallow in-context learning to emerge despite differing significantly from the data used in standardfew-shot meta-training, in that we allow items and item-label mappings to recur throughout training.
These properties are also central features of natural data including language, and thus may explain theremarkable emergence of in-context learning in large language models without explicit meta-training.
Effects of architecture.
We find that architecture does matter as well.
Transformers show asignificantly greater capacity for in-context learning than recurrent models – we were completelyunable to elicit in-context learning in recurrent models, even with the training procedure, number ofparameters, and model architecture otherwise matched to the transformer experiments.
We emphasizehowever that the transformer architecture alone was insufficient for eliciting in-context learning – itwas necessary for the training data to exhibit at least burstiness and large numbers of classes, too.
In-context vs. in-weights learning.
In most cases, we found that transformers exhibited a tradeoff intheir bias towards either in-context learning or in-weights learning, and could not maintain both inthe same model.
We characterize this behavior as a ‘bias’, because neither type of learning is ‘correct’per se.
For our training data, an in-context learning strategy and an in-weights learning strategy willgive the same answer, since the labels are fixed.
Thus, in the in-context evaluation sequences, it isambiguous (by design) whether the model should use the labels seen in training or in the currentcontext, allowing us to measure the model’s bias.
We also note that even models with an initial biastowards in-context learning can often move towards in-weights learning with enough repetition.
However, it is often important and useful for a model to exhibit both capabilities – to performslow, gradient-based in-weights learning of class information that is presented during training, whilealso being able to quickly learn (without weight updates) about new classes that appear only inevaluation.
Indeed, large language models exhibit both of these capabilities (Brown et al., 2020).
In our experiments, we discovered that an additional language-like distributional property couldallow models to maintain both capabilities as well – a skewed, Zipfian distribution over classes.
Thisallowed the models to retain information in their weights about common classes, while simultaneouslydeveloping in-context learning abilities that were presumably induced by the long tail of rare classes.
Implications for understanding language models.
Our findings have a few noteworthy implications.
First, by pointing to specific distributional properties of training data that both exist in language andalso promote in-context learning, these results may help us reach a more scientific understanding ofwhy in-context learning emerges in transformer-based language models.
This is an area of increasinginterest (e.g. Min et al., 2022; Razeghi et al., 2022; Webson and Pavlick, 2021;
Xie et al., 2021).
We emphasize that the transformers in our experiments successfully performed in-context evaluationon holdout classes, and only performed slightly better with in-context evaluation on trained classes.
These results are counter to an emerging narrative that large language models may not actually beperforming genuine in-context learning, and simply draw on examples seen in training (Min et al.
,2022; Razeghi et al., 2022;
Xie et al., 2021) – our experiments show that naturalistic distributionalproperties can give rise to a capacity for in-context learning on classes that were never seen in training.
Broader implications.
This understanding may also help us design and collect datasets to achieve incontext learning in domains outside of language, an area of ongoing research (e.g. Finn et al., 2017;
Hill et al., 2020; Wang et al., 2016).
Given that reinforcement learning environments are generallydesigned to be uniformly distributed (Chan et al., 2022), or that supervised datasets are frequentlyrebalanced to have more uniform distributions (Chawla et al., 2002; Katharopoulos and Fleuret, 2019;
Van Hulse et al., 2007), we may be missing an opportunity to endow non-language models with apowerful capability.
We may need to consider data distributions more carefully when pre-training innon-language domains, as well.
For example, recent work has shown that pre-training on languagedata was useful for offline reinforcement learning, but pre-training on vision data was not (Reid et al.,2022) – could this difference be due to the non-uniform, structured distribution of the language data? Cognition and neuroscience.
Our experiments could also potentially inspire research on the roleof non-uniformity in human cognitive development.
Infants rapidly learn statistical properties oflanguage (Saffran and Kirkham, 2018) — could these distributional features help infants to acquirean ability for rapid learning, or serve as useful pretraining for later learning? And could non-uniformdistributions in other domains (e.g., vision) also contribute to this development (cf. Smith et al., 2018)? Our results may also relate to complementary learning systems theory (Kumaran et al., 2016; McClelland and O’Reilly, 1995) and its application to language understanding in the brain (McClellandet al., 2020).
According to this theory, the neocortical part of the language system bears similaritiesto the weights of neural networks, in that both systems learn gradually through the accumulated influence of large amounts of experience.
The hippocampal system plays a role similar to the contextwindow in a transformer model, by representing the associations encountered most recently (the hippocampus generally has a time-limited window; Squire, 1992).
In this light, it is possible to see thehuman hippocampal system as a system that provides the architectural advantage of the transformer’scontext representations for in-context learning.
Future directions.
The above results suggest exciting lines of future research.
How do these datadistributional properties interact with reinforcement learning vs. supervised losses?
How mightresults differ in experiments that replicate other aspects of language and language modeling, e.g. usingsymbolic inputs, training on next-token or masked-token prediction, and having the meaning of wordsdetermined by their context?
For models that display both in-context and in-weights learning, it wouldbe interesting to understand contextual cuing of already learned information – does this increase withmore exposure?
There is also a lot more to understand about the behaviors and biases of transformersvs.
recurrent architectures – why do transformers seem to be more capable of in-context learning?
Non-uniformity.
Finally, we hope to emphasize the dual nature of non-uniformity in training data.
While it can impair both supervised and reinforcement learning (Chan et al., 2022; Van Hulse et al.
,2007), we show here that non-uniform training distributions can induce the emergence of at least oneuseful and interesting capability, and thus can be an opportunity as well as a challenge.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] All architecturaland training details have been specified in the text.
The code will be released with thecamera-ready version.
(b) Did you specify all the training details (e g., data splits, hyperparameters,
how they
were chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running experi- ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
AppendixA Model and training procedure: details
All experiments used the same model and training procedure, unless stated otherwise.
The transformerconsisted of 12 layers, with embedding dimension 64 and 8 heads.
The images were embedded by a
ResNet with two blocks per group and channels per group (16, 32, 32, 64), and which was not pretrained.
The integer labels were embedded using a standard embedding layer.
The input embeddingswere augmented with a standard sinusoidal positional encoding.
Experiments were run for 500ktraining steps on 16 TPU v2 or v3 cores.
They were trained using Adam and a learning rate schedulewith a linear warmup up to a maximum learning rate of 3e-4 at 4000 steps, followed by an inversesquare root decay.
The experiments shown in Figs 5 and 6 were run with 3 seeds each (because of thelarger number of conditions in those experiments), and all other experiments were run with 5 runseach.
In all figures, (shaded) error bars indicate standard deviation around the mean.
B Possible extensions: Generating new image labels
An important constraint of the model implementation and evaluation procedure is that we do notrequire the models to handle novel image labels, only novel image classes.
Thus, in-context learningis evaluated on labels that were previously seen in training, i.e. 0 and 1 (on the Zipfian-skewedexperiments, these corresponded to two most common labels).
Note that, if anything, this causesin-context learning to be more difficult for the model, since it must overcome existing image-labelassociations that were learned in training.
However, as future extensions, it would be possible to extend the model to handle novel labels as well.
For example, we might tie the input and output embedding layers (sometimes done in large languagemodels, though mainly for computational efficiency), or to generate novel labels as combinations ofalready-seen tokens (akin to language models that use the SentencePiece family of tokenization).
C Experiments comparing recurrent vs. transformerC.1 Architectural details
Hyperparameter sweep (15 runs for each architecture)
:• Num layers: 2 or 12• Max learning rate: Log-uniform distribution over [1e-5, 0.1]
[Table] • Num warmup steps: Log-uniform distribution over [1, 10000]
[Table]
Parameter counts:
[Table] • Transformer with 12 layers: 831, [Table] • LSTM with 12 layers: 627,
[Table]
• Transformer with 2 layers: 331,
[Table]
• LSTM with 2 layers: 297,
[Table]
C.2 In-weights learning
[Table]
Transformers exhibited similar or slightly higher in-weights learning than the recurrent models (Fig8), indicating that their superior in-context learning performance (as seen in Fig 7) cannot simply beexplained by a bias towards in-context learning and against in-weights learning.
[Figure]