Year: 2021

Abstract:
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, we generate a family of small-scale synthetic datasets (GINC) where Transformer and LSTM language models both exhibit in-context learning. Beyond the theory which focuses on the effect of the pretraining distribution, we empirically find that scaling model size improves in-context accuracy even when the pretraining loss is the same.

heading: 1 INTRODUCTION
text: 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 predict the next word given previous words. They demonstrate the surprising ability to do in-context learning, where an LM “learns” to do a 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 to make a prediction. Given the following prompt, 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 using the prompt. Further, the distribution of prompts is quite different from real text – the transitions between the independent examples can be unnatural. Our understanding of how in-context learning works 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 first 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 infer the latent concept across multiple sentences to generate coherent continuations. When conditioning on a prompt, in-context learning occurs when the LM also infers a shared latent concept across
examples to make a prediction. We assume the LM fits the pretraining distribution exactly with enough data and expressivity, so that the question of in-context learning becomes characterizing the conditional distribution of completions given prompts p(output|prompt) under the pretraining distribution, where the prompt is generated from a prompt distribution pprompt. The pretraining distribution marginalizes out the latent concepts (and HMM hidden states), so this conditional distribution is the posterior predictive distribution. Thus, in-context learning can be viewed as implicit Bayesian inference.
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 asserts that the posterior distribution of a latent variable concentrates to the MLE. However, Bernstein-von Mises typically assumes observations are independent and/or drawn from the same distribution as the model, both of which are not satisfied. However, we prove that in-context learning can emerge in the limit of infinite examples when the signal about the latent concept in each prompt example is larger than the error stemming from the mismatch between the pretraining and prompt 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 Generative IN-Context 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 find that both Transformer (Vaswani et al., 2017) and LSTM (Hochreiter & Schmidhuber, 1997) language models trained on such datasets exhibit in-context learning. We verify intuitions from the theory, showing that the accuracy of in-context learning improves 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 find 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.

heading: 2 IN-CONTEXT LEARNING SETTING
text: Pretraining distribution. In our framework, a latent concept θ from a family of concepts Θ defines a distribution over observed tokens o from a vocabularyO. To generate a document, we first 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:
p(o1,...,oT )= ∫ θ∈Θ p(o1,...,oT |θ)p(θ)dθ (1)
In our theoretical analysis, we assume p(o1,...,oT |θ) 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 h1,...,hT from a hidden state setH. Prompt distribution. The prompt distribution pprompt 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 xtest, where the goal is the predict the test output ytest by predicting the next token. In particular, the i-th training example Oi consists of an input xi =Oi[1 : k−1] (the first k−1 tokens of the example) followed by an output token yi = Oi[k] at the end1. The i-th training example is independently generated as follows:
1. Generate a start hidden state hstarti from a prompt start distribution πprompt.
2. Given hstarti , generate the example sequenceOi=[xi,yi] from p(Oi|hstarti ,θ∗), the pretraining distribution conditioned on a prompt concept θ∗.
The test example xtest =xn+1 is sampled similarly. The resulting prompt sequence concatenates the training examples (denoted Sn) and test example (xtest) separated by a special delimiter token odelim:
[Sn,xtest]=[x1,y1,o delim,x2,y2,o delim,...,xn,yn,o delim,xtest]∼pprompt. (2)
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 pprompt(y|x):
ytest∼pprompt(y|xtest)=Ehstarttest ∼πprompt(hstarttest |xtest) [ p(y|xtest,hstarttest ,θ∗) ] . (3)
where hstarttest denotes the hidden state corresponding to the first token of xtest and πprompt(h start test |xtest) is defined as pprompt(xtest|hstarttest )πprompt(hstarttest )/pprompt(xtest).
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 fn(xtest)=argmaxyp(y|Sn,xtest) and its expected 0-1 error:
L0-1(fn)=Extest,ytest∼pprompt [1[fn(xtest) 6=ytest]]. (4)

heading: 2.1 ASSUMPTIONS
text: We detail the assumptions in our framework, including the structure of delimiters and regularity assumptions. We first assume that there exists a subset of delimiter hidden statesD which generates the special delimiter token odelim deterministically. Assumption 1 (Delimiter hidden states). Let the delimiter hidden statesD be a subset ofH. For any h∈D and θ, p(odelim|h,θ)=1 and for any h′ /∈D, p(odelim|h,θ)=0.
Thus if we observe the special delimiter token odelim, then the corresponding hidden state must be in D. However, it does not reveal the particular hidden state inD nor the latent prompt concept θ∗.
1Example length k is fixed for simplicity, but in practice, it will vary — our analysis can be extended to variable k, but we leave this to future work.
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 significantly more likely under one concept rather than another. Assumption 2 (Bound on delimiter transitions). For any delimiter state hdelim ∈ D and any hidden state h ∈ H, the probability of transitioning to a delimiter hidden state under θ is upper bounded p(hdelim|h,θ)< c1 for any θ ∈Θ\{θ∗}, and is lower bounded p(hdelim|h,θ∗)> c2 > 0 for θ∗. Additionally, the start hidden state distribution over delimiters is uniform for all concepts θ: p(hdelim|θ)= 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-specification). 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 satisfies: 1) Lower bound on transition probability for the prompt concept θ∗: for any pair of hidden states h,h′ ∈H, p(h|h′,θ∗)> c3 > 0. 2) All tokens can be emitted: for every symbol o, there is some hidden state h ∈ H such that p(o|h,θ∗)>c4>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|θ).

heading: 3 THEORETICAL ANALYSIS
text: We prove that in the limit of infinite 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.

