# Data Distributional Properties Drive
Emergent In-Context Learning in Transformers

Stephanie C.Y. Chan
DeepMind
Adam Santoro
DeepMind
Andrew K. Lampinen
DeepMind
Jane X. Wang
DeepMind
Aaditya K. Singh
University College London
Pierre H. Richemond
DeepMind
James L. McClelland
DeepMind, Stanford University
Felix Hill
DeepMind

###### 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: what aspects of the training regime lead to this emergent behavior?
Here, we show that this behavior is driven by the distributions of the
training data itself. In-context learning emerges when the training data
exhibits particular distributional properties such as burstiness (items appear
in clusters rather than being uniformly distributed over time) and having
large numbers of rarely occurring classes. In-context learning also emerges
more strongly when item meanings or interpretations are dynamic rather than
fixed. These properties are exemplified by natural language, but are also
inherent to naturalistic data in a wide range of other domains. They also
depart significantly from the uniform, i.i.d. training distributions typically
used for standard supervised learning. In our initial experiments, we found
that in-context learning traded off against more conventional weight-based
learning, and models were unable to achieve both simultaneously. However, our
later experiments uncovered that the two modes of learning could co-exist in a
single model when it was trained on data following a skewed Zipfian
distribution – another common property of naturalistic data, including
language. In further experiments, we found that naturalistic data
distributions were only able to elicit in-context learning in transformers,
and not in recurrent models. In sum, our findings indicate how the transformer
architecture works together with particular properties of the training data to
drive the intriguing emergent in-context learning behaviour of large language
models, and how future work might encourage both in-context and in-weights
learning in domains beyond language.111Code is available at:
https://github.com/deepmind/emergent_in_context_learning

(a) Model, inputs, and outputs.

(b) Sequences for training.

(c) Sequences to evaluate in-context learning.

(d) Sequences to evaluate in-weights learning.

