An Explanation of In-context Learning as Implicit Bayesian Inference

1 INTRODUCTION . 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 . 2 IN-CONTEXT LEARNING SETTING . 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 ) 2.1 ASSUMPTIONS . 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|θ ) .

Summary: 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.