heading: 3.1 HIGH-LEVEL APPROACH
text: 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 argmaxyp(y|Sn,xtest)→ argmaxypprompt(y|xtest) as the number of examples n goes to infinity.
In the following, assume that the prompt has non-zero probability under the pretraining distribution p given θ∗, meaning that p(Sn,xtest|θ∗)>0. We expand p(y|Sn,xtest) to analyze its limit:
p(y|Sn,xtest)= ∫ θ p(y|Sn,xtest,θ)p(θ|Sn,xtest)dθ
∝ ∫ θ p(y|Sn,xtest,θ)p(Sn,xtest|θ)p(θ)dθ (Bayes’ rule, drop the constant 1 p(Sn,xtest) )
= ∫ θ ∑ hstarttest ∈H p(y|xtest,hstarttest ,θ)p(hstarttest |Sn,xtest,θ) p(Sn,xtest|θ) p(Sn,xtest|θ∗) p(θ)dθ (5)
(Law of total prob, Markov property, divide by p(Sn,xtest|θ∗) (a constant))
= ∫ θ ∑ hstarttest ∈H p(y|xtest,hstarttest ,θ)p(hstarttest |Sn,xtest,θ)exp(n·rn(θ))p(θ)dθ (6)
where rn(θ) = 1n log p(Sn,xtest|θ) p(Sn,xtest|θ∗) . In Theorem 1, we prove that under a distinguishability condition, exp(n ·rn(θ))→ 0 for all concepts θ except the prompt concept θ∗, where exp(n ·rn(θ∗)) = 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 first steps are bounded (dominated convergence). In Lemma 3, we show that the distribution of the hidden start state for the test example p(hstarttest |Sn,xtest,θ∗) converges to the stationary distribution πθ ∗ (hstarttest |xtest) conditioned on xtest. Putting these together with Equation 6, we have that the in-context predictor infers prompt concept θ∗:
argmax y p(y|Sn,xtest)→argmax y Ehstarttest ∼πθ∗ (hstarttest |xtest)[p(y|xtest,h start test ,θ ∗)] (7)
Recall that the test label is sampled from pprompt(y|xtest) = Ehstarttest ∼πprompt(hstarttest |xtest)[p(y|xtest,h start test ,θ ∗)]. If the prompt start distribution πprompt equals the stationary distribution πθ ∗ , then the conditional state distributions πprompt(hstarttest |xtest) = πθ ∗ (hstarttest |xtest) are the same between pretraining and prompting. Thus, the in-context predictor makes the optimal prediction when πprompt =πθ ∗ .