Figure 1: Experimental design, as described in Section 2. (1(a)) For each
experiment, a transformer model is trained on sequences of image-label pairs.
The model is trained to minimize the loss on predicting the label
corresponding to the final ‘query’ image. (1(b)) In training, image-label
mappings are fixed across sequences, in contrast to few-shot meta-training.
The training data consist of a mix of ‘bursty’ and ‘non-bursty’ sequences.
Bursty sequences, featuring multiple occurrences of the same 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. (1(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. (1(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.

##  1 Introduction

Large transformer-based language models show an intriguing ability to perform
in-context learning (Brown et al., 2020). This is the ability to generalize
rapidly from a few examples of a new concept on which they have not been
previously trained, without gradient updates to the model. In-context learning
is a special case of few-shot learning in which the output is conditioned on
examples from a ‘context’, and where there are no gradient updates. It
contrasts with ‘in-weights’ learning, which is the standard setting for
supervised learning – this is slow (requiring many examples), and depends on
gradient updates to the weights. Earlier work in the context of ‘meta-
learning’ showed how neural networks can perform few-shot learning without the
need for weight updates (Vinyals et al., 2016; Santoro et al., 2016; Wang et
al., 2016). To achieve this, the researchers explicitly 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-context learning is _emergent_. Neither the model’s transformer
architecture nor its learning objective are explicitly designed with in-
context learning in mind.

Here, we consider the question of how transformer language models are able to
acquire this impressive ability, without it being explicitly targeted by the
training setup or learning objective. The emergence of in-context learning in
language models was observed as recurrent models were supplanted by
transformers, 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 learning depends on the distributional
qualities of the training data.

This hypothesis was inspired by the observation that many natural data sources
– including natural language – 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 a distribution that is
not uniform across time, instead tending to appear in clusters (Sarkar et al.,
2005; Alvarez-Lacalle et al., 2006; Neuts, 2007; Altmann et al., 2009; Serrano
et al., 2009; Lambiotte et al., 2013). Natural data also often has the
property that the marginal distribution across entities is highly skewed,
following a Zipfian (power law) distribution with a long tail of infrequent
items (Zipf, 1949; Piantadosi, 2014; Smith et al., 2018). Finally, the
‘meaning’ of entities in natural data (such as words in natural language) is
often 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 properties may result in
training data that occupies some middle-ground between the data used in
canonical supervised learning and that used for few-shot meta-training.

In particular, standard supervised training typically consists of item classes
that recur with uniform regularity, and with item-label mappings that are
fixed throughout training – these properties allow a model to gradually learn
over time, by encoding information into its weights, e.g. via gradient
descent. By contrast, few-shot or in-context meta-training generally involves
training a model directly on specially crafted sequences of data where item
classes only recur and/or item-label mappings are only fixed within episodes –
they do not recur and are not fixed across episodes (Vinyals et al., 2016;
Santoro et al., 2016). Naturalistic data, such as language or first-person
experience, has characteristics of both of these data types. As in supervised
training, items (words) do recur, and the relationship between an entity and
its interpretation (or meaning) is fixed, to some degree at least. At the same
time, the skewed and long-tailed distribution of natural data means that some
entities recur very frequently while a large number recur much more rarely.
Importantly, however, these rare items are often bursty, making them
disproportionately likely to occur multiple times within a given context
window, somewhat like a sequence of ’meta-training’ data. We can also see the
dynamic relationship between entities and their interpretation (epitomized by
synonyms, homonyms, and polysemy, in the case of language) as weaker versions
of the completely dynamic item-label mappings that are used in few-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 and measured the effects on in-context few-shot learning. We
performed our experiments over data sequences 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 input
sequences of Omniglot images and labels, varying the natural data-inspired
distributional properties of choice. 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 when trained on data that includes both burstiness and a large
enough set of rarely occurring classes.We also 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 on the training data could bias the model more strongly towards
in-context learning. The models we tested typically exhibited a tradeoff
between rapid in-context learning vs. relying on information that was stored
through slow, gradient-based updates (‘in-weights” learning). However, we
found that models could simultaneously exhibit _both_ in-context learning and
in-weights learning when trained on 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 LSTMs and RNNs (matched on number of parameters) were
unable to exhibit in-context learning when trained on the same data
distribution. It is important to note, however, that transformer models
trained on the wrong data distributions still did fail to exhibit in-context
learning. Thus, attention is not all you need – architecture and data are both
key to the emergence of in-context learning.

(a) In-context learning on holdout classes.

(b) In-weights learning on trained classes.

Figure 2: Effects of burstiness. $P(\textrm{bursty})$ indicates the proportion
of training sequences that were bursty vs non-bursty, and models are evaluated
on the two types of evaluation sequences, over the course of training.
Burstiness in the training data increases in-context learning, and decreases
in-weights learning. Also, over the course of training, in-context learning
tends to decrease while in-weights learning increases.

(a) In-context learning on holdout classes.

(b) In-weights learning on trained classes.

Figure 3: Effects of number of classes. Increasing the number of training
classes improves in-context learning, while reducing in-weights learning.

##  2 Experimental Design

###  2.1 The training data

To investigate the factors that lead to in-context few-shot learning, we
created training and evaluation sequences using the Omniglot dataset (Lake et
al., 2019, MIT License), a standard image-label dataset for few-shot learning.
Omniglot consists of 1623 different character classes from various
international alphabets, with each class containing 20 handwritten examples.
Using the Omniglot 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 1(b)). The
first 16 elements of each sequence comprised the ‘context’, and consisted of 8
image-label pairs (where each image was always 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 was unique and fixed across training, as in typical
supervised datasets. We emphasize that this is a major departure from
conventional few-shot training, where item-label mappings are completely novel
on each 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. There are many possible ways to quantify or instantiate
burstiness (e.g., Sarkar et al., 2005; Alvarez-Lacalle et al., 2006; Neuts,
2007; Altmann et al., 2009; Serrano et al., 2009; Lambiotte et al., 2013), but
the ‘bursty’ sequences in our experiment were designed to reflect the within-
context burstiness that is observed in e.g. language. To prevent the model
from simply outputting the most common label in the sequence, a second image-
label pair also appeared 3 times in the context. For the non-bursty sequences,
the image-label pairs were drawn randomly and uniformly from the full Omniglot
set. We can continuously vary the overall degree of burstiness in a dataset by
changing the proportion of ‘bursty’ vs ‘non-bursty’ sequences.

###  2.2 The model

Each element of a sequence was first passed through an embedder (a standard
embedding layer for the integer labels, and a ResNet for the images; He et
al., 2015). These embedded tokens were passed into a causal transformer model
(Fig 1(a)) (Vaswani et al., 2017). Unless stated otherwise, we used a
transformer with 12 layers and embedding size 64. The model was trained on a
softmax cross-entropy loss 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 pairs of ‘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-shot setup. 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, the labels for these two image classes were randomly re-assigned
for each sequence. One image class was assigned to 0, and the other to 1 (Fig
1(c)). Because the labels were randomly re-assigned for each sequence, the
model must use the context in the current sequence in order to make a label
prediction for the query image. Unless stated otherwise, in-context learning
was always evaluated on holdout image classes that were never seen in
training.

Although the model is always required to perform a full multi-class
classification over all possible output labels (as in training), few-shot
accuracy is computed 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 in-
context 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 of image classes that were selected uniformly without
replacement, with the same labels that were used in training (Fig 1(d)).
Because the image classes were forced to be unique within each sequence, the
query had no support in the context. Thus, the only way for a model to
correctly predict the label was to rely on information stored in the model
weights. For this problem, where the correct query label could be any of the
labels seen in training, chance was usually $1/1600$.

##  3 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 the proportion of bursty vs non-bursty sequences in the training data
(as described in Section 2.1). These experiments replicate the finding that
transformers can acquire in-context few-shot learning even without explicit
meta-training. They further show that, as hypothesized, the model displays
better in-context learning with more burstiness in training (Fig 2(a)). We
also see that in-context learning trades off against in-weights learning –
greater burstiness simultaneously leads to lower weight-based learning (Fig
2(b)). Interestingly, the models can in some cases lose an initial bias
towards in-context learning, 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-context learning performance
depends on the number of training classes (keeping the level of burstiness
fixed at $p(\textrm{bursty})=0.9$). As we increase the number of classes from
100 to 1600 (and correspondingly decrease the frequency of each class), we see
improvement of in-context learning (Fig 3(a)). As before, we also see an
accompanying decrease in in-weights learning (Fig 3(b)). This accords with our
hypothesis about the importance of having a long tail in the distribution, or
a large vocabulary. Note that the bias against in-weights learning cannot be
explained by the number of exposures to each class – even controlling for the
number of exposures, the model trained with 1600 classes is much slower to
achieve similar levels of in-weights learning. Importantly, we need both
burstiness and a large number of classes for in-context learning to emerge. In
order to further increase the number of classes beyond the 1623 available in
the original Omniglot dataset, we rotated
($0\degree,90\degree,180\degree,270\degree$) and flipped (left-right) the
images, obtaining $8\times$ more image classes. We ensured that the holdout
set did not include transformed versions of train images. Training on these
12800 classes further improved 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 may additionally be
pushed towards in-context learning by a label-multiplicity effect, described
next.

Figure 4: Dynamic meanings improve in-context learning. Increasing the number
of labels per class (‘label multiplicity’) increases in-context learning.

Multiplicity of labels. Our third set of experiments explored the effect of
dynamic meanings, with training distributions where images did not have
completely fixed labels. Each image class was assigned to multiple possible
labels and, in the data sequences, the label shown after each image was
randomly selected among the possible labels. If a class appeared more than
once in 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 labels per class) also increases in-context
learning. Again, burstiness was fixed for these experiments at
$p(\textrm{bursty})=0.9$.

(a) In-context learning on holdout classes.

(b) In-weights learning on trained classes.

Figure 5: Effects of within-class variation. When we increase the within-class
variation (from left to right), in-context learning tends to increase (5(a))
while in-weights learning decreases (5(b)). Both effects are nonetheless
upper-bounded by the difficulty of within-class generalization, with the ‘Full
Omniglot’ problem being more difficult than the rest. For the ‘Full Omniglot’
experiments, each class contained 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.

Within-class variation. We then explored another source of dynamic variation
of meaning – the amount of variation within image classes themselves. In the
lowest-variation condition, each image class consists of only a single image,
i.e. the images for a given class were always identical. In the medium-
variation conditions, we added Gaussian pixel noise to the images (resampled
for each presentation). In the high-varation condition, we used the full
Omniglot classes (each class consists of 20 different images drawn by 20
different people). To our surprise, we found that greater within-class
variation leads to greater in-context learning (Fig 5). In other words, making
the generalization problem harder actually made in-context learning emerge
more strongly – it preferentially hampered in-weights learning more than it
hampered in-context learning.

Across all the above experiments (Figs 2-5), we also evaluated in-context
learning on training classes (rather than holdout classes), again randomly
assigning the classes to labels 0 and 1 (rather than using the ones seen in
training). Evaluations looked similar in all cases, with only slightly higher
performance (Appendix C.3.

(a) Examples of Zipfian distributions.

(b) Distribution of tokens in a natural language corpus.

(c) In-context learning on holdout classes.

(d) In-weights learning on
common classes.

(e) In-weights learning on
rare classes.

Figure 6: Effects of training on Zipfian (rather than uniform) marginal
distributions over classes. (6(a)) Examples of Zipfian distributions with
varying exponents. (6(b)) The distribution of tokens in an example English-
language corpus. In (6(c)-6(e)), bars indicate mean evaluation accuracy in the
window [400k, 500k] steps of training. (6(c)) As we increase the Zipf
exponent, i.e. increasing the skew on the class distribution, we see a
decrease in in-context learning. (6(d)) In-weights learning of the 10 most
common classes, in contrast, increases with more skew. With uniform training
(Zipf exponent = 0), the model exhibits only in-context learning and not in-
weights learning. However, if we train on skewed distributions, there is a
sweet spot where both in-context learning and in-weights learning can be
maintained at a high level in the same model (Zipf exponent = 1, for this
particular training regime). Coincidentally, a Zipf exponent of 1 corresponds
approximately to the skew in many natural languages. (6(e)) Rare items from
training are never memorized (performance is at chance for all Zipf
exponents).

###  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-weights learning – 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 to
perform rapid in-context learning on new classes that appear only in holdout.
Large language models certainly 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 if the data were non-uniform in other ways. I.e.,
each class was equally likely to appear, marginalizing across the dataset. We
postulated that we might achieve both types of learning in the same model by
instead 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 form, and are
classically described as a Zipfian (power law) distribution (Zipf, 1949):

$p(X=x)\propto\frac{1}{x^{\alpha}}$ (1)

Here, $X$ is the rank of the class (e.g. 1 for the most common class), and the
exponent $\alpha\in[0,\infty)$ determines the degree of skew. Fig 6(a) shows
some examples of Zipfian distributions with various exponents. Fig 6(b) shows
an example of token distributions in English (from the Brown corpus; Francis
and Kucera, 1979).222Plot generation adapted from
https://gist.github.com/fnielsen/7102991 This type of skewed distribution
could allow a model to learn common classes in its 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 the degree of skew. We used the same training sequences as
before, with 12800 classes and $p(\textrm{bursty})=0.9$. Our results are shown
in Figs 6(c)-6(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 is no skew, all classes are relatively rare,
and we see high levels of in-context learning but no in-weights learning.
Increasing the skew leads to the loss of in-context learning and increased in-
weights learning of common classes.333We see decreased in-weights learning for
Zipf exponent = 3, because that level of skew leads to extreme focus on a tiny
number of classes (e.g. the three most common classes form 97% of the data).
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 common classes. Intriguingly, natural languages are best
described by a Zipfian distribution with an exponent of approximately 1
(Piantadosi, 2014). Note though that the sweet spot for simultaneously
maintaining in-weights and in-context learning in transformers may differ,
depending on the training regime.

(a) Transformer.

(b) Vanilla RNN.

(c) LSTM.

Figure 7: In-context learning in transformers vs. recurrent architectures. We
compare architectures while holding fixed the number of layers, hidden layer
size, and number of parameters. Only a transformer is able to attain in-
context learning; the Vanilla RNN and LSTM never perform above chance. One run
was performed for each set of hyperparameters in a hyperparameter sweep. Each
color denotes one run, but not any particular hyperparameter values.

###  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 with either a vanilla recurrent
neural network (RNN; David E. Rumelhart et al., 1985) or a long short-term
memory network (LSTM; Hochreiter and Schmidhuber, 1997). We used the same
training sequences as before, with 1600 classes and $p(\textrm{bursty})=0.9$.
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 search for 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 the recurrent models on
in-weights learning as well (see Fig 8 in the Appendix), indicating that we
cannot explain these results by proposing that recurrent models are simply
more biased towards in-weights learning than transformers.

##  4 Discussion

In summary, we find that both data and architectures contribute significantly
to the emergence of in-context learning in transformers.

Data properties that promote in-context learning. We identify several features
of training data that can promote in-context learning – burstiness, number and
rarity of training classes, and dynamic meaning (as instantiated by multiple
labels per class or within-class variation). These data properties allow in-
context learning to emerge despite differing significantly from the data used
in standard few-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 the
remarkable 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 a significantly greater capacity for in-context learning
than recurrent models – we were completely unable to elicit in-context
learning in recurrent models, even with the training procedure, number of
parameters, and model architecture otherwise matched to the transformer
experiments. We emphasize however that the transformer architecture alone was
insufficient for eliciting in-context learning – it was 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 in their bias towards either in-context learning or in-
weights learning, and could not maintain both in the 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 will give the same answer, since the labels
are fixed. Thus, in the in-context evaluation sequences, it is ambiguous (by
design) whether the model should use the labels seen in training or in the
current context, allowing us to measure the model’s bias. We also note that
even models with an initial bias towards 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 perform slow, gradient-based in-weights learning of class
information that is presented during training, while also being able to
quickly learn (without weight updates) about new classes that appear only in
evaluation. 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 could allow models to maintain both
capabilities as well – a skewed, Zipfian distribution over classes. This
allowed the models to retain information in their weights about common
classes, while simultaneously developing 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 and also promote in-
context learning, these results may help us reach a more scientific
understanding of why in-context learning emerges in transformer-based language
models. This is an area of increasing interest (e.g. Webson and Pavlick, 2021;
Xie et al., 2021; Min et al., 2022; Razeghi et al., 2022).

We emphasize that the transformers in our experiments successfully performed
in-context evaluation on 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 be performing
genuine in-context learning, and simply draw on examples seen in training
(Razeghi et al., 2022; Min et al., 2022; Xie et al., 2021) – our experiments
show that naturalistic distributional properties 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 in-context learning in domains outside of language, an
area of ongoing research (e.g. Hill et al., 2020; Finn et al., 2017; Wang et
al., 2016). Given that reinforcement learning environments are generally
designed to be uniformly distributed (Chan et al., 2022), or that supervised
datasets are frequently rebalanced to have _more uniform_ distributions
(Chawla et al., 2002; Van Hulse et al., 2007; Katharopoulos and Fleuret,
2019), we may be missing an opportunity to endow non-language models with a
powerful capability. We may need to consider data distributions more carefully
when pre-training in non-language domains, as well. For example, recent work
has shown that pre-training on language data 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 role of non-uniformity in human cognitive development. Infants
rapidly learn statistical properties of language (Saffran and Kirkham, 2018) —
could these distributional features help infants to acquire an ability for
rapid learning, or serve as useful pretraining for later learning? And could
non-uniform distributions 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
(McClelland and O’Reilly, 1995; Kumaran et al., 2016) and its application to
language understanding in the brain (McClelland et al., 2020). According to
this theory, the neocortical part of the language system bears similarities to
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 context window in a transformer model, by
representing the associations encountered most recently (the hippocampus
generally has a time-limited window; Squire, 1992). 444While the hippocampal
system is thought to store recent context information in connection weights,
whereas transformers store such information directly in their state
representations, there is now a body of work pointing out the quantitative and
computational equivalence of weight- and state-based representations of
context state for query-based access to relevant prior information (Ramsauer
et al., 2021; Krotov and Hopfield, 2021) as implemented in transformers. In
this light, it is possible to see the human hippocampal system as a system
that provides the architectural advantage of the transformer’s context
representations for in-context learning.

Future directions. The above results suggest exciting lines of future
research. How do these data distributional properties interact with
reinforcement learning vs. supervised losses? How might results differ in
experiments that replicate other aspects of language and language modeling,
e.g. using symbolic inputs, training on next-token or masked-token prediction,
and having the meaning of words determined by their context? For models that
display both in-context and in-weights learning, it would be interesting to
understand contextual cuing of already learned information – does this
increase with more exposure? There is also a lot more to understand about the
behaviors and biases of transformers vs. 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 (Van Hulse et al., 2007; Chan et al., 2022), we show
here that non-uniform training distributions can induce the emergence of at
least one useful and interesting capability, and thus can be an opportunity as
well as a challenge.

## Acknowledgments and Disclosure of Funding

We would like to thank the following colleagues for invaluable feedback and
discussion: Kris Cao, Toni Creswell, Kevin Miller, Ivana Kajic, Andrea Banino,
Ishita Dasgupta, Kenneth Marino, Irina Higgins, Murray Shanahan, Kyriacos
Nikiforou, Richard Evans, Christos Kaplanis, David Reichert, Dave Abel, and
Drew Hudson.

We would also like to thank the following colleagues for their contributions
to the transformer implementation: Igor Babuschkin, Junyoung Chung, David
Choi, Tamara Norman, Sebastian Borgeaud, Jack Rae, David Saxton, Yujia Li,
Phil Blunsom, Maribeth Rauh, Roman Ring, Nate Kushman, Vinicius Zambaldi, Tom
Hennigan

This work was funded by DeepMind.

## References

  * Altmann et al. (2009) Eduardo G. Altmann, Janet B. Pierrehumbert, and Adilson E. Motter.  Beyond Word Frequency: Bursts, Lulls, and Scaling in the Temporal Distributions of Words.  _PLoS ONE_ , 4(11):e7678, November 2009.  ISSN 1932-6203.  doi: 10.1371/journal.pone.0007678.  URL https://dx.plos.org/10.1371/journal.pone.0007678.
  * Alvarez-Lacalle et al. (2006) E. Alvarez-Lacalle, B. Dorow, J.-P. Eckmann, and E. Moses.  Hierarchical structures induce long-range dynamical correlations in written texts.  _Proceedings of the National Academy of Sciences_ , 103(21):7956–7961, May 2006.  doi: 10.1073/pnas.0510673103.  URL https://www.pnas.org/doi/abs/10.1073/pnas.0510673103.  Publisher: Proceedings of the National Academy of Sciences.
  * Brown et al. (2020) Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared 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 M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher 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.  _arXiv:2005.14165 [cs]_ , July 2020.  URL http://arxiv.org/abs/2005.14165.  arXiv: 2005.14165.
  * Chan et al. (2022) Stephanie C. Y. Chan, Andrew K. Lampinen, Pierre H. Richemond, and Felix Hill.  Zipfian environments for Reinforcement Learning.  _arXiv:2203.08222 [cs]_ , March 2022.  URL http://arxiv.org/abs/2203.08222.  arXiv: 2203.08222.
  * Chawla et al. (2002) N. V. Chawla, K. W. Bowyer, L. O. Hall, and W. P. Kegelmeyer.  SMOTE: Synthetic Minority Over-sampling Technique.  _Journal of Artificial Intelligence Research_ , 16:321–357, June 2002.  ISSN 1076-9757.  doi: 10.1613/jair.953.  URL http://arxiv.org/abs/1106.1813.  arXiv: 1106.1813.
  * David E. Rumelhart et al. (1985) David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams.  Learning internal representations by error propagation.  Technical report, 1985.  URL https://apps.dtic.mil/dtic/tr/fulltext/u2/a164453.pdf.
  * Finn et al. (2017) Chelsea Finn, Pieter Abbeel, and Sergey Levine.  Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks.  _arXiv:1703.03400 [cs]_ , March 2017.  URL http://arxiv.org/abs/1703.03400.  arXiv: 1703.03400.
  * Francis and Kucera (1979) W. Nelson Francis and Henry Kucera.  Brown Corpus Manual, 1979.  URL http://korpus.uib.no/icame/brown/bcm.html.
  * Gale et al. (1992) William A. Gale, Kenneth W. Church, and David Yarowsky.  One Sense Per Discourse.  In _Speech and Natural Language: Proceedings of a Workshop Held at Harriman, New York, February 23-26, 1992_ , 1992.  URL https://aclanthology.org/H92-1045.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep Residual Learning for Image Recognition.  _arXiv:1512.03385 [cs]_ , December 2015.  URL http://arxiv.org/abs/1512.03385.  arXiv: 1512.03385.
  * Hill et al. (2020) Felix Hill, Olivier Tieleman, Tamara von Glehn, Nathaniel Wong, Hamza Merzic, and Stephen Clark.  Grounded Language Learning Fast and Slow.  _arXiv:2009.01719 [cs]_ , October 2020.  URL http://arxiv.org/abs/2009.01719.  arXiv: 2009.01719.
  * Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long Short-Term Memory.  _Neural Computation_ , 9(8):1735–1780, November 1997.  ISSN 0899-7667.  doi: 10.1162/neco.1997.9.8.1735.  Conference Name: Neural Computation.
  * Katharopoulos and Fleuret (2019) Angelos Katharopoulos and François Fleuret.  Not All Samples Are Created Equal: Deep Learning with Importance Sampling.  _arXiv:1803.00942 [cs]_ , October 2019.  URL http://arxiv.org/abs/1803.00942.  arXiv: 1803.00942.
  * Krotov and Hopfield (2021) Dmitry Krotov and John Hopfield.  Large Associative Memory Problem in Neurobiology and Machine Learning.  _arXiv:2008.06996 [cond-mat, q-bio, stat]_ , April 2021.  URL http://arxiv.org/abs/2008.06996.  arXiv: 2008.06996.
  * Kumaran et al. (2016) Dharshan Kumaran, Demis Hassabis, and James L. McClelland.  What Learning Systems do Intelligent Agents Need? Complementary Learning Systems Theory Updated.  _Trends in Cognitive Sciences_ , 20(7):512–534, July 2016.  ISSN 13646613.  doi: 10.1016/j.tics.2016.05.004.  URL https://linkinghub.elsevier.com/retrieve/pii/S1364661316300432.
  * Lake et al. (2019) Brenden M. Lake, Ruslan Salakhutdinov, and Joshua B. Tenenbaum.  The Omniglot Challenge: A 3-Year Progress Report.  _arXiv:1902.03477 [cs]_ , February 2019.  URL http://arxiv.org/abs/1902.03477.  arXiv: 1902.03477.
  * Lambiotte et al. (2013) Renaud Lambiotte, Lionel Tabourier, and Jean-Charles Delvenne.  Burstiness and spreading on temporal networks.  _The European Physical Journal B_ , 86(7):320, July 2013.  ISSN 1434-6028, 1434-6036.  doi: 10.1140/epjb/e2013-40456-9.  URL http://arxiv.org/abs/1305.0543.  arXiv:1305.0543 [physics, q-bio].
  * McClelland and O’Reilly (1995) James L McClelland and Randall C O’Reilly.  Why There Are Complementary Learning Systems in the Hippocampus and Neocortex:InsightsFrom the Successesand Failuresof Connectionist Models of Learning and Memory.  page 39, 1995.
  * McClelland et al. (2020) James L. McClelland, Felix Hill, Maja Rudolph, Jason Baldridge, and Hinrich Schütze.  Placing language in an integrated understanding system: Next steps toward human-level performance in neural language models.  _Proceedings of the National Academy of Sciences_ , 117(42):25966–25974, October 2020.  doi: 10.1073/pnas.1910416117.  URL https://www.pnas.org/doi/abs/10.1073/pnas.1910416117.  Publisher: Proceedings of the National Academy of Sciences.
  * Min et al. (2022) 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:2202.12837 [cs]_ , February 2022.  URL http://arxiv.org/abs/2202.12837.  arXiv: 2202.12837.
  * Neuts (2007) Marcel F. Neuts.  The burstiness of point processes.  _Stochastic Models_ , April 2007.  doi: 10.1080/15326349308807275.  URL https://www.tandfonline.com/doi/abs/10.1080/15326349308807275.  Publisher: Marcel Dekker, Inc.
  * Piantadosi (2014) Steven T. Piantadosi.  Zipf’s word frequency law in natural language: A critical review and future directions.  _Psychonomic bulletin & review_, 21(5):1112–1130, October 2014.  ISSN 1069-9384.  doi: 10.3758/s13423-014-0585-6.  URL https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4176592/.
  * Ramsauer et al. (2021) Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, Victor Greiff, David Kreil, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter.  Hopfield Networks is All You Need.  _arXiv:2008.02217 [cs, stat]_ , April 2021.  URL http://arxiv.org/abs/2008.02217.  arXiv: 2008.02217.
  * Razeghi et al. (2022) Yasaman Razeghi, Robert L. Logan IV, Matt Gardner, and Sameer Singh.  Impact of Pretraining Term Frequencies on Few-Shot Reasoning.  February 2022.  URL https://arxiv.org/abs/2202.07206v1.
  * Reid et al. (2022) Machel Reid, Yutaro Yamada, and Shixiang Shane Gu.  Can Wikipedia Help Offline Reinforcement Learning?  _arXiv:2201.12122 [cs]_ , January 2022.  URL http://arxiv.org/abs/2201.12122.  arXiv: 2201.12122.
  * Saffran and Kirkham (2018) Jenny R. Saffran and Natasha Z. Kirkham.  Infant Statistical Learning.  _Annual Review of Psychology_ , 69(1):181–203, 2018.  doi: 10.1146/annurev-psych-122216-011805.  URL https://doi.org/10.1146/annurev-psych-122216-011805.  _eprint: https://doi.org/10.1146/annurev-psych-122216-011805.
  * Santoro et al. (2016) Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap.  Meta-Learning with Memory-Augmented Neural Networks.  page 9, 2016.
  * Sarkar et al. (2005) Avik Sarkar, Paul H. Garthwaite, and Anne De Roeck.  A Bayesian mixture model for term re-occurrence and burstiness.  In _Proceedings of the Ninth Conference on Computational Natural Language Learning - CONLL ’05_ , page 48, Ann Arbor, Michigan, 2005\. Association for Computational Linguistics.  doi: 10.3115/1706543.1706552.  URL http://portal.acm.org/citation.cfm?doid=1706543.1706552.
  * Serrano et al. (2009) M. Angeles Serrano, Alessandro Flammini, and Filippo Menczer.  Modeling Statistical Properties of Written Text.  _PLOS ONE_ , 4(4):e5372, April 2009.  ISSN 1932-6203.  doi: 10.1371/journal.pone.0005372.  URL https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0005372.  Publisher: Public Library of Science.
  * Smith et al. (2018) Linda B. Smith, Swapnaa Jayaraman, Elizabeth Clerkin, and Chen Yu.  The Developing Infant Creates a Curriculum for Statistical Learning.  _Trends in Cognitive Sciences_ , 22(4):325–336, April 2018.  ISSN 1364-6613.  doi: 10.1016/j.tics.2018.02.004.  URL https://www.sciencedirect.com/science/article/pii/S1364661318300275.
  * Squire (1992) Larry R. Squire.  Memory and the Hippocampus: A Synthesis From Findings With Rats, Monkeys, and Humans.  _Psychological review_ , 1992.
  * Van Hulse et al. (2007) Jason Van Hulse, Taghi M. Khoshgoftaar, and Amri Napolitano.  Experimental perspectives on learning from imbalanced data.  In _Proceedings of the 24th international conference on Machine learning_ , ICML ’07, pages 935–942, New York, NY, USA, June 2007\. Association for Computing Machinery.  ISBN 978-1-59593-793-3.  doi: 10.1145/1273496.1273614.  URL https://doi.org/10.1145/1273496.1273614.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is All you Need.  page 11, 2017.
  * Vinyals et al. (2016) Oriol Vinyals, Charles Blundell, and Timothy Lillicrap.  Matching Networks for One Shot Learning.  page 9, 2016.
  * Wang et al. (2016) Jane X. Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z. Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick.  Learning to reinforcement learn.  _arXiv:1611.05763 [cs, stat]_ , November 2016.  URL http://arxiv.org/abs/1611.05763.  arXiv: 1611.05763.
  * Webson and Pavlick (2021) Albert Webson and Ellie Pavlick.  Do Prompt-Based Models Really Understand the Meaning of their Prompts?  _arXiv:2109.01247 [cs]_ , September 2021.  URL http://arxiv.org/abs/2109.01247.  arXiv: 2109.01247.
  * Xie et al. (2021) Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma.  An Explanation of In-context Learning as Implicit Bayesian Inference.  _arXiv:2111.02080 [cs]_ , December 2021.  URL http://arxiv.org/abs/2111.02080.  arXiv: 2111.02080.
  * Zipf (1949) George Kingsley Zipf.  Human Behavior and the Principle of Least Effort - Google Books, 1949.

## Checklist

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes]

    2. (b)

Did you describe the limitations of your work? [Yes]

    3. (c)

Did you discuss any potential negative societal impacts of your work? [N/A]

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you ran experiments…

    1. (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 architectural and training details have been specified in the text. The
code will be released with the camera-ready version.

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes]

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes]

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]

  3. 3.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes]

    2. (b)

Did you mention the license of the assets? [Yes]

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [N/A]

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [N/A]

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [N/A]

  4. 4.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]

