A N
E XPLANATION
OF I N - CONTEXT L EARNING AS
I MPLICIT B AYESIAN I NFERENCE
Anonymous authors
Paper under double-blind review
A BSTRACT
Large pretrained language models such as GPT-3 have the surprising ability to do in-context learning, where the model learns to do a downstream task simply by conditioning on a prompt consisting of input-output examples.
Without being explicitly pretrained to do so, the language model learns from these examples during its forward pass without parameter updates on “out-of-distribution” prompts.
Thus, it is unclear what mechanism enables in-context learning. In this paper, we study the role of the pretraining distribution on the emergence of in-context learning under a mathematical setting where the pretraining texts have long-range coherence.
Here, language model pretraining requires inferring a latent document-level concept from the conditioning text to generate coherent next tokens.
At test time, this mechanism enables in-context learning by inferring the shared latent concept between prompt examples and applying it to make a prediction on the test example.
Concretely, we prove that in-context learning occurs implicitly via Bayesian inference of the latent concept when the pretraining distribution is a mixture of HMMs.
This can occur despite the distribution mismatch between prompts and pretraining data.
In contrast to messy large-scale pretraining datasets for in-context learning in natural language, wegenerateafamilyofsmall-scalesyntheticdatasets(GINC)whereTransformerand LSTM language models both exhibit in-context learning.
Beyond the theory which focuses on the effect of the pretraining distribution, we empirically ﬁnd that scaling model size improves in-context accuracy even when the pretraining loss is the same.
I NTRODUCTION
Large language models (LMs) such as GPT-3 (Brown et al., 2020; Lieber et al., 2021; Wang & Komatsuzaki, 2021; Radford et al., 2019) are pretrained on massive text corpora to simply predict the next word given previous words.
They demonstrate the surprising ability to do in-context learning , where a LM “learns” to do a downstream task simply by conditioning on a prompt containing input-output pairs.
In-context learning in GPT-3 has achieved SOTA results in the LAMBADA (Paperno et al., 2016) completion task and the TriviaQA question answering task (Joshi et al., 2017) (18% and 3% over previous SOTA (Brown et al., 2020)).
For example, consider the task of predicting nationalities from names.
A prompt (Figure 1) is constructed by concatenating independent “training” examples (e.g., “Albert Einstein was German”) followed by a “test example” (“Marie Curie was”).
The LM conditions on this prompt, presented as a contiguous sequence of examples, to make a prediction.
In this example, GPT-3 places the largest probability on the correct output p ( “Polish” | “Albert Einstein was German \ n Mahatma Gandhi was Indian \ n Marie Curie was” ) seemingly by inferring the task (names to nationalities) from the examples.
Intuigingly, GPT-3 was not explicitly pretrained to learn from examples and does not optimize any parameters on the downstream task.
Further, the distribution of prompts are quite different from real text – because the examples are independently generated, the transitions between examples can be very unnatural. How does in-context learning work?
Our understanding is limited since (i) real pretraining data is messy and (ii) in-context learning has so far required large-scale datasets and models.
In this paper, we introduce a simple pretraining distribution which gives rise to in-context learning.
To generate a document, we ﬁrst draw a latent concept θ , which parameterizes the transitions of a Hidden
Markov Model (HMM) (Baum & Petrie, 1966), then sample a sequence of observations (tokens) from the HMM (see Figure 2).
This latent variable structure is common in topic models such as LDA (Blei et al., 2003; Gruber et al., 2007).
To capture long-range coherence in the pretraining data, an LM
must
Ideally the posterior predictive distribution should concentrate on the correct output as the number of in-context examples increases but the main challenge is that prompts are not presented as a set of independent examples but as a concatenated sequence (Figure 1).
Thus, the prompts are from a different distribution than the pretraining distribution.
The canonical asymptotic tool in Bayesian methods is the Bernstein-von Mises theorem (van der Vaart, 1998; Kleijn & van der Vaart, 2012; Gunst & Shcherbakova, 2008), which typically assumes observations are independent and/or drawn from the same distribution as the model, both of which are not satisﬁed.
However, we prove that in-context learning can emerge in the limit of inﬁnite examples when the signal about the latent concept in each promptexampleislargerthantheerrorstemmingfromthemismatchbetweenthepretrainingandprompt distributions.
Additionally, we prove that the in-context learning error can be bounded by a function of the length of each example—longer example sequences provide more signal on the latent concept.
As a companion to this theory, we created the G enerative IN - C ontext learning dataset (GINC), which is a small-scale synthetic dataset which we can use to study in-context learning without relying on large-scale, messy real-world data.
We ﬁnd that both Transformer (Vaswani et al., 2017) and
LSTM (Hochreiter & Schmidhuber, 1997) language models trained on such datasets exhibit in-context learning.
Weverifyintuitionsfromthetheory, showingthattheaccuracyofin-contextlearningimproves with the number of examples and example length.
Ablations of the GINC dataset show that the mixtureover-concepts structure in the pretraining distribution is crucial to the emergence of in-context learning.
The experiments also bring up open questions which go beyond our theory, which only studies the pretraining distribution.
We ﬁnd that scaling up the number of model parameters steadily improves the in-context accuracy despite achieving the same pretraining loss, showing that larger models may improve in-context learning beyond increasing the capacity for memorizing the training data better.
I N - CONTEXT LEARNING SETTING
Pretraining distribution.
In our framework, a latent concept θ from a family of concepts Θ deﬁnes a distribution over observed tokens o from a vocabulary O .
To generate a document, we ﬁrst sample a concept from a prior over concepts p ( θ ) and then sample a sequence given the concept.
We model the pretraining distribution as a mixture of token distributions described by latent concepts.
Each pretraining document is a length T sequence:
[Equation]
In our theoretical analysis, we assume p ( o 1 ,...,o T | θ ) is described by the distribution over observations from a Hidden Markov Model (HMM).
The concept θ determines the transition probability matrix of the HMM hidden states h 1 ,...,h T from a hidden state set H .
Prompt distribution.
The prompt distribution p prompt generates prompts for in-context learning.
Each prompt is generated with respect to a prompt concept θ ∗ , which is shared by all the prompt examples.
A prompt example is composed of an input token sequence x (e.g., Albert Einstein was) followed by an output token y (e.g., German).
The prompt is a concatenation of n independent training examples and 1 test input x test , where the goal is the predict the test output y test by predicting the next token.
In particular, the i -th training example
O i consists of an input x i
= O i [1: k − 1] (the ﬁrst k − 1 tokens of the example) followed by an output token y i
= O i
[ k ] at the end 1 . The i -th training example is independently generated as follows:
Generate a start hidden state h start i from a prompt start distribution π prompt .
Given h start i , generate the example sequence O i =[
x i ,y i ] from p ( O i | h start i ,θ ∗ ) , the pretraining distribution conditioned on a prompt concept θ ∗ .
The test example x test = x n +1 is sampled similarly.
The resulting prompt sequence concatenates the training examples (denoted S n ) and test example ( x test ) separated by a special delimiter token o delim : [ S n ,x test ]=[ x 1 ,y 1 ,o delim
[Equation]
Prompts are out-of-distribution.
We provide a simple illustration for why prompts are OOD.
Suppose a pretraining document is a b c d...x y z a b c d...
y z... where the latent concept enforces alphabetical order with high probability.
A prompt based on this concept might be a b \ n x y \ n f g \ n c , with independent examples and a test input c .
The transitions between examples are unnatural (e.g., x rarely follows closely after b ).
Thus, prompts are low probability (OOD) sequences under the pretraining distribution.
Even though the examples were independently generated, we can no longer factorize the examples under the pretraining distribution.
In-context predictor and task.
For in-context learning, the output target y for each example x is sampled according to p prompt ( y | x ) :
[Equation]
(3) where h start test denotes the hidden state corresponding to the ﬁrst token of x test and π prompt ( h start test | x test ) is deﬁned as p prompt ( x test | h starttest ) π prompt ( h starttest ) /p
prompt ( x test ) .
The target distribution differs from the pretraining distribution p on the distribution of h starttest .
The central question is, when the number of examples grows, does the posterior predictive distribution of the pretraining distribution given the prompt concentrate on the correct output?
Towards this, we study the in-context predictor f n ( x test )=
argmax y p ( y | S n ,x test ) and its expected 0-1 error:
[Equation] [Equation]
2. A SSUMPTIONS
We detail the assumptions in our framework, including the structure of delimiters and regularity assumptions.
We ﬁrst assume that there exists a subset of delimiter hidden states D which generates the special delimiter token o delim deterministically.
Assumption 1 (Delimiter hidden states) .
Let the delimiter hidden states D be a subset of H .
For any h ∈D and θ , p ( o delim | h,θ )=1 and for any h (cid:48) / ∈D , p ( o delim | h,θ )=0 .
Thus if we observe the special delimiter token o delim , then the corresponding hidden state must be in
D .
However, it does not reveal the particular hidden state in D nor the latent prompt concept θ ∗ .
The delimiter token is usually a token that can appear in a broad range of contexts (e..g, newline).
The delimiter ideally does not distract from the examples — for example, an adversarial delimiter could look like part of the input x .
To mitigate these scenarios, we assume that no delimiter (e.g., newline) is signiﬁcantly more likely under one concept rather than another.
Assumption 2 (Bound on delimiter transitions) .
For any delimiter state h delim ∈ D and any hidden state h ∈ H , the probability of transitioning to a delimiter hidden state under θ is upper bounded p ( h delim | h,θ ) <
c 1 for any θ ∈ Θ \{ θ ∗ } , and is lower bounded p ( h delim | h,θ ∗ ) > c 2 > 0 for θ ∗ .
Additionally, the start hidden state distribution over delimiters is uniform for all concepts θ : p ( h delim | θ )= 1 |D| for all θ ∈ Θ and h delim ∈D .
We also assume the prompt concept θ ∗ is in the family Θ , which is a broad set of concepts.
Assumption 3 (Well-speciﬁcation) .
The prompt concept θ ∗ is in Θ .
Even though the pretraining distribution is broad, the prompt is still out-of-distribution for the pretraining distribution since it concatenates independent examples.
Finally, if the prompt has zero probability under the prompt concept θ ∗ , then Bayesian inference will not be able to infer the prompt concept as in Section 3.1.
The following are regularity assumptions which mainly ensure that the prompt is not zero probability under θ ∗ .
Assumption 4 (Regularity) .
The pretraining distribution p satisﬁes: 1) Lower bound on transition probability for the prompt concept θ ∗ : for any pair of hidden states h,h (cid:48) ∈ H , p ( h | h (cid:48) ,θ ∗ ) > c 3 >
All tokens can be emitted: for every symbol o , there is some hidden state h ∈ H such that p ( o | h,θ ∗ ) >c 4 > 0 . 3)
The hidden state Markov chain for concept θ has stationary distribution π θ over H . 4)
The prior p ( θ ) has support over the entire concept family Θ and is bounded above everywhere.
We denote the stationary distribution conditioned on x to be π θ ( h | x )=( p ( x | h,θ ) π θ ( h )) /p
( x | θ ) .
T HEORETICAL ANALYSIS
We prove that in the limit of inﬁnite examples, conditioning the pretraining distribution on the prompt results in in-context learning provided that a distinguishability condition holds — the prompt concept θ ∗ is distinct enough from the other concepts in Θ (e.g., when Θ is a discrete set).
In this case, the expected error is optimal.
When distinguishability does not hold (e.g, Θ is continuous-valued), we show that the expected error still decreases with example length.
3. H IGH - LEVEL APPROACH
In-context learning in GPT-3 is often done with greedy decoding or nucleus sampling (Holtzman et al., 2020), which samples high probability completions.
Following this practice, we analyze the most likely prediction over the pretraining distribution conditioned on the prompt, which is sampled from the prompt distribution.
Our goal is to show that argmax y p ( y | S n ,x test ) → argmax y p prompt ( y | x test ) as the number of examples n goes to inﬁnity.
In the following, assume that the prompt has non-zero probability under the pretraining distribution p given θ ∗ , meaning that p ( S n ,x test | θ ∗ ) > 0 .
We expand p ( y | S n ,x test ) to analyze its limit:
[Equation]
(Bayes’ rule, drop the constant
[Equation]
(Law of total prob, Markov property, divide by p [Equation] S n ,x test | θ ∗ [Equation]
(a constant))
[Equation]
where r n ( θ ) = 1 n log p ( S n ,x test | θ ) p ( S n ,x test | θ ∗ ) .
In Theorem 1, we prove that under a distinguishability condition, exp( n · r n ( θ ))
→ 0 for all concepts θ except the prompt concept θ ∗ , where exp( n · r n ( θ ∗ ))
= 1 .
The only nonzero term in the integral is when θ = θ ∗ , and thus the prompt concept is “selected” as a consequence of Bayesian inference.
Note that we can exchange limits and integrals here since the probabilities involved in the ﬁrst steps are bounded (dominated convergence).
In Lemma 3, we show that the distribution of the hidden start state for the test example p ( h starttest | S n ,x test ,θ ∗ ) converges to the stationary distribution π θ ∗ ( h starttest | x test ) conditioned on x test .
Putting these together with Equation 6, we have that the in-context predictor infers prompt concept θ ∗ :
[Equation]
Recall that the test label is sampled from p prompt ( y | x test ) =
E h starttest ∼ π prompt
( h starttest | x test )
[ p ( y | x test ,h starttest ,θ ∗ )] .
If the prompt start distribution π prompt equals the stationary distribution π θ ∗ , then the conditional state distributions π prompt ( h starttest | x test )
=
π θ ∗ ( h starttest | x test ) are the same between pretraining and prompting.
Thus, the in-context predictor makes the optimal prediction when π prompt =
π θ ∗ .
3. H EURISTIC DERIVATION
Recall from Section 3.1 that if exp( n · r n ( θ )) → 0 for all θ (cid:54) = θ ∗ , then Bayesian inference “selects” the prompt concept through marginalization.
To do this, we focus on showing that r n ( θ ) , the average log-likelihood ratio between θ and θ ∗ , converges to a negative constant, and thus nr n goes to −∞ .
The main technical challenge is to handle the sequence-of-examples structure of the prompt, which makes all the examples dependent with respect to the pretraining distribution.
Our approach uses properties of delimiter tokens to approximately factorize the examples, with constant error per example.
We let O ex i =[
o delim i − 1 ,
O i ] be the i -th input-output pair and the previous delimiter together for i> 1 and deﬁne O
ex 1 = O 1 .
Expanding the likelihood term inside r n ( θ ) , our goal is to show
[Equation]
To show this, we expand p ( S n | θ ) with the chain rule, and with Assumption 4 (to bound p ( x test | S n ,θ ) by O (1) )
it can be shown that
[Equation]
We then marginalize p ( O ex i | O ex 1: i − 1 ,θ ) over the hidden state h delim
i − 1 corresponding to the delimiter in
O ex i =[
o delim i − 1 ,O i ] :
[Equation]
While summing over H would be a trivial equality, we can replace H with the set of delimiter hidden states D since p ( h delim i − 1 | O ex 1: i − 1 ,θ )=0 for non-delimiter hidden states (Assumption 1).
We used in the ﬁrst equality that O ex 1: i − 1 → h delim i − 1 → O ex i forms a Markov chain and p ( o delim i − 1 | h delim i − 1 )=
1 (Assumption 1) to change O ex i to O i .
Finally, we can show using properties of delimiter hidden states (Assumption 2) that p ( h delim i − 1 |
O ex 1: i − 1 ,θ )
≈ O (1) p ( h delim
i − 1 | θ ) , which allows us to marginalize out h delim i − 1 in the second step.
Therefore, we can upper bound r n ( θ ) as
[Equation]
Since all tokens after the start token in an example O are drawn from the pretraining distribution p conditioned on the prompt concept θ ∗ , the expectation in Equation 11 becomes a negative KL divergence for tokens 2 through k .
There are O ( k ) KL terms and only O (1) other error terms, which come from the distribution mismatch between the prompt and pretraining distributions.
If the KL terms are larger than the error terms, then r n ( θ ) has a negative limit.
If this holds for all θ (cid:54)
= θ ∗ , then we have exp( n · r n ( θ ))
→ 0 for all θ (cid:54)
= θ ∗ , enabling in-context learning.
3. F ORMAL RESULTS 3.3. I N - CONTEXT LEARNING UNDER DISTINGUISHABILITY
We deﬁne a distinguishability condition which formalizes when in-context learning occurs.
Letting p jθ be the output distribution of the j -th token given the previous tokens, the condition depends on KL divergence between p jθ ∗ and p jθ as well as error terms (cid:15) θ start and (cid:15) θ delim coming from the distribution mismatch between the prompt and pretraining distributions at the start and delimiter token for each example:
[Equation]
Condition 1 (Distinguishability) .
We deﬁne θ ∗ to be distinguishable if for all θ ∈ Θ ,θ (cid:54)
= θ ∗ ,
[Equation]
This condition operates on the level of one example.
When the signal from KL divergence (LHS) is larger than the error terms, Equation 14 is satisﬁed.
For longer example lengths k , the LHS increases and the condition is more likely satisﬁed.
Intuitively, longer example lengths increases the proportion of the prompt sampled from the pretraining distribution.
Under Condition 1, we show that the in-context predictor asymptotically achieves the optimal expected error.
Theorem 1.
Assume the assumptions in Section 2.1 hold.
If Condition 1 holds, then as n → ∞ the prediction according to the pretraining distribution is
[Equation]
As a corollary, if the prompt start distribution π prompt is the stationary distribution π θ ∗ of the hidden state chain for θ ∗ , then the in-context predictor f n achieves the optimal expected 0-1 error:
[Equation]
[Equation] When the prompt start distribution π prompt matches the stationary distribution π θ ∗ , the in-context predictor predicts the most likely token from the same predictive distribution as the test example (Equation 3).
When this is not true, there may be additional distribution shift.
3.3. N ON - DISTINGUISHABLE CASE
Even when the distinguishability condition does not hold, can we give an error guarantee?
Note that
Condition 1 fails when there is some θ (cid:54)
= θ ∗ for which the KL between θ and θ ∗ is less than the error terms.
However, this means that the output distributions of θ and θ ∗ are close in KL.
We leverage this to prove that the expected 0-1 error decreases with the example length k
under two different settings.
Continuity.
Our ﬁrst result relies on a continuity assumption between the concept parameter and its corresponding output distribution.
Our assumption is based on prior works (Kleijn & van der Vaart, 2012), where the KL divergence is assumed to have a 2nd-order Taylor expansion.
Theorem 2.
Let the set of θ which does not satisfy Equation 14 in Condition 1 to be B and assume π prompt =
π θ ∗ .
Assume that KL divergences have a 2nd-order Taylor expansion 2 around θ ∗ :
[Equation]
(17) where I j,θ ∗ is the Fisher information matrix of the j -th token distribution with respect to θ ∗ .
Let γ θ ∗ = max j λ max ( I j,θ ∗ ) min jλ min ( I j,θ ∗ ) where λ max ,λ min return the largest and smallest eigenvalues.
Then as n → ∞ , the 0-1 risk of the in-context learning predictor f n is bounded as
[Equation]
where g ( δ )= 12 ((1 − δ ) log(1 − δ )
+(1+ δ )log(1+ δ )) is a calibration function (Steinwart, 2007; ´Avila
Pires & Szepesv´ari, 2016) for the multiclass logistic loss for δ ∈ [0 , 1) .
Varying-length test examples.
In the setting where the length of x test is random (uniformly from 2 to k ), we can give a similar error guarantee without the continuity assumption from Theorem 2.
Theorem 3.
Let the set of θ which does not satisfy Equation 14 in Condition 1 to be B and assume π prompt =
π θ ∗ .
Let the length of the test example x test be uniformly distributed between 2 and k , for k ≥ 2 .
As n →∞ , the 0-1 risk of the in-context learning predictor f n is bounded as
[Equation]
Instead of measuring only the error at the k -th token, we average the prediction error on the 2nd to k -th tokens.
However, we leave bridging the mismatch between training examples, which are consistently length k , and test examples, which have random length, to future work.
S IMULATIONS
We generate the GINC dataset from a mixture of HMMs and show that Transformers (Vaswani et al., 2017) and LSTMs (Hochreiter & Schmidhuber, 1997) trained on GINC exhibit in-context learning.
4. GINC PRETRAINING DATASET
We consider a pretraining distribution that is a mixture of HMMs with an interpretable hidden state structure and emission distribution (see Appendix E for details and a concrete example).
The HMM hidden state h t =[ s t ,v t ] at time t is composed of an entity v t ∈{ 1 ,..., |V|} (e.g., Einstein) and a property s t ∈{ 1 ,..., |S|} (e.g., nationality, ﬁrst name, last name, other grammatical tokens).
We model the entities and properties as independent Markov chains (i.e., a factorial HMM (Ghahramani & Jordan, 1997)), while the emissions depend on both.
In pretraining documents, we expect that the entities (e.g., Einstein) change slowly over time while and their properties (e.g., nationality) change quickly to generate natural sentences.
We implement this by ensuring that the probability of transitioning to the same entity index in the next step is at least 0.9.
The emission distribution depends on a memory matrix M with |V| rows and |S| columns (Figure 2).
At step t , we use the entity v t and property s t to index into the memory matrix.
In particular, the observed tokens are deterministic with p ( o t | h t )
=
if o t = M [ v t ,s t ] .
Concept parameter.
The concept parameter is the property transition matrix, which encodes many relations between properties to generate a diverse pretraining document.
The entity transition matrix is the same for all concepts.
During in-context learning, the prompt start distribution and the concept [Figure] [Figure] together determine the in-context task.
While the prompt concept θ ∗ may have a high probability transition between names and nationalities, another concept may tend to transition to birthplaces after names. Here, if the prompt start distribution focuses on names (with random entities for each example), then the prompt examples (generated according to θ ∗ ) will largely focus on the transition between names and nationalities, visiting part of the total hidden state space. In our experiments, we deﬁne a uniform mixture of HMMs over 5 concepts to generate 1000 documents with ˜10 million tokens total
[Figure]
4. GINC PROMPT DISTRIBUTION
We generate prompts according to the theoretical setup, where we ﬁrst sample a concept θ uniformly at random, then use it to generate all the prompt examples.
The prompt start distribution is chosen to be uniform over entities but with a ﬁxed starting property that is chosen randomly for each prompt, for consistency in the task.
We generate prompts with 0 to 64 training examples and example lengths k
= { 3 , 5 , 8 , 10 } (2500 prompts for each setting).
The target token y test is taken to be the most likely output argmax y p prompt ( y | x test ) instead of sampling to reduce the intrinsic error.
4. S IMULATION RESULTS
We train GPT-2-based Transformer (Radford et al., 2019) and LSTM language models on three versions of the GINC dataset with vocabulary sizes 50, 100, and 150, then evaluate the in-context accuracy (see Appendix E.2, E.3 for model details).
We average all results over 5 pretraining runs.
Figure 3 shows that for both Transformer and LSTMs, in-context accuracy improves as the number of prompt examples n and the example length k increase, verifying our theory.
Ablationsandout-of-distributionconcepts.
Weablatetheroleofthemixture-of-conceptsstructure in GINC.
In Figure 4 (left), we pretrain a 4 layer Transformer on data with only one concept (removing the prior) from Θ , resulting in in-context accuracy curves that do not increase with the number of examples.
Figure 4 (middle) shows that when pretraining on random data, which contains all possible transitions between tokens, in-context learning also fails.
Therefore, the mixture-of-concepts structure is important to in-context learning and simply seeing a diverse set of token transitions does not enable incontext learning.
Finally, we test a 4 layer Transformer pretrained on GINC on randomly generated concepts that are not in the pretraining family of concepts.
We generate 5 additional concepts and use these to generate prompts.
Figure 4 (right) shows that in-context learning also fails for these novel concepts.
Effect of model size and architecture.
Figure 5 shows that increasing the Transformer’s size (4, 12, 16 layers) increases in-context accuracy, as in Brown et al.
Table 6 shows that even though larger Transformers may have the same pretraining loss (e.g., 12 and 16 layer Transformers both get 1.33 validation loss for vocab size 50), the in-context accuracy still improves (81% to 85% from 12 to
[Figure] [Figure]
Accuracy improves with scale even though the pretraining loss may be the same.
layers), suggesting that larger models can improve in-context learning beyond increasing memorization capacity. Finally, the architecture also plays a role — LSTMs consistently outperform Transformers on GINC despite having fewer parameters, perhaps due to the similarity between HMMs and LSTMs.
We leave analysis of the effect of model architecture and model scaling as an open question.
D ISCUSSION AND RELATED WORK
Learning via Bayesian inference and misspeciﬁcation.
The Bernstein-von Mises theorem states that under some conditions, the posterior distribution of a parameter in a Bayesian parametric model is asymptotically normal and centered on the MLE (van der Vaart, 1998).
However, this does not apply for in-context learning since the prompt examples are not presented as independent samples but as one contiguous sequence.
We also focus on the posterior predictive distribution rather than the parameter.
Gunst & Shcherbakova (2008) show a Bernstein-von Mises-type result for observations from an HMM, but do not handle observations from a different distribution.
However, we hope to show more precise asymptotic results in the future about the posterior distribution and give results under misspeciﬁcation (Kleijn & van der Vaart, 2012).
Topic models.
Topic models such as LDA (Blei et al., 2003) also use mixture models of text, but learning is typically done through algorithms such as Expectation-Maximization (Dempster et al., 1977), variational inference (Jordan et al., 1999), or MCMC (Metropolis et al., 1953; Hastings, 1970).
Since in-context learning occurs as a byproduct of a forward pass, we focus on learning as a natural result of marginalization in Bayesian inference without an explicit inference algorithm.
Modeling the pretraining distribution with HMMs.
Wei et al.
(2021a) also use an HMM in their generative model of the pre-training distribution.
However, their goal is different – they analyze how pre-trained representations learned with masked language modeling (Devlin et al., 2019; Liu et al., 2019; Lewis et al., 2020; Clark et al., 2020) can be adapted for downstream tasks through linear probing and prompt tuning (Li & Liang, 2021; Lester et al., 2021) rather than in-context learning.
Bridging the mismatch between pretraining and prompting.
Prior work shows that ﬁnetuning LMs on text with a prompting format improves its zero-shot performance (Wei et al., 2021b; Sanh et al., 2021).
Relatedly, Jiang et al.
; Schick & Sch¨utze (2021); Shin et al. ; Gao et al. (2021) optimize prompt templates for better few-shot ﬁnetuning.
Zhao et al.
; Holtzman et al.
(2021) improve in-context accuracy via calibration or renormalization, a form of adaptation to the prompt distribution.
These methods support our theoretical intuitions that reducing the prompt distribution mismatch would improve in-context learning.
Meta-learning.
Meta-learning methods that train a sequence model to learn from examples (Ravi &
Larochelle, 2017) also performs learning during the forward pass of a model.
However, meta-learning models are trained to learn, while in-context learning emerges from LM pretraining.
C ONCLUSION
We cast in-context learning as implicit Bayesian inference, where the pretrained LM implicitly infers a concept when making a prediction.
We show that in-context learning occurs when the pre-training distribution is a mixture of HMMs.
Our work provides a ﬁrst step towards understanding in-context learning, which we hope will provide insight for improving pretraining and prompting.
E THICS
Since the datasets in this paper are synthetically generated, there are no privacy/security issues to report.
We investigate the emergent behaviors in large models, which will ultimately require careful consideration of environmental factors (energy usage), bias stemming from pretraining data, and others.
R EPRODUCIBILITY
The source code for the experiments are included as a supplementary material.
Proofs for theoretical results are included in the appendix, and assumptions are in Section 2.1.
The datasets used in the experiments can be generated from the supplementary code.
The experimental results are averaged over 5 seeds and 95% intervals are reported.
A F RAMEWORK DETAILS
Prompt distribution details. For in-context learning, we sample a prompt from a new distribution p prompt , which consists of n independent training examples and 1 test example.
We ﬁrst sample n hidden segments H of length k by sampling the ﬁrst element h start
= H [1] from a prompt start distribution π prompt .
Then, we sample the rest of the segment H seg
=
H [2: k ] from the hidden transition distribution of the pretraining distribution p corresponding to a particular concept θ ∗ :
[Equation]
To end each example (except the test example), we sample n delimiters h delim ∈D from p delimprompt :
[Equation]
Conditioned on hidden variables H i and h delim i , we sample the observed tokens
O i
=[
o i, 1 ,...,o i,k ] and o delim i respectively from the pre-training distribution:
[Equation]
[Equation]
The “input” for each example is x i
= O i [1: k − 1] and the “output” is y i
= O i
[ k ] .
Taking S to be the sequence of training examples (without the test example), the resulting prompt sequence is [ S n ,x test ]=[
O 1 ,o delim,...
,O n ,o delim n ,x test ]=[ x 1 ,y 1 ,o delim,x 2 ,y 2 ,o delim,.
[Equation]
where x test = x n +1 =
O n +1 [1: k − 1] is sampled via the same process but with k − 1 elements.
B L EMMAS FOR T HEOREM
The following propositions, which lower bound the probability of a delimiter token and probability of an example under θ ∗ , are direct corollaries of the assumptions.
Proposition 1.
For all i , we have p ( h delim i
[Equation]
o delim
[Equation]
[Equation]
>c 2 and p ( h
[Equation]
i
[Equation]
O
[Equation]
o delim
[Equation]
,...,O i
[Equation]
< c 1 .
Proof.
By Assumption 2,
[Equation]
Similarly, p ( h delim
i
[Equation]
[Equation]
Proposition 2.
The density of an example is lower bounded for θ ∗ : there is some c 5 > 0 such that p ( O i | h start i ,h j,l ,θ ∗ ) >
c 5 for all i and future hidden states h j,l , for any l and j >i .
Proof.
By Assumption 4, we have
[Equation]
for some H i .
We have
[Equation]
which lower bounds the terms in the numerator by c 3 (marginalizing over previous hidden states), and upper bounding the denominator by 1.
Setting c 5 =( c 4 ) k c 23 ﬁnishes the proof.
B. M IXING TIME LEMMA
The following lemma is used to show that conditioning on a variable that is at the end of the Markov chain has small effect on variables that are far away from the end of the chain (further than the mixing time of the chain).
Lemma 1.
Suppose Z 1 ,...,Z t are sampled from a Markov chain with stationary distribution π and (cid:15) - mixing time t mix .
Suppose A is any random variable where p ( Z t | Z i ,A )=
p ( Z t | Z i ) .
Then for i<t − t mix , [Equation] 32) where [Equation]
[Equation]
Proof.
Using the deﬁnition of mixing time we have
[Equation]
so that
[Equation]
Using Bayes’ rule,
[Equation]
using Equation 36 and p ( Z t )=
E Z 1 [ p ( Z t | Z 1 )] >π ( Z t ) − (cid:15) by the deﬁnition of mixing time.
We also have
[Equation]
Therefore, conditioning on Z t does not affect the distribution of Z i much, and as the error (cid:15) decreases, the effect becomes smaller.
B. L EMMA FOR LIMITING DISTRIBUTION OF HIDDEN START STATE
We aim to show a limiting distribution for the hidden start state. As an intermediate step, we ﬁrst show lower and upper bounds for the likelihood of the prompt examples given a hidden start state for the test example.
Lemma 2.
Deﬁne (cid:15) θ ( t ) to be the bound on the TV distance between the distribution on h starttest and π θ for (cid:15) θ ( t ) -mixing time t .
This bound (cid:15) θ ( t ) monotonically decreases to 0 as t →∞ .
Let O ex i =[
o delim i − 1 ,
O i ] be the i -th observation segment and the previous delimiter together for i> 1 and deﬁne
O ex 1 = O 1 .
In the setting of Section 2.1, we have the following upper and lower bounds for any t ≤ n :
[Equation]
Proof.
We analyze the likelihood term:
[Equation]
where in the last step we use our assumption p ( o delim | h delim )
= 1 (Assumption 1) for any delimiter token to change O ex i to O i and o delim n to h delim n .
Note that we can take the sum over hidden states to be over D because of our assumptions on the delimiter hidden states.
First work with
[Equation]
Focusing on the last term, we use Lemma 1 to show that for i<n − t ,
[Equation]
where C θ 1 ( (cid:15) θ ( t )) ,C θ 2 ( (cid:15) θ ( t )) are as deﬁned in Lemma 1, and both quantities converge to 1 as (cid:15) ( t ) → 0 .
Thus, for i<n − t ,
[Equation]
For the other term, we also invoke Lemma 1 to show that for i<n − t
[Equation]
We can put these together to bound the terms in the original sum (Equation 49) for i<n − t :
[Equation]
While these steps result in an upper bound, using C θ 1 ( (cid:15) θ ( t )) yields the lower bound.
The following lemma is used to show that the hidden start state distribution converges to the stationary distribution with inﬁnite examples in the prompt.
Lemma 3.
In the setting of Section 2.1, we have as n →∞ that
[Equation]
Proof.
First, by Bayes’ rule, we have
[Equation]
We will analyze the following ratio for h starttest where p ( S n | h starttest ,θ ) > 0 and show it converges to a constant wrt h starttest :
[Equation]
which implies that
[Equation]
Plugging this into Equation 58 (after taking the reciprocal) and noting that p ( h starttest | θ )
→ π θ ( h starttest ) show that
[Equation]
Note that we do not take the limit of the x test terms since x test is observed.
Since the limiting distribution is a valid distribution (the simplex is closed) and should sum to 1, we can apply this constraint to show that the constant factor is 1, which shows the result.
Now it remains to show that
[Equation]
Redeﬁning the ratio, we have
[Equation]
As in Lemma [Equation] let O ex i =[
o delim i − 1 ,
O i ] be the i -th observation segment and the previous delimiter together for i> 1 and deﬁne
O ex 1 = O 1 .
Deﬁne (cid:15) θ ( t ) to be the bound on the TV distance between the distribution on h starttest vs. π θ , for (cid:15) θ ( t ) -mixing time t .
Using Lemma [Equation] we have the following upper and lower bounds:
[Equation]
We used both Propositions 1 and 2 in the terms involving c 1 ,c 2 and c 5 .
Note that in the ﬁrst line, the sum can must be over the set of delimiter states D by using the assumption that observing a delimiter o delim implies that the corresponding hidden state h delim must be in D .
where we note that t, C θ 2 ( (cid:15) θ ( t )) , C θ 1 ( (cid:15) θ ( t )) are constants.
The ﬁrst of the two log ratios involving a marginalization over h delim i − 1 is 0 since the numerator and denominator are the same.
Flipping the order of C θ 2 ( (cid:15) θ ( t )) ,C θ 1 ( (cid:15) θ ( t )) , c 1 ,c 2 and 1 vs. c 5 gives a lower bound:
[Equation]
Together, this shows that
[Equation]
as desired.
C P ROOF OF T HEOREM
Proof.
We analyze the most likely prediction over the pretraining distribution conditioned on the prompt argmax y p
( y | S n ,x test ) .
[Equation]
Deﬁning the following quantity,
[Equation]
we will show that under distinguishability for all θ (cid:54) = θ ∗ , r n ( θ ) converges to a negative constant such that
[Equation]
for θ (cid:54) = θ ∗ [Equation] whereas this ratio is always 1 for θ = θ ∗ .
This will then “select” the desired prompt concept through marginalization.
Taking the limit of Equation
[Equation]
we have
[Equation]
as desired, where we can swap the limit and integral by the dominated convergence theorem (probabilities are bounded above by 1, and p ( θ ) are bounded above by Assumption 4), and we used Lemma 3 to show the convergence to the stationary distribution.
Now it remains to show that r n ( θ ) converges to a negative constant for θ (cid:54)
= θ ∗ .
Let O ex i =[
o delim i − 1 ,
O i ] be the i -th observation segment and the previous delimiter together for i > 1 and deﬁne
O ex 1 = O 1 .
Expanding the numerator of the ratio in r n ( θ ) , we have
[Equation]
Note that in the last line, the sum must be over the set of delimiter states D by using the assumption that observing a delimiter o delim implies that the corresponding hidden state h delim must be in D .
We restrict our attention to θ where p ( S n ,x test | θ ) > 0 , since otherwise θ does not affect the prediction.
Expanding r n ( θ ) [Equation] we have the following upper bound:
[Equation]
[Equation] We used both Propositions 1 and 2 in the terms involving c 1 ,c 2 (bounding the probability of h delim hidden states) and c 5 (bounding the probability of x test ).
Note that in the second line, the sum can must be over the set of delimiter states D by using the assumption that observing a delimiter o delim implies that the corresponding hidden state h delim must be in D .
In the second to last line, we used the assumption that p ( h delim | θ )= 1 |D| for all θ,h delim , multiplying both top and bottom of the ratio by this constant.
The expectation becomes a KL divergence for token 2 through k of each segment (the ﬁrst token is different due to the starting distribution π prompt ).
Aiming to decompose the expectation into a sum over the k tokens, we write p jθ ( o )
=
p ( O [ j ]= o | O [1: j − 1] ,θ ) .
Then we have that
[Equation]
The second term is an error term that depends on how different the starting prompt distribution π prompt (which is part of p prompt ) is to the pre-training distribution.
The third term is an error term that comes from the delimiter transitions.
The bound is negative when the sum of KL terms is larger in magnitude than the error terms.
Note that as k becomes larger, the number of observations of θ ∗ “overpowers” the OOD transitions in the prompt distribution.
This condition is equivalent to the disinguishability condition (Condition 1).
By assumption, for θ (cid:54)
=
θ ∗ the Condition 1 holds, and thus
[Equation]
since r n ( θ ) has a negative, constant limit.
Note that exp( n · r n ( θ ∗ ))
=1 for θ ∗ .
D N ON - DISTINGUISHABLE CASE
When Condition 1 is unsatisﬁed, Equation 14), gives an upper bound on the sum of KL divergences for the next token distributions given different-length histories.
In contrast, the in-context task only measures the accuracy of the last ( k -th) token.
The main challenge is to relate the different-length histories to each other to give a more precise bound for the error on the in-context task (last token).
Before addressing this challenge, we give the following lemma, which leverages the result of ´Avila
Pires & Szepesv ´ ari (2016); Steinwart (2007) to relate a bound on the KL divergence to 0-1 loss.
Lemma 4.
Let the set of θ which does not satisfy Condition 1 to be B .
Assume that KL
[Equation]
( y test | x test ,θ ) is bounded above for all θ .
If
[Equation]
then
[Equation]
where
[Equation]
is a calibration function for the multiclass logistic loss for δ ∈ [0 , 1] .
Proof.
First, we note that we can study the 0-1 risk of the limiting predictor:
[Equation]
where in the last step we use that since the output space of f n is discrete and the probabilities that the in-context predictor takes an argmax over converges, then for N large enough, f N ( x test )=
lim n →∞ f n ( x test ) .
Note that for every input x test , the limiting in-context learning predictor outputs the argmax of a predictive distribution which can be a mixture of predictive distributions over B :
[Equation]
for some distribution q over B .
The KL divergence between this mixture and the prompt concept is bounded by the KL divergence of any one θ ∈B , due to the convexity of KL:
[Equation]
where we can exchange the order of expectations since the KL is bounded (dominated convergence).
From the KL bound KL ( p ( y test | x test ,θ ∗ ) (cid:107) p
( y test | x test ,θ ) , we thus have
[Equation]
(111) where L CE ( θ ) = − E x test ∼ p prompt
[ p ( y test | x test ,θ ∗ )log p ( y test | x test ,θ )] is the multiclass logistic risk, and
L CE ( θ ∗ ) is the entropy of the predictive distribution.
Applying Theorem 2.2 and 5.11 of ´Avila Pires & Szepesv ´ ari (2016), g is a calibration function for the multiclass logistic loss, and allows us to convert the surrogate risk bound to a bound on the 0-1 loss, giving the result.
Note that we have zero approximation error here, since θ ∗ ∈ Θ .
Note that g − 1 is roughly linear in (cid:15) for (cid:15) smaller than 0.7, where the bound is non-vacuous.
D. P ROOF OF T HEOREM
Proof.
By the continuity assumption, we have for any θ in B that
[Equation]
To bound the last KL term,
[Equation]
[Equation] [Equation] Rearranging and noting that KL k
[Equation]
= E x [Equation]
p
[Equation]
y
[Equation]
| x test ,θ ))] , we have
[Equation]
Plugging into Lemma 4 gives the result.
D. P ROOF OF T HEOREM
Note that Condition 1 ensures that the sum of KL divergences between positions within a k -length input is bounded.
This means that we have a bound over not only the last-position KL divergence, but also for all the intermediate tokens.
Intuitively, the random length test example allows the in-context predictor to “take credit” for ﬁtting the intermediate tokens.
The proof is immediate given the KL bound and Lemma 4, given that the length of x test is uniformly random between 2 to k .
Proof.
Let the set of θ that does not satisfy Condition 1 to be B .
We have for any θ in B that
[Equation]
by Theorem 1 and Condition 1.
Plugging this into Lemma 4 gives the result.
E E XPERIMENTAL DETAILS E. GINC DATASET
The GINC dataset is generated from a mixture of HMMs.
These HMMs output tokens from a vocabulary of size in { 50 , 100 , 150 } .
The vocabulary contains a special delimiter token (backslash – see Figure ?? , designated to be index 1.
The vocabulary is generated as combinations of letters starting from a to z, then aa to az, and so on.
All sequences are tokenized by splitting on whitespaces.
Memory matrix.
The shared memory matrix has 10 entities and 10 properties, totaling 100 entries (corresponding to 100 hidden states).
The ﬁrst column of the memory matrix is ﬁxed to be the delimiter token, while each remaining entry of the shared memory matrix is populated with a token sampled uniformly from the vocabulary.
Transition matrix for properties.
We generate 5 property transition matrices, one for each component of the HMM mixture.
We generate each transition matrix via a convex combination of 100 random permutation matrices.
The weights of the convex combination are randomly generated as
[Equation]
where u ∈ R 100 has uniform random entries in [0 , 1] and t is a temperature parameter, set to 0.1. [Figure]
Transition matrix for entities.
The entity transition matrix is shared between all the HMMs that consistute the mixture.
The entity transition matrix is generated in the same way as the property transition matrices, except with one additional step.
Letting T be a transition matrix sampled in the same way as a property transition matrix,
In pretraining documents, we expect that the entities (e.g., Einstein) change slowly over time while and the properties of the entity (e.g., their occupation) change quickly with some pattern to generate natural sentences.
We implement this by ensuring that the probability of transitioning to the same entity index in the next step is at least 0.9.
The ﬁnal entity transition matrix is then 0 .
1 T +0 .
9 I where
I is the identity matrix.
Although we add the diagonal component for added realism, we also consider not adding this component.
Figure 8 shows in-context learning curves for a small (4 layer)
Transformer trained on data that does not add the diagonal component (we check this for vocabulary sizes 50, 100, and 150).
In-context learning still works in this case, although not as well for the 50 vocab size case.
Start distribution.
The starting distribution for the hidden states in all HMMs in the mixture are close to uniform.
We generate the start distribution as softmax (( u − 0 . 5) /t ) for random vector u with entries uniformly from [0 , 1] and temperature t =10 .
In the pretraining documents, we only sample from the start distribution in the beginning of the document.
Prompt distribution.
The prompts are sampled according to the theoretical setup.
The prompt start distribution for each prompt is chosen to deterministically output one property that was randomly chosen to be the starting property for each example in the prompt.
Each prompt has an independently sampled starting property.
Each prompt also generates examples from one particular HMM in the mixture (which HMM this is is chosen randomly).
Given the starting property, we sample k tokens from the HMM.
Finally, we append the delimiter token for the example.
We repeat this process for each example in the prompt, concatenating all examples.
The label is generated as
[Equation]
under the prompt concept θ ∗ .
This differs from the theory, which samples y test instead of taking it to be the most likely token.
However, there can be a large amount of intrinsic error that sampling introduces.
We deﬁne the label this way in the simulations to remove the intrinsic error from sampling.
Example of prompt generation.
In the example in Figure 7 (right), the starting property is ﬁxed to be 5 (for example).
The ﬁrst token (l) is generated by sampling a random entity index (3), and indexing into the memory matrix returns l. Running the hidden state chain of the HMM forward gives the next pair of property and entity.
Since the entity Markov chain changes slowly, the entity is still 3 in the next step – however, the property has changed to 4, and indexing into the memory matrix outputs the next token (aw). Following this same process to generate the third token (the output for the ﬁrst example), we ﬁnish generatingoneexample. Toendtheexample, weappendadelimiter(backslash). Werepeatthisexample generation process for all the examples, except for the test example at the end, where we do not generate the last token.
We condition the HMM on the generated prompt to compute the posterior distribution over the next token p ( y | x test ,θ ∗ ) .
We take the argmax of this distribution to be the ground truth label.
Dataset details. The dataset contains 1000 training documents and 100 validation documents, where training documents have 10240 tokens and validation documents have 1024 tokens.
Each document is generated by ﬁrst selecting one of the HMMs from the mixture uniformly at random, then generating 10240 tokens from the HMM.
We also generate 2500 in-context prompts for each (example length,number of examples) pair, for example lengths k
= [3 , 5 , 8 , 10] and number of examples n = [0 , 1 , 2 , 4 , 8 , 16 , 32 , 64] .
Each prompt is generated using a random HMM in the mixture.
E. T RANSFORMER DETAILS
Our Transformer models are based on the GPT-2 architectures with 4, 12, and 16 layers respectively, with 12 attention heads, 768 dimensional embeddings, residual/embedding/attention dropout set to 0.1, and a context window of 1024.
Other than the number of layers, the other parameters are the default settings from the HuggingFace library (Wolf et al., 2019).
We train for 5 epochs using the AdamW optimizer (Loshchilov & Hutter, 2019; Kingma & Ba, 2015) with a batch size of 8 and a linear learning rate schedule (with 1000 step warmup) up to a learning rate of 8e-4 for the 4 layer and 12 layer model, while for the 16 layer model we start with a constant learning rate of 8e-4 and reduce by a factor of 0.25 whenever the best validation loss does not improve. We tried both learning rate strategies for all models and take the most consistent.
We tuned these models so that the training loss curves between seeds have smaller variability between the runs in terms of the curve shape and when the loss decreases – we found that this is an important indication of stable results.
The models took 50 minutes, 2 hours, 3 hours to train respectively.
The hardware was mainly Titan Xp GPUs, trained and evaluated using 16-bit precision.
All the results are reported with 5 pretraining runs (5 different seeds).
E. LSTM DETAILS
We train an LSTM language model with embedding size 768, hidden layer size 768, and 6 layers.
We use dropout 0.2 and weight decay 1e-5.
The optimizer is AdamW starting with a learning rate of 1e-3, then reducing by a factor of 0.25 whenever the best validation loss does not go down.
We train for a total of 10 epochs, with gradient clipping at norm 1.0.
We use a batch size of 8 and backpropagate through time for 1024 steps (each pretraining data segment is also 1024 tokens).
Each model takes roughly 2 hours to train on Titan Xp GPUs.
[Figure]
[Figure]
E. V ARYING THE VOCABULARY SIZE
To do well on the in-context learning task, the model must both infer the prompt concept and the last
HMM hidden state.
In general, increasing the number of observable symbols makes the in-context task easier by making the inference of the HMM hidden state easier.
With more symbols, each hidden state is more likely to output a different symbol, making the inference problem easier.
This improvement comes despite the number of output classes in the problem (same as the vocabulary size) increasing.
[Figure]
[Figure]
[Figure]
such as LAMBADA (Brown et al., 2020).
The few-shot setting introduces the OOD prompt structure, which can initially lower accuracy.
Figures 9, 10, 11, 12 show in-context learning curves for vocabulary sizes 50, 100, and 150, keeping other hyperparmeters of the dataset the same.
E. S ENSITIVITY TO EXAMPLE ORDERING
In Figure 13 (left), we test the sensitivity of the in-context accuracy on GINC to the ordering of the prompt examples, following Zhao et al.
For this experiment, we consider prompts generated from a single concept and prompt start distribution.
We sample 10 different sets (leading to 10 training set IDs) of 4 examples and generate all 24 possible permutations for each set of examples.
Similarly to the behavior of GPT-3 (Zhao et al., 2021), there is a signiﬁcant variation between example orderings, with a 10–40% gap between different permutations of the same set of examples.
E. Z ERO - SHOT IS SOMETIMES BETTER THAN FEW - SHOT
On some hyperparameter settings of GINC, we ﬁnd that zero-shot performance can be better than few-shot performance in our experiments. This ﬁnding mirrors the behavior of GPT-3 on some datasets (e.g., LAMBADA, HellaSwag, PhysicalQA, RACE-m, CoQA/SAT analogies for smaller [Table]
models (Brown et al., 2020)), where zero-shot performance is better than one-shot performance.
This occurs especially when the transition probabilities in GINC are lower entropy (controlled via a temperature parameter).
For this experiment, we consider GINC with transition matrix temperature parameter 0.01 (instead of 0.1), 12 concepts, and vocabulary size 100.
Zero-shot completions in this case are more possible to predict with only the test input.
Figure 13 (right) shows that in this case, the in-context accuracy is initially worse than zero-shot with a few examples, but can still later stabilize to a higher value.
We hypothesize that this phenomenon occurs since the few-shot setting introduces the OOD prompt structure, thus initially decreasing the accuracy.
We conduct an additional experiment which shows that longer examples improve in-context learning in GPT-3 on the LAMBADA (Paperno et al., 2016) completion task.
Data.
In this experiment, we deﬁne a short version of the LAMBADA test dataset (LAMBADA test-short) which contains only test examples with up to 200–300 characters in length.
We also deﬁne two “training” datasets from which to sample examples for the in-context prompts from.
The short training dataset (LAMBADA train-short) contains examples from the training set that are 200–300 characters in length, which matches the distribution of test-short.
The long training dataset (LAMBADA train-long) contains training examples that are 500–600 characters long.
We cut the number of examples in the larger of the two training datasets so that the two training datasets are equally sized (47 examples).
For each test example, we sample 5 random training examples (5-shot learning).
We also consider equalizing the total length of the prompts in two ways.
First, we consider duplicating the 5 short examples (if the examples are [1,2,3,4,5], duplicating refers to [1,2,3,4,5,1,2,3,4,5]).
This allows for equalizing the total length without increasing the number of examples.
As a skyline comparison, we also consider sampling 10 independent short examples, which contains more input-output pairs for the task.
Result.
Table 1 shows that when evaluating only on LAMBADA test-short, 5-shot in-context learning using LAMBADA train-long improves the test accuracy by almost 1% compared to LAMBADA train-short, despite the long/short distribution mismatch between train and test.
This supports intuitions from our theory.
In comparison, simply increasing the total prompt length by duplicating the short examples does not improve the accuracy.
Intuitively, the longer examples have additional information that is not directly related to mapping between the input and output, but can be leveraged to improve in-context learning by helping the model infer the latent concept.
Using 5 long examples (as opposed to 5 short examples) closes about 56% of the gap between using 5 short examples and 10 independent short examples despite not adding additional examples or task-related information.