heading: 3.2 HEURISTIC DERIVATION
text: Recall from Section 3.1 that if exp(n·rn(θ))→0 for all θ 6=θ∗, then Bayesian inference “selects” the prompt concept through marginalization. To do this, we focus on showing that rn(θ), the average log-likelihood ratio between θ and θ∗, converges to a negative constant, and thus nrn 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 letOexi =[o delim i−1 ,Oi] be the i-th input-output pair and the previous delimiter together for i>1 and defineOex1 =O1. Expanding the likelihood term inside rn(θ), our goal is to show
p(Sn,xtest|θ)=p(xtest|Sn,θ)p(Sn|θ)≈ n∏ i=1 O(1)p(Oi|θ) (8)
To show this, we expand p(Sn|θ) with the chain rule, and with Assumption 4 (to bound p(xtest|Sn,θ) byO(1)) it can be shown that
p(xtest|Sn,θ)p(Sn|θ)≈ n∏ i=1 O(1)p(Oexi |Oex1:i−1,θ). (9)
We then marginalize p(Oexi |Oex1:i−1,θ) over the hidden state hdelimi−1 corresponding to the delimiter in Oexi =[o delim i−1 ,Oi]:
n∏ i=1 O(1)p(Oexi |Oex1:i−1,θ)= n∏ i=1 O(1) ∑
hdelimi−1∈D
p(Oi|hdelimi−1 ,θ)p(hdelimi−1 |Oex1:i−1,θ)≈ n∏ i=1 O(1)p(Oi|θ)
While summing overHwould be a trivial equality, we can replaceHwith the set of delimiter hidden statesD since p(hdelimi−1 |Oex1:i−1,θ)=0 for non-delimiter hidden states (Assumption 1). We used in the first equality thatOex1:i−1→hdelimi−1 →Oexi forms a Markov chain and p(odelimi−1 |hdelimi−1 )=1 (Assumption 1) to changeOexi toOi. Finally, we can show using properties of delimiter hidden states (Assumption 2) that p(hdelimi−1 |Oex1:i−1,θ)≈O(1)p(hdelimi−1 |θ), which allows us to marginalize out hdelimi−1 in the second step. Therefore, we can upper bound rn(θ) as
rn(θ)≤ 1
n
( O(n)+
n∑ i=1 log p(Oi|θ) p(Oi|θ∗) ) →O(1)+EO∼pprompt [ log p(O|θ) p(O|θ∗) ] . (10)
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 10 becomes a negative KL divergence for tokens 2 through k. There areO(k) KL terms and onlyO(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 rn(θ) has a negative limit. If this holds for all θ 6=θ∗, then we have exp(n·rn(θ))→0 for all θ 6=θ∗, enabling in-context learning.

heading: 3.3 FORMAL RESULTS

heading: 3.3.1 IN-CONTEXT LEARNING UNDER DISTINGUISHABILITY
text: We define a distinguishability condition which formalizes when in-context learning occurs. Letting pjθ be the output distribution of the j-th token given the previous tokens, the condition depends on KL divergence between pjθ∗ and p j θ as well as error terms θ start and θ delim coming from the distribution mismatch
between the prompt and pretraining distributions at the start and delimiter token for each example:
pjθ(o)=p(O[j]=o |O[1 : j−1],θ), KLj(θ ∗‖θ)=EO[1 : j−1]∼pprompt [KL(p j θ∗‖p j θ)] (11)
θstart =EO[1]∼pprompt
[ log
p(O[1]|θ) p(O[1]|θ∗)
] , θdelim =(log(c1)−log(c2)). (12)
Condition 1 (Distinguishability). We define θ∗ to be distinguishable if for all θ∈Θ,θ 6=θ∗, k∑ j=2 KLj(θ ∗‖θ)> θstart+ θdelim. (13)
This condition operates on the level of one example. When the signal from KL divergence (LHS) is larger than the error terms, Equation 13 is satisfied. For longer example lengths k, the LHS increases and the condition is more likely satisfied. 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
argmax y p(y|Sn,xtest)→argmax y Ehstarttest ∼πθ∗ (hstarttest |xtest)[p(y|xtest,h start test ,θ ∗)]. (14)
As a corollary, if the prompt start distribution πprompt is the stationary distribution πθ ∗
of the hidden state chain for θ∗, then the in-context predictor fn achieves the optimal expected 0-1 error:
lim n→∞ L0-1(fn)=inf f L0-1(f). (15)
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.

heading: 3.3.2 NON-DISTINGUISHABLE CASE
text: Even when the distinguishability condition does not hold, can we give an error guarantee? Note that Condition 1 fails when there is some θ 6=θ∗ 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 first 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 13 in Condition 1 to be B and assume πprompt =π θ∗ . Assume that KL divergences have a 2nd-order Taylor expansion2 around θ∗:
∀j, KLj(θ∗‖θ)= 1
2 (θ−θ∗)>Ij,θ∗(θ−θ∗)+O(‖θ−θ∗‖3) (16)
where Ij,θ∗ is the Fisher information matrix of the j-th token distribution with respect to θ∗. Let γθ∗ = maxjλmax(Ij,θ∗ )
minjλmin(Ij,θ∗ ) where λmax,λmin return the largest and smallest eigenvalues. Then as n→∞,
the 0-1 risk of the in-context learning predictor fn is bounded as
lim n→∞ L0-1(fn)≤ inf f L0-1(f)+g
−1 ( O ( γθ∗supθ∈B( θ start+ θ delim)
k−1
)) (17)
where g(δ)= 12 ((1−δ)log(1−δ)+(1+δ)log(1+δ)) is a calibration function (Steinwart, 2007; Ávila Pires & Szepesvári, 2016) for the multiclass logistic loss for δ∈ [0,1). Since the inverse calibration function g−1 is roughly linear in for ≤ 0.7, the excess risk roughly decreases asO(1/k). When the “worst-case condition number” γθ∗ of the Fisher information matrices is smaller (well-conditioned), the error decreases. Intuitively, this means that there is no direction to vary θ∗ in which the output distribution will sharply change. As a consequence, the concepts θ that are not distinguishable (via Equation 13) from the prompt concept θ∗ parameterize distributions that produce similar outputs to the prompt concept and thus achieve a small error.
2This requires the KL divergence to be twice differentiable.
Varying-length test examples. In the setting where the length of xtest 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 13 in Condition 1 to be B and assume πprompt = π
θ∗ . Let the length of the test example xtest be uniformly distributed between 2 and k, for k≥2. As n→∞, the 0-1 risk of the in-context learning predictor fn is bounded as
lim n→∞ L0-1(fn)≤ inf f L0-1(f)+g
−1 ( supθ∈B( θ start+ θ delim)
k−1
) . (18)
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.

heading: 4 SIMULATIONS
text: 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. In the theory, we assumed that the pretrained LM fits the pretraining distribution exactly. Here, we pretrain LMs to approximate the pretraining distribution, showing that the in-context learning properties transfer.

heading: 4.1 GINC PRETRAINING DATASET
text: 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 ht=[st,vt] at time t is composed of an entity vt∈{1,...,|V|} (e.g., Einstein) and a property st∈{1,...,|S|} (e.g., nationality, first 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 matrixM with |V| rows and |S| columns (Figure 2). At step t, we use the entity vt and property st to index into the memory matrix. In particular, the observed tokens are deterministic with p(ot|ht) = 1 if ot=M [vt,st]. This construction satisfies the structure on delimiter states (Assumption 1). We ensure that all the transitions have nonzero probability and use a uniform prior over concepts, satisfying Assumptions 2 and 4.
Concept parameter. The concept parameter is the property transition matrix, while the entity transition matrix is fixed for all concepts. The prompt start distribution and the concept together determine the in-context task. We define a uniform mixture of HMMs over a family Θ of 5 concepts to generate 1000 documents with 1̃0 million tokens total.

heading: 4.2 GINC PROMPT DISTRIBUTION
text: We generate prompts according to the theoretical setup, where we first sample a concept θ uniformly at random from Θ (well-specification, Assumption 3), then use it to generate all the prompt examples.
The prompt start distribution is chosen to be uniform over entities but with a fixed 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 ytest is taken to be the most likely output argmaxypprompt(y|xtest) instead of sampling to reduce the intrinsic error.

heading: 4.3 SIMULATION RESULTS
text: 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.
Ablations and out-of-distribution concepts. We ablate the role of the mixture-of-concepts structure 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. (2020). 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 16 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. This may be related to phenomena from overparameterization and overtraining (Zhang et al., 2017; Power et al., 2021). We leave analysis of this as an open question.

heading: 5 DISCUSSION AND RELATED WORK
text: Learning via Bayesian inference and misspecification. 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 independent under the pretraining
Model # Params Train loss(pretraining) Val loss (pretraining) In-context Acc Vocab size 50, k=10,n=64 Transformer (4 layer) 29M 1.49 1.50 60.2± 5.7 Transformer (12 layer) 85M 1.31 1.33 81.2± 7.1 Transformer (16 layer) 115M 1.31 1.33 84.7± 3.4 LSTM 28M 1.31 1.35 95.8± 1.11 Vocab size 100, k=10,n=64 Transformer (4 layer) 29M 1.58 1.59 67.4± 4.7 Transformer (12 layer) 85M 1.40 1.42 84.6± 3.0 Transformer (16 layer) 115M 1.41 1.43 88.7± 1.6 LSTM 28M 1.43 1.44 95.8± 1.54 Vocab size 150, k=10,n=64 Transformer (4 layer) 29M 1.44 1.45 92.8± 1.9 Transformer (12 layer) 85M 1.27 1.28 98.4± 0.4 Transformer (16 layer) 115M 1.27 1.28 98.1± 0.5 LSTM 28M 1.26 1.31 99.2± 1.06
Figure 6: In-context accuracies (95% intervals) on GINC with vocab sizes (50, 100, 150) for Transformers and LSTMs. Accuracy improves with scale even though the pretraining loss may be the same.
distribution. We also focus on the posterior predictive distribution rather than the posterior of 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 misspecification (Kleijn & van der Vaart, 2012). A possible avenue for extrapolating to unseen concepts is to have an additional latent variable γ that modifies concepts θ. It may be possible to generalize to unseen pairs of (θ, γ) if all the individual values of θ and γ occur in the pretraining data.
Topic models and HMMs. 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. 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 finetuning LMs on text with a prompting format improves its zero-shot performance (Wei et al., 2021b; Sanh et al., 2021). Relatedly, Jiang et al. (2020); Schick & Schütze (2021); Shin et al. (2020); Gao et al. (2021) optimize prompt templates for better few-shot finetuning. Zhao et al. (2021); 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.
Takeaways for real data and models. Although HMMs are not the most realistic model of text, our framework is more general – the main component is the presence of a latent variable (concept) which captures long-range coherence in the pretraining data. Beyond the theory for understanding in-context learning, our GINC experiments show that in-context learning and related phenomena can also occur at a small scale. This is because the complexity of the data (vocabulary size, HMM hidden state size, number of latent concepts) is small, such that the data and models are relatively larger. Because we generate a synthetic dataset, we can also control the latent properties of the data (e.g., out-of-distribution concepts), which can help make predictions about large language models while working at a smaller scale.

heading: 6 CONCLUSION
text: 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 first step towards understanding in-context learning, which we hope will provide insight for improving pretraining and prompting.

heading: 7 ETHICS
text: 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.

heading: 8 REPRODUCIBILITY
text: 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.

heading: A FRAMEWORK DETAILS
text: Prompt distribution details. For in-context learning, we sample a prompt from a new distribution pprompt, which consists of n independent training examples and 1 test example. We first sample n hidden segments H of length k by sampling the first element hstart = H[1] from a prompt start distribution πprompt. Then, we sample the rest of the segmentHseg =H[2 :k] from the hidden transition distribution of the pretraining distribution p corresponding to a particular concept θ∗:
H1,...,Hn, Hi=[hi,1,...,hi,k] (19)
hstarti =Hi[1]∼πprompt, H seg i =Hi[2 :k]∼p(H seg i |h start,θ∗). (20)
To end each example (except the test example), we sample n delimiters hdelim∈D from pdelimprompt:
hdelim1 ,...,h delim n , h delim i ∼pdelimprompt. (21)
Conditioned on hidden variablesHi and hdelimi , we sample the observed tokensOi=[oi,1,...,oi,k] and odelimi respectively from the pre-training distribution:
O1,...,On, Oi∼p(Oi|Hi) (22) odelim1 ,...,o delim n , o delim i ∼p(odelimi |hdelimi ,θ∗) (23)
The “input” for each example is xi=Oi[1 :k−1] and the “output” is yi=Oi[k]. Taking S to be the sequence of training examples (without the test example), the resulting prompt sequence is
[Sn,xtest]=[O1,o delim 1 ,...,On,o delim n ,xtest]=[x1,y1,o delim 1 ,x2,y2,o delim 2 ,...,xn,yn,o delim n ,xtest]∼pprompt
(24)
where xtest =xn+1 =On+1[1 :k−1] is sampled via the same process but with k−1 elements.

heading: B LEMMAS FOR THEOREM 1
text: 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(hdelimi |O1,odelim1 ,...,Oi,θ∗)>c2 and p(hdelimi |O1,odelim1 ,...,Oi,θ)< c1.
Proof. By Assumption 2, p(hdelimi |O1,odelim1 ,...,Oi,θ)= ∑ hi,k p(hdelimi |hi,k)p(hi,k|O1,odelim1 ,...,Oi,θ) (25)
< ∑ hi,k c1p(hi,k|O1,odelim1 ,...,Oi,θ)=c1. (26)
Similarly, p(hdelimi |O1,odelim1 ,...,Oi,θ∗)= ∑ hi,k p(hdelimi |hi,k)p(hi,k|O1,odelim1 ,...,Oi,θ∗) (27)
> ∑ hi,k c2p(hi,k|O1,odelim1 ,...,Oi,θ∗)=c2. (28)
Proposition 2. The density of an example is lower bounded for θ∗: there is some c5 > 0 such that p(Oi|hstarti ,hj,l,θ∗)>c5 for all i and future hidden states hj,l, for any l and j>i.
Proof. By Assumption 4, we have p(Oi|hstarti ,hj,l,θ∗)= ∑ Hi p(Oi|Hi)p(Hi|hstarti ,hj,l,θ∗)>(c4)k (29)
for someHi. We have
p(Hi|hstarti ,hj,l,θ∗)= p(hj,l|H,hstarti ,θ∗)p(H|hstarti ,θ∗)
p(hj,l|hstarti ,θ∗) >c23 (30)
which lower bounds the terms in the numerator by c3 (marginalizing over previous hidden states), and upper bounding the denominator by 1. Setting c5 =(c4)kc23 finishes the proof.

heading: B.1 MIXING TIME LEMMA
text: 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. SupposeZ1,...,Zt are sampled from a Markov chain with stationary distribution π and - mixing time tmix. SupposeA is any random variable where p(Zt|Zi,A)=p(Zt|Zi). Then for i<t−tmix,
p(Zi|A,Zt)∈ [C1( )p(Zi|A),C2( )p(Zi|A)] (31)
where
C1( )= π(Zt)− π(Zt)+
(32)
C2( )= π(Zt)+
π(Zt)− (33)
Proof. Using the definition of mixing time we have
|p(Zt|Zi)−π(Zt)|< , (34)
so that
p(Zi,Zt|A)=p(Zi|A)p(Zt|Zi)∈ [p(Zi|A)(π(Zt)− ),p(Zi|A)(π(Zt)+ )]. (35)
Using Bayes’ rule,
p(Zi|A,Zt)= p(Zt,Zi|A) p(Zt)
(36)
< p(Zi|A)(π(Zt)+ )
π(Zt)− . (37)
using Equation 35 and p(Zt)=EZ1 [p(Zt|Z1)]>π(Zt)− by the definition of mixing time. We also have
p(Zi|A,Zt)= p(Zt,Zi|A) p(Zt)
(38)
> p(Zi|A)(π(Zt)− )
π(Zt)+ . (39)
Therefore, conditioning onZt does not affect the distribution ofZi much, and as the error decreases, the effect becomes smaller.

heading: B.2 LEMMA FOR LIMITING DISTRIBUTION OF HIDDEN START STATE
text: We aim to show a limiting distribution for the hidden start state. As an intermediate step, we first show lower and upper bounds for the likelihood of the prompt examples given a hidden start state for the test example.
Lemma 2. Define θ(t) to be the bound on the TV distance between the distribution on hstarttest and πθ for θ(t)-mixing time t. This bound θ(t) monotonically decreases to 0 as t→∞. LetOexi =[odelimi−1 ,Oi] be the i-th observation segment and the previous delimiter together for i>1 and defineOex1 =O1. In the setting of Section 2.1, we have the following upper and lower bounds for any t≤n:
p(Sn,|hstarttest ,θ)≤p(hdelimn |Oex1:n,hstarttest ,θ) n−t∏ i=1 [ Cθ2 ( θ(t))2 ∑
hdelimi−1∈D
p(Oi|hdelimi−1 ,θ)p(hdelimi−1 |Oex1:i−1,θ) ]
(40)
n∏ i=n−t [ ∑ hdelimi−1∈D p(Oi|hdelimi−1 ,hstarttest ,θ)p(hdelimi−1 |Oex1:i−1,hstarttest ,θ) ]
(41)
p(Sn,xtest|hstarttest ,θ)≥p(xtest|hstarttest ,θ)p(hdelimn |Oex1:n,hstarttest ,θ) n−t∏ i=1 [ Cθ1 ( θ(t))2 ∑
hdelimi−1∈D
p(Oi|hdelimi−1 ,θ)p(hdelimi−1 |Oex1:i−1,θ) ]
(42)
n∏ i=n−t [ ∑ hdelimi−1∈D p(Oi|hdelimi−1 ,hstarttest ,θ)p(hdelimi−1 |Oex1:i−1,hstarttest ,θ) ]
(43)
(44)
Proof. We analyze the likelihood term:
p(Sn|hstarttest ,θ)=p(Sn|hstarttest ,θ) (45)
=p(odelimn |Oex1:n,hstarttest ,θ) n∏ i=1 p(Oexi |Oex1:i−1hstarttest ,θ) (46)
=p(odelimn |Oex1:n,hstarttest ,θ) n∏ i=1 [ ∑ hdelimi−1∈D p(Oexi |hdelimi−1 ,hstarttest ,θ)p(hdelimi−1 |Oex1:i−1,hstarttest ,θ) ]
(47)
=p(hdelimn |Oex1:n,hstarttest ,θ) n∏ i=1 [ ∑ hdelimi−1∈D p(Oi|hdelimi−1 ,hstarttest ,θ)p(hdelimi−1 |Oex1:i−1,hstarttest ,θ) ]
(48)
where in the last step we use our assumption p(odelim|hdelim) = 1 (Assumption 1) for any delimiter token to changeOexi toOi and o delim n to h delim n . Note that we can take the sum over hidden states to be overD because of our assumptions on the delimiter hidden states.
First work with p(Oi|hdelimi−1 ,hstarttest ,θ)= ∑ Hi p(Oi|Hi,hdelimi−1 ,θ)p(Hi|hdelimi−1 ,hstarttest ,θ). (49)
Focusing on the last term, we use Lemma 1 to show that for i<n−t,
p(Hi|hdelimi−1 ,hstarttest ,θ)∈ [Cθ1 ( θ(tn))·p(Hi|hdelimi−1 ,θ),Cθ2 ( θ(t))·p(Hi|hdelimi−1 ,θ)] (50)
whereCθ1 ( θ(t)),Cθ2 ( θ(t)) are as defined in Lemma 1, and both quantities converge to 1 as (t)→0. Thus, for i<n−t,
p(Oi|hdelimi−1 ,hstarttest ,θ)≤Cθ2 ( θ(t))·p(Oi|hdelimi−1 ,θ). (51)
For the other term, we also invoke Lemma 1 to show that for i<n−t
p(hdelimi−1 |Oex1:i−1,hstarttest ,θ)≤Cθ2 ( θ(t))·p(hdelimi−1 |Oex1:i−1,θ). (52)
We can put these together to bound the terms in the original sum (Equation 48) for i<n−t:∑ hdelimi−1∈D p(Oi|hdelimi−1 ,hstarttest ,θ)p(hdelimi−1 |Oex1:i−1,hstarttest ,θ) (53)
≤Cθ2 ( θ(t))2 ∑
hdelimi−1∈D
p(Oi|hdelimi−1 ,θ)p(hdelimi−1 |Oex1:i−1,θ). (54)
While these steps result in an upper bound, usingCθ1 ( θ(t)) yields the lower bound.
The following lemma is used to show that the hidden start state distribution converges to the stationary distribution with infinite examples in the prompt.
Lemma 3. In the setting of Section 2.1, we have as n→∞ that
p(hstarttest |Sn,xtest,θ)→πθ(hstarttest |xtest). (55)
Proof. First, by Bayes’ rule, we have
p(hstarttest |Sn,xtest,θ)= p(Sn,xtest|hstarttest ,θ)p(hstarttest |θ)
p(Sn,xtest|θ) (56)
= p(Sn|hstarttest ,θ)p(xtest|hstarttest )p(hstarttest |θ)
p(Sn,xtest|θ) . (57)
We will analyze the following ratio for hstarttest where p(Sn|hstarttest ,θ) > 0 and show it converges to a constant wrt hstarttest :
p(Sn,|hstarttest ′ ,θ) p(Sn,|hstarttest ,θ) →Θ(1) (58)
which implies that
p(Sn,xtest|θ) p(Sn|hstarttest ,θ)p(xtest|hstarttest ) =
∑ hstarttest ′p(Sn|hstarttest ′ ,θ)p(xtest|hstarttest ′ )p(hstarttest
′|θ) p(Sn|hstarttest ,θ)p(xtest|hstarttest )
(59)
→Θ(1) p(xtest|θ) p(xtest|hstarttest ) . (60)
Plugging this into Equation 57 (after taking the reciprocal) and noting that p(hstarttest |θ)→πθ(hstarttest ) show that
p(hstarttest |Sn,xtest,θ)→Θ(1) p(xtest|hstarttest )πθ(hstarttest )
p(xtest|θ) (61)
=Θ(1)πθ(hstarttest |xtest). (62)
Note that we do not take the limit of the xtest terms since xtest 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
p(Sn,|hstarttest ′ ,θ) p(Sn,|hstarttest ,θ) →Θ(1). (63)
Redefining the ratio, we have
p(Sn,|hstarttest ′ ,θ) p(Sn,|hstarttest ,θ) =exp ( n· 1 n log p(Sn |hstarttest ′ ,θ) p(Sn |hstarttest ,θ) ) . (64)
As in Lemma 2, letOexi =[o delim i−1 ,Oi] be the i-th observation segment and the previous delimiter together for i>1 and defineOex1 =O1. Define θ(t) to be the bound on the TV distance between the distribution onhstarttest vs. π θ, for θ(t)-mixing time t. Using Lemma 2, we have the following upper and lower bounds:
1 n log
p(Sn |hstarttest ′ ,θ) p(Sn |hstarttest ,θ) (65)
≤ 1 n (n−t∑ i=1 [ log Cθ2 ( θ(t))2 Cθ1 ( θ(t))2 +log ∑ hdelimi−1∈D p(Oi|hdelimi−1 ,θ)p(hdelimi−1 |Oex1:i−1,θ)∑ hdelimi−1∈D p(Oi|hdelimi−1 ,θ)p(hdelimi−1 |Oex1:i−1,θ) ]
+log p(hdelimn |Oex1:n,hstarttest
′ ,θ)
p(hdelimn |Oex1:n,hstarttest ,θ) + n∑ i=n−t log
∑ hdelimi−1∈D p(Oi|hdelimi−1 ,hstarttest ′ ,θ)p(hdelimi−1 |Oex1:i−1,hstarttest
′ ,θ)∑
hdelimi−1∈D p(Oi|hdelimi−1 ,hstarttest ,θ)p(hdelimi−1 |Oex1:i−1,hstarttest ,θ)
)
≤ 1 n (n−t∑ i=1 log Cθ2 ( θ(t))2 Cθ1 ( θ(t))2 +(log(c1)−log(c2))+ n∑ i=n−t log ∑ hdelimi−1∈D 1·p(hdelimi−1 |Oex1:i−1,hstarttest ′ ,θ)∑ hdelimi−1∈D c5p(h delim i−1 |Oex1:i−1,hstarttest ,θ) )
≤ 1 n (n−t∑ i=1 log Cθ2 ( θ(t))2 Cθ1 ( θ(t))2 +(log(c1)−log(c2))− n∑ i=n−t log(c5) ) ≤ 1 n ( (log(c1)−log(c2))+O(t) ) (66)
=O(1/n) (67)
We used both Propositions 1 and 2 in the terms involving c1,c2 and c5. Note that in the first line, the sum can must be over the set of delimiter states D by using the assumption that observing a
delimiter odelim implies that the corresponding hidden state hdelim must be in D. where we note that t, Cθ2 ( θ(t)), Cθ1 ( θ(t)) are constants. The first of the two log ratios involving a marginalization over hdelimi−1 is 0 since the numerator and denominator are the same. Flipping the order of Cθ2 ( θ(t)),Cθ1 ( θ(t)), c1,c2 and 1 vs. c5 gives a lower bound:
1 n log
p(Sn,|hstarttest ′ ,θ)
p(Sn,|hstarttest ,θ) ≥ 1 n (n−t∑ i=1 log Cθ1 ( θ(t))2 Cθ2 ( θ(t))2 +(log(c2)−log(c1))+ n∑ i=n−t log(c5) ) ≥ 1 n ( (log(c2)−log(c1))+Ω(t) ) (68)
=Ω(−1/n) (69) Together, this shows that
p(Sn |hstarttest ′ ,θ) p(Sn |hstarttest ,θ) =exp ( n· 1 n log p(Sn |hstarttest ′ ,θ) p(Sn |hstarttest ,θ) ) =Θ(1) (70)
as desired.

heading: C PROOF OF THEOREM 1
text: Proof. We analyze the most likely prediction over the pretraining distribution conditioned on the prompt argmaxyp(y|Sn,xtest).
p(y|Sn,xtest)= ∫ θ p(y,|Sn,xtest,θ)p(θ|Sn,xtest)dθ (71)
∝ ∫ θ p(y,|Sn,xtest,θ)p(Sn,xtest|θ)p(θ)dθ (72)
∝ ∫ θ p(y,|Sn,xtest,θ) p(Sn,xtest|θ) p(Sn,xtest|θ∗) p(θ)dθ (73)
= ∫ θ ∑ hstarttest ∈H p(y,|Sn,xtest,hstarttest ,θ)p(hstarttest |Sn,xtest,θ) p(Sn,xtest|θ) p(Sn,xtest|θ∗) p(θ)dθ (74)
Defining the following quantity,
rn(θ)= 1
n log p(Sn,xtest|θ) p(Sn,xtest|θ∗) . (75)
we will show that under distinguishability for all θ 6=θ∗, rn(θ) converges to a negative constant such that p(Sn,xtest|θ) p(Sn,xtest|θ∗) =exp(n·rn(θ))=0 (76) for θ 6=θ∗, whereas this ratio is always 1 for θ=θ∗. This will then “select” the desired prompt concept through marginalization.
Taking the limit of Equation 74, we have
lim n→∞ ∫ θ ∑ hstarttest ∈H p(y,|Sn,xtest,hstarttest ,θ)p(hstarttest |Sn,xtest,θ) p(Sn,xtest|θ) p(Sn,xtest|θ∗) p(θ)dθ (77)
= lim n→∞ ∫ θ ∑ hstarttest ∈H p(y,|Sn,xtest,hstarttest ,θ)p(hstarttest |Sn,xtest,θ)exp(n·rn(θ))p(θ)dθ (78)
= ∫ θ ∑ hstarttest ∈H lim n→∞ p(y,|Sn,xtest,hstarttest ,θ)p(hstarttest |Sn,xtest,θ)exp(n·rn(θ))p(θ)dθ (79)
= ∑
hstarttest ∈H
p(y,|Sn,xtest,hstarttest ,θ∗)πθ ∗ (hstarttest |xtest)p(θ∗) (80)
∝ ∑
hstarttest ∈H
p(y,|Sn,xtest,hstarttest ,θ∗)πθ ∗ (hstarttest |xtest) (81)
=Ehstarttest ∼πθ∗ (hstarttest |xtest) [ p(y|xtest,hstarttest ,θ∗) ] (82)
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 rn(θ) converges to a negative constant for θ 6=θ∗. LetOexi =[odelimi−1 ,Oi] be the i-th observation segment and the previous delimiter together for i > 1 and define Oex1 =O1. Expanding the numerator of the ratio in rn(θ), we have
p(Sn,xtest|θ)=p(xtest|Sn,θ)p(Sn|θ) (83)
= ∑ hstarttest p(xtest|hstarttest ,θ)p(hstarttest |Sn,θ)p(odelimn |Oex1:n,θ) n∏ i=1 p(Oexi |Oex1:i−1,θ) (84)
= ∑ hstarttest p(xtest|hstarttest ,θ)p(hstarttest |Sn,θ) (85)
∑ hdelimn ∈D p(odelimn |hdelimn )p(hdelimn |Oex1:n,θ) n∏ i=1 ∑ hdelimi−1∈D p(Oi|hdelimi−1 ,θ)p(hdelimi−1 |Oex1:i−1,θ)
(86) = ∑ hstarttest p(xtest|hstarttest ,θ)p(hstarttest |Sn,θ) (87)
∑ hdelimn ∈D p(hdelimn |Oex1:n,θ) n∏ i=1 ∑ hdelimi−1∈D p(Oi|hdelimi−1 ,θ)p(hdelimi−1 |Oex1:i−1,θ) (88)
Note that in the last line, the sum must be over the set of delimiter statesD by using the assumption that observing a delimiter odelim implies that the corresponding hidden state hdelim must be inD. We restrict our attention to θ where p(Sn,xtest|θ)>0, since otherwise θ does not affect the prediction. Expanding rn(θ), we have the following upper bound:
rn(θ)= 1
n
( log
p(Sn,xtest|θ) p(Sn,xtest|θ∗)
) (89)
= 1
n
( log ∑ hstarttest
p(xtest|hstarttest ,θ)p(hstarttest |Sn,θ)∑ hstarttest p(xtest|hstarttest ,θ∗)p(hstarttest |Sn,θ∗) +log
∑ hdelimn ∈D
p(hdelimn |Oex1:n,θ)∑ hdelimn ∈D p(hdelimn |Oex1:n,θ∗) (90)
+ n∑ i=1 log
∑ hdelimi−1∈D
p(Oi|hdelimi−1 ,θ)p(hdelimi−1 |Oex1:i−1,θ)∑ hdelimi−1∈D p(Oi|hdelimi−1 ,θ∗)p(hdelimi−1 |Oex1:i−1,θ∗)
) (91)
≤ 1 n
( log ∑ hstarttest
1·p(hstarttest |Sn,θ)∑ hstarttest c5 ·p(hstarttest |Sn,θ∗) +n(log(c1)−log(c2))+ n∑ i=1 log
∑ hdelimi−1∈D
p(Oi|hdelimi−1 ,θ)∑ hdelimi−1∈D p(Oi|hdelimi−1 ,θ∗) ) (92)
= 1
n
( −log(c5)+n(log(c1)−log(c2))+ n∑ i=1 log ∑ hdelimi−1∈D p(Oi|hdelimi−1 ,θ)∑ hdelimi−1∈D p(Oi|hdelimi−1 ,θ∗) ) (93)
= 1
n
( −log(c5)+n(log(c1)−log(c2))+ n∑ i=1 log ∑ hdelimi−1∈D p(hdelimi−1 |θ)p(Oi|hdelimi−1 ,θ)∑ hdelimi−1∈D p(hdelimi−1 |θ∗)p(Oi|hdelimi−1 ,θ∗) ) (94)
= 1
n
( −log(c5)+n(log(c1)−log(c2))+ n∑ i=1 log p(Oi|θ) p(Oi|θ∗) ) (95)
→n→∞EO∼pprompt [ log
p(O|θ) p(O|θ∗)
] +log(c1)−log(c2) (96)
We used both Propositions 1 and 2 in the terms involving c1,c2 (bounding the probability of hdelim hidden states) and c5 (bounding the probability of xtest). 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 odelim implies that the corresponding hidden state hdelim must be in D. In the second to last line, we used the assumption that p(hdelim|θ)= 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 first token is different due to the starting distribution πprompt).
Aiming to decompose the expectation into a sum over the k tokens, we write pjθ(o)=p(O[j]=o|O[1 : j−1],θ). Then we have that
lim n→∞ rn(θ)<− k∑ j=2 EO[1:j−1]∼pprompt [KL(p j θ∗‖p j θ)]+EO[1]∼pprompt [ log p(O[1]|θ) p(O[1]|θ∗) ] +(log(c1)−log(c2))
(97)
The second term is an error term that depends on how different the starting prompt distribution πprompt (which is part of pprompt) 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 θ 6=θ∗ the Condition 1 holds, and thus
lim n→∞ p(Sn,xtest|θ) p(Sn,xtest|θ∗) = lim n→∞ exp(n·rn(θ))=0 (98)
since rn(θ) has a negative, constant limit. Note that exp(n·rn(θ∗))=1 for θ∗.

heading: D NON-DISTINGUISHABLE CASE
text: When Condition 1 is unsatisfied, Equation 13), 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 Ávila Pires & Szepesvári (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(p(ytest|xtest,θ∗)‖p(ytest|xtest,θ) is bounded above for all θ. If
Extest∼pprompt [KL(p(ytest|xtest,θ∗)‖p(ytest|xtest,θ))]≤ θ for all θ∈B, (99)
then
lim n→∞ L0-1(fn)≤ inf f L0-1(f)+g
−1 (
sup θ∈B θ
) (100)
where
g(δ)= 1
2 ((1−δ)log(1−δ)+(1+δ)log(1+δ)) (101)
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:
lim n→∞ L0-1(fn)= lim n→∞ Extest,ytest∼pprompt [1[fn(xtest) 6=ytest]] (102)
=Extest,ytest∼pprompt [ lim n→∞ 1[fn(xtest) 6=ytest]] (dominated convergence, boundedness of indicator) (103)
=Extest,ytest∼pprompt [1[ lim n→∞ fn(xtest) 6=ytest]] (104)
where in the last step we use that since the output space of fn is discrete and the probabilities that the in-context predictor takes an argmax over converges, then for N large enough, fN (xtest)=limn→∞fn(xtest).
Note that for every input xtest, the limiting in-context learning predictor outputs the argmax of a predictive distribution which can be a mixture of predictive distributions over B:
lim n→∞ fn(xtest)=argmax y
Eθ∼q[p(y|xtest,θ)] (105)
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:
Extest∼pprompt [KL(p(y|xtest,θ∗)‖Eθ∼q[p(y|xtest,θ)]] (106) ≤Extest∼pprompt [Eθ∼q[KL(p(y|xtest,θ∗)‖p(y|xtest,θ))]] (107) =Eθ∼q[Extest∼pprompt [KL(p(y|xtest,θ∗)‖p(y|xtest,θ))]] (108) ≤sup θ∈B Extest∼pprompt [KL(p(y|xtest,θ∗)‖p(y|xtest,θ))] (109)
where we can exchange the order of expectations since the KL is bounded (dominated convergence).
From the KL boundKL(p(ytest|xtest,θ∗)‖p(ytest|xtest,θ), we thus have
Extest∼pprompt [KL(p(ytest|xtest,θ∗)‖p(ytest|xtest,θ))]=LCE(θ)−LCE(θ∗)≤sup θ∈B θ (110)
where LCE(θ) = −Extest∼pprompt [p(ytest|xtest,θ∗)logp(ytest|xtest,θ)] is the multiclass logistic risk, and LCE(θ
∗) is the entropy of the predictive distribution. Applying Theorem 2.2 and 5.11 of Ávila Pires & Szepesvári (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 for smaller than 0.7, where the bound is non-vacuous.

heading: D.1 PROOF OF THEOREM 2
text: Proof. By the continuity assumption, we have for any θ in B that k∑ j=2 KLj(θ ∗‖θ)≥ 1 2 k∑ j=2 (θ−θ∗)>Ij,θ∗(θ−θ∗)+(k−1)O(‖θ−θ∗‖3) (111)
≥ 1 2 (k−1)λmin(Ij,θ∗)‖θ−θ∗‖2 (112)
=⇒‖θ−θ∗‖2≤ θ start+ θ delim
1 2 (k−1)(minj λmin(Ij,θ∗))
. (113)
To bound the last KL term,
KLk(θ ∗‖θ)= 1
2 k∑ j=2 (θ−θ∗)>Ij,θ∗(θ−θ∗)+O(‖θ−θ∗‖3) (114)
≤ 1 2 (max j λmax(Ij,θ∗))‖θ−θ∗‖2+O(‖θ−θ∗‖2) (115)
≤ ( θ start+ θ delim)(maxj λmax(Ij,θ∗)+O(1))
(k−1)minj λmin(Ij,θ∗) . (116)
Rearranging and noting thatKLk(θ∗‖θ)=Extest∼pprompt [KL(p(ytest|xtest,θ∗)‖p(ytest|xtest,θ))], we have
Extest∼pprompt [KL(p(ytest|xtest,θ∗)‖p(ytest|xtest,θ))]≤ ( θstart+ θ delim)(maxj λmax(Ij,θ∗)+O(1))
(k−1)minj λmin(Ij,θ∗) (117)
Plugging into Lemma 4 gives the result.

heading: D.2 PROOF OF THEOREM 3
text: 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 fitting the intermediate tokens. The proof is immediate given the KL bound and Lemma 4, given that the length of xtest 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
Extest∼ppromptKL(p(ytest|xtest,θ∗‖p(ytest|xtest,θ)) (118)
≤ 1 k−1 k∑ j=2 EO[1:j−1]∼ppromptKL(p(O[j]|O[1 :j−1],θ ∗‖p(O[j]|O[1 :j−1],θ) (119)
≤ supθ( θ start+ θ delim)
k−1 (120)
by Theorem 1 and Condition 1. Plugging this into Lemma 4 gives the result.

heading: E EXPERIMENTAL DETAILS

heading: E.1 GINC DATASET
text: 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 first column of the memory matrix is fixed 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
softmax((u−0.5)/t) (121)
where u∈R100 has uniform random entries in [0,1] and t is a temperature parameter, set to 0.1.
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 final entity transition matrix is then 0.1T+0.9I 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 uwith 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
argmax y
p(y|xtest,θ∗) (122)
under the prompt concept θ∗. This differs from the theory, which samples ytest instead of taking it to be the most likely token. However, there can be a large amount of intrinsic error that sampling introduces. We define 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 fixed to be 5 (for example). The first 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 first example), we finish generating one example. To end the example, we append a delimiter (backslash). We repeat this example 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|xtest,θ∗). 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 first 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.

heading: E.2 TRANSFORMER DETAILS
text: 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).

heading: E.3 LSTM DETAILS
text: 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.
E.4 VARYING 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.
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.

heading: E.5 SENSITIVITY TO EXAMPLE ORDERING
text: 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. (2021). 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. We use a 4-layer Transformer and GINC with vocab size 50, and consider test examples generated from the concept and prompt start distribution that matches the training set. Similarly to the behavior of GPT-3 (Zhao et al., 2021), there is a significant variation between example orderings, with a 10–40% gap between different permutations of the same set of examples.

heading: E.6 ZERO-SHOT IS SOMETIMES BETTER THAN FEW-SHOT
text: On some hyperparameter settings of GINC, we find that zero-shot performance can be better than few-shot performance in our experiments. This finding mirrors the behavior of GPT-3 on some
datasets (e.g., LAMBADA, HellaSwag, PhysicalQA, RACE-m, CoQA/SAT analogies for smaller 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.

heading: E.7 EXPERIMENT ON GPT-3
text: 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 define 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 define 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.

Review summaries:
The submission studies the phenomenon of the "in-context learning" behavior of large language models (LLMs). Two results are presented: First, the submission argues that in-context learning can be formalized as Bayesian inference in a mixture of Hidden Markov Models (HMMs). Second, the submission introduces small-scale settings in which the in-context learning behavior of LLMs is reproduced empirically.

The paper attempts to bring clarification to In-Context Learning (ICL) by relating it to Bayesian Learning over mixtures of HMMs. Each component of the mixture generates examples of facts of the same conceptual type, such as "Albert Einstein was German", "Marie Curie was Polish", etc (here the conceptual type is "nationality"). The first part of the paper addresses a theoretical question, namely whether the mixture model, on being given a prompt, that is, a prefix such as "Albert Einstein was German. Marie Curie was Polish. Isaac Newton was ", is able to implicitly identify the relevant component of the mixture (i.e. "nationality"), and consequently, whether it is able to correctly complete the prompt, here with "British". The second part of the paper consists in simulations, where the mixture model is used to produce data on which a neural language model (LSTM or Transformer) is trained, and on evaluations of how well the neural LM is able to perform the completions, under different configurations and ablations.

This paper studies the question of why few-shot prompting works for large language models (LM) such as GPT? The paper is mainly focused on understanding why prompting works given that concatenating multiple examples separated by a delimiter introduces out-of-distribution prompts that the LM would most likely not have seen during its pre-training. The paper shows that when the pre-training distribution is a mixture of HMMs, prompting works as a result of Bayesian inference. The authors show the validity of their theoretical results by training transformers and LSTM on synthetic datasets sampled from some HMMs that they devised.

The authors propose to study the phenomenon of "in-context learning" through the lens of Bayesian prediction. They introduce a framework in which the data generating distribution is given as a mixture of HMMs parameterized by a latent "concept," and they show that under suitable conditions the posterior predictive distribution over a completion given a structured prompt asymptotically selects the "concept" behind this particular structure. They illustrate this framework and reproduce the phenomenon in a simplified language modeling example.