Appendix

##  Appendix A Model and training procedure: details

All experiments used the same model and training procedure, unless stated
otherwise. The transformer consisted 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 pre-trained. The
integer labels were embedded using a standard embedding layer. The input
embeddings were augmented with a standard sinusoidal positional encoding.
Experiments were run for 500k training steps on 16 TPU v2 or v3 cores. They
were trained using Adam and a learning rate schedule with a linear warmup up
to a maximum learning rate of 3e-4 at 4000 steps, followed by an inverse
square root decay. The experiments shown in Figs 5 and 6 were run with 3 seeds
each (because of the larger number of conditions in those experiments), and
all other experiments were run with 5 runs each. In all figures, (shaded)
error bars indicate standard deviation around the mean.

##  Appendix B Possible extensions: Generating new image labels

An important constraint of the model implementation and evaluation procedure
is that we do not require the models to handle novel image labels, only novel
image classes. Thus, in-context learning is evaluated on labels that were
previously seen in training, i.e. 0 and 1 (on the Zipfian-skewed experiments,
these corresponded to two most common labels). Note that, if anything, this
causes in-context learning to be more difficult for the model, since it must
overcome existing image-label associations 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 language models, though mainly for
computational efficiency), or to generate novel labels as combinations of
already-seen tokens (akin to language models that use the SentencePiece family
of tokenization).

##  Appendix C Experiments comparing recurrent vs. transformer

###  C.1 Architectural details

Hyperparameter sweep:

  * •

Max learning rate: 15 samples log-uniform distributed over the range [1e-5,
0.1]

  * •

Num warmup steps: 15 samples log-uniform distribution over the range [1,
10000]

We performed 15 runs for each architecture (Transformer with 2 or 12 layers,
LSTM with 2 or 12 layers, Vanilla RNN with 2 or 12 layers), i.e. 90 runs
total.

Parameter counts:

  * •

Transformer with 12 layers: 831,479

  * •

LSTM with 12 layers: 627,959

  * •

Transformer with 2 layers: 331,639

  * •

LSTM with 2 layers: 297,719

###  C.2 In-weights learning

(a) Transformer.

(b) Vanilla RNN.

(c) LSTM.

Figure 8: In-weights learning in transformers vs. recurrent architectures. We
compare architectures while holding fixed the number of layers, hidden layer
size, and number of parameters. One run was performed for each set of
hyperparameters in a hyperparameter sweep. Each color denotes one run, but not
any particular hyperparameter values.

Transformers exhibited similar or slightly higher in-weights learning than the
recurrent models (Fig 8), indicating that their superior in-context learning
performance (as seen in Fig 7) cannot simply be explained by a bias towards
in-context learning and against in-weights learning.

###  C.3 In-context evaluation on trained classes

Fig 9 shows results of evaluating in-context learning on classes that were
seen in training, rather than on holdout classes (the standard evaluation
setting for few-shot learning, as described in Sec 2.3. The pattern of results
is very similar between the two settings, with just slightly higher
performance when evaluating on training classes.

(a) Burstiness.

(b) Num training classes*

(c) Within-class variation.

(d) Dynamic meanings*

Figure 9: In-context learning accuracy, evaluated on classes that were
observed in training, rather than holdout classes. Patterns of results are
very similar to those shown in the main text, with overall slightly higher
performance when evaluated on training classes.

###  C.4 Multi-class in-context evaluation

For completeness, we also report the in-context evaluation results by
computing accuracy fully multi-class across all possible outputs of the model
(Fig. 10). This is in contrast to the evaluations that were reported in the
main text (as described Sec 2.3), across just the two labels that appeared in
context; the two-choice evaluation provides a more sensitive measure of
performance, ensuring that all experimental conditions have the same levels of
chance, and also ensuring that the model cannot achieve above-chance
performance simply by randomly selecting from the labels in context. Note
that, across training and both types of evaluation (in-context and in-
weights), the model is the same – it is trained to perform multi-class
classification.

Multi-class evaluation shows the same patterns of results as the two-way
evaluation from the main text. Note that the multi-class evaluation results
showing the effects of the number of training classes (Fig 10(b)) and dynamic
meanings (Fig 10(d)) need to be interpreted with caution, because the number
of model outputs changes in the different conditions, so that task difficulty
and chance levels differ for each.

The multi-class evaluation uncovers one additional interesting result, for the
models trained on Zipfian distributions (Fig 10(e)). As in the two-choice
evaluation setting, a Zipfian distribution with an exponent of 1 is the only
one able to elicit significantly above chance accuracy on both in-context
evaluation and in-weights evaluation on common classes. However, Zipf 1 models
have relatively lower few-shot performance when evaluated in the fully multi-
class setting. Further investigation revealed that this was because those
models have overall less tendency to output labels from context (Fig 11).
Nonetheless, the Zipf 1 models do perform significantly above chance in both
settings, and when forced to choose between the two labels that are shown in
context, the model performs very well on these sequences. This indicates that,
interestingly, a model can attain both in-weights and in-context learning
abilities and process an input sequence in both ways, even if it is unsure
which of those two processes it should output the result for.

(a) Burstiness.

(b) Num training classes*

(c) Within-class variation.

(d) Dynamic meanings*

(e) Zipfian distributions.

Figure 10: In-context learning accuracy, evaluated fully multi-class across
all possible outputs of the model, rather than considering outputs on just the
two labels that appeared in context. Patterns of results are qualitatively
similar to those shown in the main text. *Figures (10(b)) and (10(d)) should
be interpreted with caution, because the total number of classes differ for
each experimental condition, and therefore chance levels. We include them for
completeness. Figure 11: Frequency of outputting any of the two labels that
appear in context, for a model trained on Zipfian distributions and evaluated
on in-context evaluation sequences.


mainReference

# Virtual Node Graph Neural Network for Full Phonon Prediction

Ryotaro Okabe  Quantum Measurement Group, Massachusetts Institute of
Technology, Cambridge, MA, USA  Department of Chemistry, Massachusetts
Institute of Technology, Cambridge, MA, USA  These authors contributed equally
Abhijatmedhi Chotrattanapituk  Quantum Measurement Group, Massachusetts
Institute of Technology, Cambridge, MA, USA  Department of Electrical
Engineering and Computer Science, Massachusetts Institute of Technology,
Cambridge, MA, USA  These authors contributed equally  Artittaya Boonkird
Quantum Measurement Group, Massachusetts Institute of Technology, Cambridge,
MA, USA  Department of Nuclear Science and Engineering, Massachusetts
Institute of Technology, Cambridge, MA, USA  Nina Andrejevic  Argonne National
Laboratory, Lemont, IL, USA  Xiang Fu  Department of Electrical Engineering
and Computer Science, Massachusetts Institute of Technology, Cambridge, MA,
USA  Tommi S. Jaakkola  Department of Electrical Engineering and Computer
Science, Massachusetts Institute of Technology, Cambridge, MA, USA  Qichen
Song  Department of Chemistry and Chemical Biology, Harvard University,
Cambridge, MA, USA  Thanh Nguyen  Quantum Measurement Group, Massachusetts
Institute of Technology, Cambridge, MA, USA  Department of Nuclear Science and
Engineering, Massachusetts Institute of Technology, Cambridge, MA, USA  Nathan
Drucker  Quantum Measurement Group, Massachusetts Institute of Technology,
Cambridge, MA, USA  Applied Physics, School of Engineering and Applied
Sciences, Harvard University, Cambridge, MA, USA  Sai Mu  Department of
Physics and Astronomy, University of South Carolina, Columbia, South Carolina,
USA  Bolin Liao  Department of Materials, University of California, Santa
Barbara, Santa Barbara, CA, USA  Yongqiang Cheng  Chemical Spectroscopy Group,
Spectroscopy Section, Neutron Scattering Division Oak Ridge National
Laboratory, Oak Ridge, TN, USA  Mingda Li  Quantum Measurement Group,
Massachusetts Institute of Technology, Cambridge, MA, USA  Department of
Nuclear Science and Engineering, Massachusetts Institute of Technology,
Cambridge, MA, USA  e-mail<EMAIL_ADDRESS>
###### Abstract

The structure-property relationship plays a central role in materials science.
Understanding the structure-property relationship in solid-state materials is
crucial for structure design with optimized properties. The past few years
witnessed remarkable progress in correlating structures with properties in
crystalline materials, such as machine learning methods and particularly graph
neural networks as a natural representation of crystal structures. However,
significant challenges remain, including predicting properties with complex
unit cells input and material-dependent, variable-length output. Here we
present the virtual node graph neural network to address the challenges. By
developing three types of virtual node approaches - the vector, matrix, and
momentum-dependent matrix virtual nodes, we achieve direct prediction of
$\Gamma$-phonon spectra and full dispersion only using atomic coordinates as
input. We validate the phonon bandstructures on various alloy systems, and
further build a $\Gamma$-phonon database containing over 146,000 materials in
the Materials Project. Our work provides an avenue for rapid and high-quality
prediction of phonon spectra and bandstructures in complex materials, and
enables materials design with superior phonon properties for energy
applications. The virtual node augmentation of graph neural networks also
sheds light on designing other functional properties with a new level of
flexibility.

## Introduction

The structure-property relationship defines one of the most fundamental
questions in materials science[oganov2006crystal, le2012quantitative]. The
ubiquitous presence of structure-property relationships profoundly influences
almost all branches of materials sciences, such as structural
materials[cheng2011atomic], energy harvesting and conversion and energy
storage materials[mishra2009metal, dresselhaus2007new,
liu2018antiferroelectrics], catalysts[zheng2021metal] and
polymers[jancar2010current], and quantum materials[kumar2020topological].
However, despite its central importance to materials design, building an
informative structure-property relationship can be nontrivial. On the one
hand, the number of stable structures grows exponentially with unit cell
size[oganov2019structure], and the structure design efforts have been largely
limited to crystalline solids with relatively small unit cells. On the other
hand, certain material properties are challenging to acquire due to
experimental or computational complexities.

In the past few years, data-driven and machine-learning methods play an
increasingly important role in materials science and significantly boost the
research on building structure-property relationships[zhu2021charting,
dunn2020benchmarking, peng2022human]. Complex structures such as porous
materials[altintas2021machine, schwalbe2021priori], nanoalloys[yao2022high,
hart2021machine], and grain boundaries[wagih2020learning] are becoming more
feasible to handle, and properties ranging from mechanical strength to quantum
ordering can be learned with increased confidence[guo2021artificial,
stanev2021artificial]. One particular powerful approach is the graph neural
networks (GNNs)[fung2021benchmarking]. By representing atoms as graph nodes
and interatomic bonds as graph edges, GNNs provide a natural representation of
molecules and materials. For crystalline solids, crystallographic symmetry
offers a further boost on the GNN performance, with a few symmetry-augmented
GNNs being proposed [xie2018crystal, Thomas2018TFN, geiger2022e3nn]. A few
fundamental challenges still exist. For one thing, many materials properties
are not naturally represented as a weighted aggregation of each atom in real
space, such as reciprocal and energy space properties. For another thing, the
output property length is usually fixed, like the heat
capacity[delaire2009phonon] as a single scalar. In contrast, many materials’
properties have unique degrees of dimensions, such as the number of electronic
and phononic bands[baroni:phonon:2001], frequency ranges with optical
responses, and the features of magnetic structures like propagation vectors.

In this work, we propose Virtual Node Graph Neural Network (VGNN) as a
generically applicable approach to augment GNN. In contrast to symmetry-
augmented GNN which focuses on reducing the input data volume, VGNN focuses on
handling the output properties with variable or even arbitrary dimensions. We
study materials’ phonon spectra and dispersion relations, given that phonons
bands are challenging to compute or measure with high computational cost and
limited experimental resources. By using the phonon spectra as examples, we
present three versions of VGNN: the vector virtual nodes (VVN), the matrix
virtual nodes (MVN), and the momentum-dependent matrix virtual nodes
($k$-MVN). All three VGNN models take atomic structures as input without prior
knowledge of interatomic forces. The VVN is the simplest VGNN that takes in a
crystal structure with $m$ atoms and outputs $3m$ branches $\Gamma$-phonon
energies. The MVN is a more involved VGNN that shows higher accuracy for
complex materials with slightly higher computational cost. Finally, the
$k$-MVN is a VGNN that can predict full phonon band structure at arbitrary $k$
points in the Brillouin zone. To achieve so, the crystal graphs contain
"virtual-dynamical matrices", which are matrix structures that resemble phonon
dynamical matrices[kong2011phonon]. Instead of performing direct ab initio
calculations on each material, all matrix elements are learned from the neural
network optimization process using training data comprised of all other
materials. Our work offers an efficient technique that can compute zone-center
phonon energies and full phonon band structures directly from atomic
structures in complex materials and enables phonon property optimization
within a larger structure design space. The prediction methods has enabled us
to acquire relevant information of materials such as group velocities, heat
capacities, density of states as by-products. Meanwhile, the virtual node
structures also shed light on future flexible GNN design, that to put
intermediate crucial quantities (e.g. dynamical matrix) as key learning
parameters without having to put target properties (e.g. phonon band
structures) as output.

Figure 1: Overview of virtual node graph neural network (VGNN). a. Atomic
structure of a crystalline material with $m$ atoms per unit cell. b. A GNN
converts the atomic structures into a crystal graph. After layers of graph
convolutions (omitted for simplicity), the final node features are aggregated
into a single fixed-sized output feature. c. A flexible of $n$ virtual atoms
are added into the crystal structure. d. After forming the crystal graph with
both real and virtual nodes, the flexibility of virtual nodes enables the
choices of output not necessarily from real-node aggregation but can have
variable length and in different spaces.

## Results

Virtual node augmentation for graph neural networks. Figure 1 gives an
overview of the VGNN method as a generic approach to augment GNN. For a
crystal with $m$ atoms per unit cell (Figure 1a), a typical GNN model converts
the crystal into a crystal graph, where each graph node represents an atom,
and each graph edge represents the interatomic bonding as shown in Figure 1b.
The node features associated with each atomic node (gray arrays in Figure 1b)
are updated by neighborhood nodes and edges connecting the nodes (gray arrows
in Figure 1b). After iterative layers of graph convolutions, $m$ final-layer
node features are obtained that represent the atomic features from each of the
$m$ atoms. The final graph output can be obtained by aggregating the final-
layer node features into one fixed-sized output.

Figures 1c,d describe the general idea of VGNN that endows a GNN with greater
flexibility for prediction. On top of the conventional, real-node GNN, virtual
atoms are added into crystal (yellow nodes in Figure 1c), which become the
virtual nodes in the corresponding GNN (yellow nodes in Figure 1d). As Figure
1d illustrates, just like the bi-directional message passing between real
atomic nodes (double-arrow gray lines), the message passing (double-arrow
yellow lines) between virtual nodes is also bi-directional. On the other hand,
to preserve the structure of the conventional GNN, the messages from real
nodes to virtual (single-arrow gray-to-yellow gradient lines) are uni-
directional. Given the flexibility of the choice of the virtual nodes, a VGNN
gains huge flexibility to predict materials-dependent outputs with arbitrary
lengths and in spaces. We will introduce three VGNN methods for phonon
prediction with increased levels of predictive power and complexity.

Vector virtual nodes for $\Gamma$-phonon prediction. As illustrated in Figure
1, VGNN makes it possible to adjust output dimension based on input
information with flexibility. We first introduce the vector virtual node (VVN)
method, which is the simplest approach to acquire $3m$ phonon branches when
inputting a crystal with $m$ atoms per unit cell. (See Methods for more
detail) Figure 2 shows the VVN approach to predict $\Gamma$-phonon spectra.
Since the virtual nodes do not pass information to real nodes, there is
additional flexibility in choosing the position of the virtual node. Without
loss of generality, we assign the position of the virtual nodes evenly spaced
along the diagonal line of the unit cell. The crystal graph is constructed
with virtual and real nodes (Figure 2a). After updating node features in each
convolution layer, the feature vectors pass a linear layer so that virtual
node features $V_{i},i\in[1,3m]$ are converted to $3m$ scalars, which
represent the predicted $\Gamma$-phonon energies. Throughout this work, the
GNN part is implemented through the Euclidean neural networks[geiger2022e3nn]
that are aware of the crystallographic symmetry. Data preparation, neural
network architectures, and optimizations are described in Supplementary
Information I-III.

The main results using the VVN for $\Gamma$-phonon prediction are shown in
Figure 2b. The three-row spectral comparison plots are randomly selected
samples from the test set within each error tertile (top-to-bottom rows are
top-to-bottom performance tertiles, respectively). The first four columns are
taken from the same database as the training set from high-quality density-
functional perturbation theory (DFPT) calculations [petretto2018high], and the
fifth column contains additional test examples with much larger unit cells
from a frozen phonon database[togo2015kyoto]. It is worthwhile mentioning that
although for very complicated materials, the VNN-predicted phonons tend to
have higher frequencies (e.g., third row, fourth column of Ba12I36Y4), the
resulted phonon density-of-states over the entire Brillouin zone can still be
largely comparable. However, the prediction loss becomes larger and
distributed broader as the input materials are more complicated (Figure 2c).
From the correlation plot of predicted and ground-truth phonon frequencies
(Figure 2d), most data points are along the diagonal line, indicating good
prediction between VNN prediction and ground-truth from DFPT calculations with
the number of atoms per unit cell $m\leq 24$ (blue dots). For complex
materials, the correlation performance could be degraded (orange dots). More
test results are shown in Supplementary Information IV.

Figure 2: The vector virtual node (VVN) method to predict $\Gamma$-point
phonons. a. Schematic of VVN model construction and prediction. For material
with $m$ atoms per unit cell, $3m$ Virtual nodes are augmented along the
diagonal vector $\vec{v}=\vec{a}+\vec{b}+\vec{c}$ of the unit cell. We
embedded the components of the crystal when building the GNN model. For
instance, atomic numbers of the $m^{\text{th}}$ real atom ($A_{R_{m}}$) and
that of the $3m^{\text{th}}$ virtual atom ($A_{V_{3m}}$) are embedded as the
attributes of each nodes. The atomic mass of $m^{\text{th}}$ real atom
($Z_{R_{m}}$) is set as the initial feature of that node. The relative
position of the node $V_{1}$ with respect to $R_{m}$ is
$\vec{r}_{V_{1}R_{m}}$, which is used to embed the edge attribute between the
two nodes. The model predicts $\Gamma$-phonon spectra by sorting the scalar
output features from virtual nodes. b. Spectral prediction samples in the test
set within each error tertile compared with ground truth (black): Test from
the same database as the training set (blue), and a different database
containing complex materials (orange). c-d. Evaluation of the test accuracy
through the distribution of loss function and correlation plot between ground-
truth and predicted average phonon frequencies, respectively. The heavy
distribution at low loss regime of the distribution plot and the agreement
along the diagonal line of the correlation plot for the test set (blue)
indicates a high-quality phonon prediction at least for relatively simple
materials with the number of atoms per unit cell $m\leq 24$. The loss becomes
higher with reduced performance for complex materials (orange).

Matrix virtual nodes for $\Gamma$-phonon with enhanced performance. In this
section, we introduce another type of virtual nodes approach, the matrix
virtual nodes (MVN). The MVN approach performs better $\Gamma$-phonon
prediction than VVN, especially for complex materials, with a slightly higher
computational cost. Moreover, the structure of MVN lays the groundwork for the
full phonon band structures to be discussed in the next section. In MVN, $m$
copies of virtual crystals are generated for material with $m$ atoms per unit
cell, and each copy contains $m$ virtual nodes that share the same crystal
structure as the real crystal (Figure 3a). This results in a total of $m^{2}$
virtual nodes $V_{ij}$, $i,j\in[1,m]$ with more involved node connectivity.
(See Methods for more detail).

With this graph construction scheme, after the neural network training, the
virtual nodes $V_{ij}$ would capture the essence of the connection between
$R_{i}$, and $R_{j}$. Hence, after the message passes in each convolutional
layer, each virtual node feature is further converted into a three-by-three
matrix. Each of $V_{ij}$ is assembled to form $(i,j)$ block of a supermatrix
$\tilde{D}$ of shape $(3m,3m)$. Given the structural similarity of this matrix
and the dynamical matrix expressed in Equation (2) with $\vec{k}=0$, we
predict $\Gamma$-point phonon energies by solving for $3m$ eigenvalues of the
matrix $\tilde{D}$. It is still worthwhile mentioning that although the matrix
shares a similar feature with the dynamical matrix, the matrix elements are
learned from neural network training and are not necessarily the matrix
elements from the real dynamical matrix. An intuitive comparison is that the
edge of GNN does not necessarily reflect true chemical bonding, but is more
like an atomic neighbor connection.

The predicted phonons using MVN are summarized in Figure 3b, which shares the
same structure with Figure 2b as error tertile plots from the high-quality
DFPT database (blue) and database for complex materials (orange). MVN shows
comparable performance with VVN for simple materials (blue curves in Figure 2c
and Figure 3c), but shows significant performance improvement for complex
materials. The prediction loss distribution of MVN shows a heavier
distribution toward a lower loss regime compared to VVN (orange curves in
Figure 2c and Figure 3c), and the average phonon frequencies in the
correlation plot align better toward ground truth (orange dots in Figure 2d
and Figure 3d). More results and correlation plots are shown in Supplementary
Information IV.

Figure 3: The matrix virtual node (MVN) method to predict $\Gamma$-point
phonons. a. Augment $m^{2}$ virtual nodes as $m$ sets of virtual crystals
(left) and the message passing scheme and post-processing of virtual node
features (right). The legends are the same as Figure 2. In contrast to VVN,
where each node $V_{j}$ is a scalar, here, each node $V_{ij}$ is a $3\times 3$
matrix. The phonon spectra in MVN are obtained by solving the eigenproblems
instead of direct output, as done in VVN. b. Selected test examples within
each error tertile. Tests from the same dataset as the training set and
additional tests containing complex materials are predicted in blue and
orange, respectively. c. Comparison of prediction loss distribution with
several examples of materials. d. The correlation plots of average phonon
frequencies with the graph $y=x$ as reference. Better performance for MVN is
achieved than VVN for complex materials (orange color), which can be seen from
both the loss distribution and the average phonon frequencies.

Momentum-dependent matrix virtual nodes for predicting full phonon band
structures. The structure of MVN inspires us to take one step further and
construct full momentum-dependent virtual dynamical matrices by taking into
account the unit cell translation, termed momentum-dependent matrix virtual
nodes ($k$-MVN). We construct virtual-dynamical matrices following Equation
(2). In contrast to the MVN, which focuses on $\Gamma$-point phonons by taking
$\vec{k}=0$, here in $k$-MVN, we include the phase factor
$e^{i\vec{k}\cdot\vec{T}}$ when defining the virtual dynamical matrices, where
$\vec{T}$ is the relative unit-cell translation of a neighboring unit cell
origin relative to the chosen reference unit cell $\vec{T_{0}}$ (Figure 4a).
If a total number of $t$ neighboring unit cells are included, each with
translation $\vec{T_{h}},h\in[0,t-1]$ (reference cell included), then a total
$t$ copies of MVN-type virtual nodes matrices will be generated, with a total
number of $tm^{2}$ virtual nodes $V^{h}_{ij},h\in[0,t-1],i,j\in[1,m]$ in
$k$-MVN. To obtain the phonon band structure, each set of virtual nodes at a
given $\vec{T_{h}}$ needs to multiply by the phase factor
$e^{i\vec{k}\cdot\vec{T_{h}}}$, and all virtual nodes at each $\vec{T_{h}}$
are summed in Equation 3. Thanks to the graph connectivity within the cutoff
radius (see Methods), only a small number of $t$ is needed as long as crystal
graph connectivity can be maintained. In practice, $t$ is materials dependent,
and $t=27$ (nearest neighbor unit cells) is sufficient for many materials and
does not need to go beyond $t=125$ (next-nearest neighbor unit cells) in all
cases. Intuitively, such a supercell approach resembles the ab initio band
structure calculations with frozen phonons. To facilitate the training,
phonons from selected high-symmetry points are included in the training data,
without the need to use full phonon energies in the entire Brillouin zone.
This significantly facilitates the training process while maintaining
accuracy. More details are discussed in Methods and Supplementary Information
V.

Figure 4b shows the prediction results of phonon band structures. Here 12
materials are selected from the same dataset for training (blue color) and the
additional dataset for complex materials (orange color). Despite the
complexity of a generic phonon band structure, the $k$-MVN model could predict
the positions and the shapes of the phonon bands, such as gaps between
different optical branches. The dispersion relations of the acoustic phonons
are also well generated around the $\Gamma$-points on the left three columns,
even though we do not enforce that acoustic $\Gamma$-phonons have to be
gapless with zero-energy known as acoustic sum rule[sham1969electronic]. This
may enable the prediction of crystal stability for future works. While there
are risks that prediction performance could be degraded for the phonon bands
of at higher frequencies, most of the predicted phonons follow the references,
including the complex materials with more than 40 atoms per unit cell.

Figure 4: The momentum-dependent matrix virtual nodes ($k$-MVN) to predict
full phonon band structures. a. (Top) Augment $m^{2}$ virtual nodes for each
translation vector $\vec{T}$, and with a total $t$ neighboring unit cells, a
total $tm^{2}$ virtual nodes are generated. (Bottom) By multiplying a phase
factor by each translated unit cell, a full virtual dynamical matrix can be
constructed. b. Selected examples in the test set within each error tertile,
for the high-quality DFPT database (blue) and additional complex materials
test (orange). $\Gamma$-point positions are labeled for each spectrum.

## Discussion

We demonstrate the prediction of phonons directly from the materials’ atomic
coordinates, using three different types of virtual nodes – the VVN, the MVN,
and the $k$-MVN – to augment the symmetry-aware Euclidean neural networks. The
comparison between the three virtual node approaches is summarized in Table 1.
VVN directly acquires the phonon spectra from the virtual nodes. The
assignment of $3m$ virtual nodes ensures that the output phonon band number is
always $3m$ for a crystal with a primitive unit cell containing $m$ atoms. In
MVN, instead of computing phonon energies directly, a virtual dynamical matrix
(VDM) is constructed first, from which the phonon energies are solved as an
eigenvalue problem. This step is crucial to gain robustness for complex
materials prediction since intermediate quantities like force constants and
dynamical matrices are considered more “fundamental” than final phonon
energies to reflect the interatomic interactions. The $k$-MVN goes one step
further, using the unit-cell translations to generate the momentum dependence
that could be used to obtain the full phonon band structure.

Today, the ab initio calculations like frozen-phonon and DFPT remain the most
accurate methods for phonon calculations. Even so, since the VGNN-based phonon
calculation skips the direct calculation of the material-by-material dynamical
matrix, it shows significantly faster computation speed while maintaining
reasonable accuracy. Additional tests on SiGe alloys, FeCoNi alloys, and other
high-energy alloys are performed, which agree well with existing literature
(Supplementary Information VI). Finally, by using MVN, we build a database
containing the $\Gamma$-phonon spectra for over 140,000 materials listed in
Materials Project (Supplementary Data and Supplementary Information VII). It
took an eight-GPU system less than five hours to obtain all results, even
though some materials contain over 400 atoms per unit cell. Such efficiency
enables the material design, searching, and optimization in a much larger
design space, including alloys, interfaces, and even amorphous solids, with
superior engineered phonon properties for thermal storage, energy conversion
and harvesting, and superconductivity. In parallel, by taking advantage of the
flexibility endowed by virtual nodes, other properties that are challenging to
predict for a conventional GNN can be predicted similarly, such as electronic
band structures and tight-binding and $k\cdot p$ effective Hamiltonian with a
variable number of bands, optical properties like flexible optical absorption
peaks as in the Lorentz oscillator model, and magnetic properties such as the
number of propagation vectors.

Table 1: Comparison of how the virtual nodes contribute to phonon prediction in terms of physics and computational costs. Here $m$, $N_{train}$, $t$, $n$ indicate the number of atoms per unit cell, the average of that in training data, the number of the unit cell counts, and an arbitrary, not large number respectively. | VVN | MVN | k-MVN
---|---|---|---
Force constants | - | - | Reflected in VDM
Dynamical matrices | - | VDM | VDM
Phonon data | Virtual nodes | Eigenvalues | Eigenvalues
Run time | $O(m^{2})$ | $O(m^{2.37})$ | $O(t\times m^{2.37})$
Storage | $O(m)$ | $O(m^{2})$ | $O(t\times m^{2})$
Generalization to larger systems | False | True | True

* abbrv ms

## Methods

### Phonon data preparation

We trained all of our models against an ab initio DFPT computational database
for phonon dispersion in harmonic model[petretto2018high]. The data set
contains material structures (the same as the primitive structure obtained
from the Material Project[jain2013commentary]), second-order derivatives of
energies with respect to atomic perturbations for regular points inside the
irreducible zone, and phonon dispersion along highly symmetric paths of 1,521
crystalline inorganic materials. These materials have 2 to 40 atoms per unit
cell, with an average of 7.38. For this work, we only used the highly
symmetric path phonon dispersion as our training data. The dispersion is
between wave vectors $\vec{k}$ in the fractional reciprocal unit and response
spectra in cm-1. All models randomly split the data into 90% training (1,365
materials), and 10% testing (156 materials) sets. Furthermore, we trained our
models with a 5-fold cross-validation scheme.

We also got phonon dispersion of complex (more number of atoms per unit cell)
materials from Atsushi Togo’s phonon database[togo2015kyoto]. We used
$seekpath$ [hinuma2017band, togo2018texttt] module to get the highly symmetric
path of each material. Then, we fed it alongside $POSCAR$, $FORCE\\_SET$, and
$phonopy.config$ files from the database to $Phonopy$[togo2015first]’s python
command to calculate the phonon dispersion along such path. To quality control
the data, we selected materials whose lowest $\Gamma$-phonon band is higher
than $-0.07$ cm-1. We also filtered the material to get only the ones with
more than 40 atoms per unit cell. Finally, we randomly selected 156 (the same
as the number of data in the testing set for ease of comparison) out of 505
filtered materials. We used them as our complex material data set.

### Computation environments

We coded the models in Python 3.9.13 and trained them on our GPU cluster with
CUDA version 10.2. To facilitate the model implementation, and training, we
used some important python modules: $Pymatgen$[ong2013python] and
$ase$[mortensen2017atomic] for handling material structure files ($.CIF$),
$PyTorch$[pytorch2018pytorch] for managing model training framework,
$e3nn$[geiger2022e3nn] for implementing our neural network models in the form
that is equivariant for Euclidean group.

### Virtual node graph neural network (VGNN)

We have developed a scheme for a graph neural network (GNN) for it to be able
to have variable output dimensions depending on the input size. For ease of
understanding, we will explain the method with our work on phonon prediction.

Considering a material with $m$ atoms per unit cell, we add $n$ additional
virtual atoms. We can adjust the number $n$ depending on the model
architecture. Using both real and virtual atoms, we convert the crystal
structures into periodic graphs with $m$ real nodes for the actual atoms and
$n$ virtual nodes for the added virtual atoms. Then, we connect nodes with
edges indicating the message-passing process. To preserve the structural
information of the materials and limit the computational cost, we apply the
following rules for connections. First, if the distance between the two real
nodes is within a specified cutoff radius $r_{\text{max}}$, the real nodes are
connected through bi-directed edges. We also set up an edge between a real
node and a virtual node according to the model description, but this edge is
directed from real to virtual nodes. Lastly, we embed the information of
radial distance vector, e.g., $\vec{r}_{ab}$ from atom $b$ to $a$, in the form
of radial basis functions and spherical harmonics on the corresponding edge as
edge attributes, which represent the distance and the direction of
$\vec{r}_{ab}$ respectively.

Since each node represents an atom in the unit cell, we embedded the atomic
numbers $A$ information as node attributes $\mathscr{A}$ by passing one-hot
representation vectors of length 118 through an embedding layer. As for the
model’s input, we embedded the atomic masses $Z$ information as input node
features $\mathscr{Z}$ by passing the product of atomic mass and one-hot
representation of atomic number through an embedding layer.

The constructed graph is then passed through the model message passing that
operates on the features with multiple convolutions and gated activation
layers[miller2020relevance]. After the final layer, which consists of only a
convolution (no gated activation), each of the $n$ virtual node features is
collected, and passed through the post-processing block, which output the $3m$
predicted phonon branches. The post-processing block is different and will be
explained in detail in the subsequent section of each model.

The model is optimized by minimizing the mean squared error (MSE) loss
function between the phonon of the training data set and the one predicted by
the model after normalizing them by the maximum phonon frequency of each
material. The full network structure is provided in the supplementary
Information.

### Vector virtual node method (VVN)

VVN is a VGNN we designed for learning to predict $\Gamma$-phonon spectra from
material structures. Since, for a material with $m$ atoms per unit cell, there
are $3m$ phonon bands, one sensible choice of adding virtual atoms is to add
$3m$ virtual atoms each outputs the prediction of one of the bands. Hence,
when there are $m$ atoms in the unit cell of crystalline material, we assign
the position $\vec{r}_{V_{i}}$ of the virtual nodes $V_{i},i\in[1,3m]$
following equation (1). We can set the atomic species of the virtual node as
anything, and we use Fe after optimization.

$\vec{r}_{V_{i}}=\frac{i-1}{3m}(\vec{a}+\vec{b}+\vec{c}).$ (1)

Here $\vec{a}$, $\vec{b}$, $\vec{c}$ indicates the unit cell vector of the
material. In other words, $3m$ virtual atoms are placed along the diagonal
line from $(0,0,0)$ to $\vec{a}+\vec{b}+\vec{c}$ with equal spacing. By
keeping the distances between the virtual nodes in the real space, it is
possible to give position dependencies to the feature updating process. In
that sense, equation (1) can consistently keep virtual nodes away from each
other and enables us to use the virtual $3m$ virtual nodes as the output nodes
of the network. To get information from the whole structure, each of the $3m$
virtual nodes is connected to all of the real nodes via directed edges from
real to virtual nodes. After each convolution layer, the virtual node features
are passed to a linear layer, converted to a scalar output, and sorted based
on their magnitudes. The outputted $3m$ scalars represent the predicted
$\Gamma$-phonon.

### Matrix virtual node method (MVN)

MVN is a VGNN we designed with the influence of the dynamic matrix
representation of a periodic harmonic system for learning to predict
$\Gamma$-phonon spectra from material structures. Given the momentum vector
$\vec{k}$, the dynamical matrix element $\tilde{D}_{ij}(\vec{k})$, which is a
three-by-three matrix representing 3D harmonic interaction between atom
$R_{i}$ and $R_{j}$, can be written as the Fourier transform of the force
constant matrix $\Phi^{\alpha\beta}_{ij}$ following equation (2). Here,
$Z_{R_{i}}$ is $R_{i}$ atom’s atomic mass, and $\vec{T}_{\alpha}$ is the
$\alpha^{\text{th}}$ unit cell position. Note that, for each $k$-vector, the
system has $3m$ degrees of freedom and frequencies where $m$ is the number of
atoms per unit cell. We can get the phonon dispersion relations
$\omega(\vec{k})$ by solving eigenvalues $\omega^{2}(\vec{k})$ of
$\tilde{D}(\vec{k})$, which is a matrix with shape $(3m,3m)$ that composed of
$m^{2}$ blocks of $\tilde{D}_{ij}(\vec{k})$ for $i,j\in[1,m]$,

$\tilde{D}_{ij}(\vec{k})=\sum_{\alpha,\beta}{\frac{\Phi^{\alpha\beta}_{ij}}{\sqrt{Z_{R_{i}}Z_{R_{j}}}}e^{i\vec{k}\cdot(\vec{T}_{\alpha}-\vec{T}_{\beta})}}.\\\
$ (2)

In the MVN method, we generate a matrix that could work like a dynamical
matrix as is written in equation (2). Here, we focus on the prediction of
$\Gamma$-phonon, i.e. $\vec{k}=\vec{0}$. So, the contributions of the same
atom pair, e.g., $R_{i}$, and $R_{j}$ from every unit cell separation
$\vec{T}_{\alpha}-\vec{T}_{\beta}$ are summed without the $\vec{k}$-dependent
exponential phase factor. Hence, the model needs to predict a matrix with
shape $(3m,3m)$ representing such summation. In order to do that, while
preserving the relation of each matrix element, we generate $m$ virtual
crystals $C_{j},j\in[1,m]$ each of which has $m$ virtual nodes
$V_{ij},i\in[1,m]$ of the same atomic species and at the same positions as the
real atoms $R_{i},i\in[1,m]$. Here, a virtual node $V_{ij}$ represents the
interaction term from a real node $R_{j}$ to another real node $R_{i}$ by
adding a directed edge from $R_{j}$ to $V_{ij}$ whenever there is an edge
connecting $R_{j}$ to $R_{i}$. After each convolution layer, the virtual node
features are passed to a linear layer and converted to complex-valued output
vectors with length 9. For each output feature, we reshape the output features
into three-by-three matrices and arrange them such that $V_{ij}$’s matrix is
the $(i,j)$ block of $\tilde{D}$ supermatrix with shape ($3m$, $3m$). Finally,
we solve $\tilde{D}$ for its $3m$ eigenvalues, which work as the
$\Gamma$-phonon prediction.

### Momentum-dependent matrix virtual node method ($k$-MVN)

$k$-MVN is a generalization of MVN model with non-zero $\vec{k}$. Unlike the
MVN case, the $k$-MVN model needs to predict matrices representing
interactions between atoms from a unit cell, e.g., $\vec{T}_{\beta}$, to the
different unit cell, e.g., $\vec{T}_{\alpha}$. Since the phase factor depends
only on the difference in unit cell positions, we can redefine $\vec{T}$ to be
such a difference and simplify equation (2) into

$\tilde{D}_{ij}(\vec{k})=\sum_{\vec{T}}{\frac{\Phi^{\vec{T}}_{ij}}{\sqrt{Z_{R_{i}}Z_{R_{j}}}}e^{i\vec{k}\cdot\vec{T}}}\coloneqq\sum_{\vec{T}}{D_{ij}^{\vec{T}}e^{i\vec{k}\cdot\vec{T}}}.$
(3)

With this simplification, for each $\vec{T}$, we generate $m$ virtual crystal
$C_{j\in[1,m]}^{\vec{T}}$ the same way as in MVN. However, in this case,
$V_{ij}^{\vec{T}}$ represents the interaction term from a real node $R_{j}$ to
another real node $R_{i}$ that is in the unit cell with unit cell position
$\vec{T}$ with respect to $R_{j}$’s. In other words, we add a directed edge
from $R_{j}$ to $V_{ij}^{\vec{T}}$ whenever there is an edge connecting
$R_{j}$ to $R_{i}$ and that edge represent
$\vec{r}_{i}-\vec{r}_{j}=\vec{r}^{\prime}_{i}+\vec{T}-\vec{r}^{\prime}_{j}$.
Here, $\vec{r}^{\prime}$ is the atomic position relative to its unit cell.
Since GNN only considers edges with interatomic distance less than
$r_{\text{max}}$, the model can generate, with this scheme, a non-zero matrix
for only a finite number of $\vec{T}$ that satisfy

$\min_{i,j\in[1,m]}|\vec{r}^{\prime}_{i}+\vec{T}-\vec{r}^{\prime}_{j}|\leq
r_{\text{max}}.$ (4)

Hence, before the virtual crystal generations, the model also iterates through
atom pairs to find all viable $\vec{T}$.

Similar to the MVN model, we convert virtual node features into three-by-three
matrices and merge them into a matrix with shape $(3m,3m)$ representing
$\tilde{D}^{\vec{T}}$ for each $\vec{T}$. Finally, we weight sum these
matrices with their phase factor to get $\tilde{D}$ and solve for its $3m$
eigenvalues as phonon spectrum at wave vector $\vec{k}$.

Acknowledgements
RO and AC contribute equally to this work. RO, AC, AB, and ML thank M Geiger,
S Fang, T Smidt, and K Persson for helpful discussions, and acknowledge the
support from the U.S. Department of Energy (DOE), Office of Science (SC),
Basic Energy Sciences (BES), Award No. DE-SC0021940, and National Science
Foundation (NSF) Designing Materials to Revolutionize and Engineer our Future
(DMREF) Program with Award No. DMR-2118448. BL acknowledges the support of NSF
DMREF with Award No. DMR-2118523. TN, ND, and ML are partially supported by
DOE BES Award No. DE-SC0020148. TN acknowledges support from Mathworks
Fellowship and Sow-Hsin Chen Fellowship. ML acknowledges the support from the
Class of 1947 Career Development Chair and discussions with S. Yip.

Competing interests
The authors declare no competing interests

Data Availability Statement
The data that support the findings of this study are openly available in
GitHub at https://github.com/RyotaroOKabe/phonon_prediction. The
$\Gamma$-phonon database generated with the MVN method is available at
https://osf.io/k5utb/