Title: Large Language Models Are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Large language models

Bayesian explanation

in-context learning
Abstract: In recent years, pre-trained large language models (LLMs) have demonstrated remarkable efficiency in achieving an inference-time few-shot learning capability known as in-context learning. However, existing literature has highlighted the sensitivity of this capability to the selection of few-shot demonstrations. Current understandings of the underlying mechanisms by which this capability arises from regular language model pretraining objectives remain disconnected from the real-world LLMs. This study aims to examine the in-context learning phenomenon through a Bayesian lens, viewing real-world LLMs as latent variable models. On this premise, we propose an algorithm to select optimal demonstrations from a set of annotated data with a small LM, and then directly generalize the selected demonstrations to larger LMs. We demonstrate significant improvement over baselines, averaged over eight GPT models on eight real-world text classification datasets. We also demonstrate the real-world usefulness of our algorithm on GSM8K, a math word problem dataset. Our empirical findings support our hypothesis that LLMs implicitly infer a latent variable containing task information.

Large Language Models Are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning

Abstract

In recent years, pre-trained large language models (LLMs) have demonstrated remarkable efficiency in achieving an inference-time few-shot learning capability known as in-context learning. However, existing literature has highlighted the sensitivity of this capability to the selection of few-shot demonstrations. Current understandings of the underlying mechanisms by which this capability arises from regular language model pretraining objectives remain disconnected from the realworld LLMs. This study aims to examine the in-context learning phenomenon through a Bayesian lens, viewing real-world LLMs as latent variable models. On this premise, we propose an algorithm to select optimal demonstrations from a set of annotated data with a small LM, and then directly generalize the selected demonstrations to larger LMs. We demonstrate significant improvement over baselines, averaged over eight GPT models on eight real-world text classification datasets. We also demonstrate the real-world usefulness of our algorithm on GSM8K, a math word problem dataset. Our empirical findings support our hypothesis that LLMs implicitly infer a latent variable containing task information. 1

1 Introduction

Transformer-based [41] pre-trained large language models (LLMs) have demonstrated significant advancements in a variety of natural language processing (NLP) tasks. As the size of these LLMs increases, they gain 'in-context learning' capabilities, whereby the models achieve state-of-the-art (SOTA) or near-SOTA performance by conditioning on a small number of demonstration examples at inference time, without any need for updating model parameters [4]. Below is an example input sequence for semantic analysis with in-context learning:

Great movie. Positive.\n The worst movie ever. Negative.\n Can't wait to see the second movie!

The first two lines are two demonstrations, and the third line is a test input. We expect an LLM to output the correct label Positive as a continuation.

In-context learning has been demonstrated to be an effective technique for a wide range of NLP tasks. However, it is sensitive to the choice, format, and even the order of the demonstrations used [29, 20]. This makes achieving optimal performance with in-context learning a significant challenge, requiring real human effort to adjust the format and selection of demonstration examples. Heuristic solutions, such as selecting demonstrations based on the similarity between the demonstrations and test input

[19, 37] have been proposed, but a comprehensive understanding of why certain demonstrations are effective while others are not remains elusive. Additionally, the mechanisms by which LLMs acquire in-context learning capabilities through training on natural text under the standard language model pre-training objective are not fully understood. Recent works on understanding in-context learning provide valuable insights and theoretical results [5, 1, 42, 14, 12], but are limited in scope, focusing on synthetic experiments to validate their hypotheses, while it remains unclear if these results generalize to LLMs pre-trained on real-world natural language data. Xie et al. [50] introduced a prominent result providing a latent topic (concept) variable interpretation for in-context learning. They showed that the in-context learning predictor approaches the Bayes optimal predictor when the number of demonstrations approaches infinity, under the assumption that both the pre-training data distribution and task-specific data distribution are Hidden Markov Models (HMM). However, the assumption that the data generation process is Hidden Markovian makes extrapolation of the result to natural language questionable, and restricts empirical verification to synthetic data with toy models.

We are inspired by this prior work and introduce a more general and natural explanation built on realistic assumptions, which gives rise to a practical demonstration selection algorithm. Our explanation is inspired by the generation process of a topic model, i.e. a simple latent variable model:

Where θ ∈ Θ represents a potentially high dimensional topic/concept variable, Θ is the space of the topic/concept variable, and w 1: T refers to the token sequence of a piece of text. Note that the topic model here refers to the modern neural topic models [23, 22]. On the other hand, generative LLMs model text data according to the general probabilistic decomposition:

While in practice, LLMs generate new tokens based on all previous tokens, we investigate whether a simplified assumption similar to that of topic models can be made for LLMs:

In this scenario, the generated tokens are assumed to be conditionally independent of previous tokens, given the latent topic (concept) variable that acts like an approximate sufficient statistic for the posterior information related to the prompt w 1: t . For in-context learning, this concept variable includes format and task information. By conditioning on an appropriate latent concept variable, LLMs would generate the desired continuation with P ( w t +1: T | θ ) . As LLMs do not explicitly learn a latent variable distribution like LDA-style topic models [3], we can instead utilize this formulation under an Empirical Bayesian formulation inspired by Lester et al. [17] to only approximate the optimal latent variable value for a desired task, using a small LLM (with less than 1B parameters), which is computationally efficient.

We empirically validate our explanation by selecting examples ( w 1: t in the equations) that are most likely to infer the optimal latent variable value (those with the highest posterior probability P ( θ | w t +1: T ) ). We then directly use them as demonstrations for in-context learning with other larger LLMs (up to 175B parameters) and observed a significant performance improvement. The generalization of demonstrations between LLMs is likely a result of similar pre-training data distributions.

While our work is inspired by that of Xie et al. [50], our approach differs significantly in both theoretical analysis and experimental settings. Our main contributions are as follows:

We assume a general data generation process specified by a three-variable causal graph, without constraints on the distribution function or the number of demonstrations.

We prove under these realistic assumptions that the in-context learning predictor can reach the Bayes optimal predictor with a finite number of demonstrations chosen using the latent concept variable.

We introduce an efficient, practical demonstration selection algorithm based on our theoretical results, which can select demonstrations using a small LLM and then directly generalize the demonstrations to other LLMs. The effectiveness of our algorithm is empirically validated using real-world LLMs on both text classification tasks and math word problems.

Our goal is to close the gap between theoretical understandings and real-world LLMs. To the best of our knowledge, our proposed latent variable explanation of in-context learning is the first Bayesian explanation that yields an effective algorithm in real-world scenarios.

2 Theoretical Analysis

In in-context learning, the prompt w 1: t is composed of several demonstrations and a test input. The generated tokens w t +1: T represent the model's prediction for the test input.

2.1 Notations and Problem Setting

Suppose the objective of our task is to predict a discrete target variable Y ∈ Y , given a token sequence X ∈ X , where X is the space of all possible token sequences. θ ∈ Θ is a potentially high dimensional latent variable, where Θ is the high dimensional space of the variable. Unlike the traditional topic model, θ is not assumed to be discrete, but continuously distributed over Θ . To define the data generation process, we posit the existence of an underlying causal relation between X , Y , and θ . We examine two potential directions of this causal relation, namely X GLYPH<1> Y GLYPH<0> θ and Y GLYPH<1> X GLYPH<0> θ , which can be represented mathematically as the following structural equations:

Here ϵ ∈ E is an independent noise variable, f : X × Θ ×E → Y and g : Y × Θ ×E → X are two deterministic functions. Furthermore, we denote the joint data distribution by X,Y, θ ∼ P , and assume that Y is sampled from a uniform distribution over Y . The distinction between these two directions is crucial, as it allows us to utilize the direction in which the child variable ( Y or X ) is independent of the other variables, given its parents.

We hypothesize that the causal direction depends on the nature of the task. For instance, in the task of predicting the sentiment ( Y ) of a movie review ( X ), it is reasonable to assume that the opinion about the movie is formed before writing the review, thus making Y the cause of X , along with the task concept of 'writing a passage to express one's opinion about the movie" ( θ ). Conversely, for the task of classifying whether a product review ( X ) is helpful to other customers ( Y ), it is the quality of the review ( X ) that cause other customers to upvote it ( Y ), along with the task concept of 'rating the helpfulness of this review" ( θ ). In the rest of the paper, we will focus on the X GLYPH<1> Y GLYPH<0> θ direction and leave a detailed discussion of the other direction in the Appendix.

Suppose we are interested in a task (e.g. semantic analysis) denoted by d ∈ T , where T is the space of all possible tasks. We assume there is an injective function between T and Θ . i.e. for each task d , there is a concept variable θ d , such that each data ( X d , Y d ) sampled from task d is generated by:

To perform in-context learning with an LLM (generically denoted by model label M ), we condition on a fixed set of k demonstration examples ( X d 1 , Y d 1 ) , ( X d 2 , Y d 2 ) , ..., ( X d k , Y d k ) sampled from task d .

Following previous works [24, 26], as we are not using any instruction fine-tuned models, we do not include a task description in the prompt, with the aim of focusing on the examination of the demonstrations. To naturally project Y into the token space X , we define injective mappings τ d : Y → X , which are typically defined by human understanding of the task d . e.g. for sentiment analysis, τ d map positive class to the token 'positive" and negative class to the token 'negative". Additionally, a delimiter token w d is defined, typically an empty space or a new line token, to separate the demonstrations when concatenated. We denote the LLM output probability of X , Y , and θ , with the aforementioned preprocessing applied, by P d M :

2.2 Problem Analysis and Theoretical Results

Suppose a set of observed data sampled from task d , denoted as D d , is available, allowing for the selection of the k most suitable demonstrations from it. For any incoming test example X , we have:

LatentConceptLearning

DemonstrationSelection

工

In-contextLearning

TestY

LLM

LLM

LLM

X

compute

X

Y

compute

(X,Y),(X2,Y2),,(XkYk)

Test X

Update

Cross entropy loss

Dataset

Select (X,Y) pairs

Language model probability

infer

embeddings

log PM(Y|0,x)

PM(0|X,Y)

Here, we assume the sampling of the test example is independent of the sampling of the demonstrations, so Y is independent of the demonstrations given θ and X . We also assume that the pre-trained data distribution P d M is a suitable approximation of the assumed data distribution P :

Assumption 2.1. Assume that P M ( X ) = P ( X ) , and P d M ( Y | θ , X ) ∝ P ( Y | θ , X ) for X GLYPH<1> Y GLYPH<0> θ .

Note that the assumption that a large language model captures the true distribution of language is fairly common in the literature studying LLMs [50, 34, 47]. With this assumption, we establish:

Proposition 2.2. If task d follows the X GLYPH<1> Y GLYPH<0> θ direction, then arg max y ∈Y P d M ( Y = y | θ d , X ) is the Bayes optimal classifier.

In this case, only when P d M ( θ | X d 1 , Y d 1 , ..., X d k , Y d k , X ) completely concentrate on θ d , can the incontext learning classifier become the Bayes optimal classifier [11]:

Theorem 2.3. If task d follows the X Y θ direction, then the in-context learning classifier

always has a higher or equal probability of misclassification to the Bayes optimal classifier arg max y ∈Y P d M ( Y = y | θ d , X ) . Equality only holds when

A similar argument can be made for the Y X θ direction. 2 Here, Equation (1) would become:

Note that the left-hand side of Equation (1) and Equation (2) are similar to the direct and channel method introduced by Min et al. [24]. However, our analysis differs from theirs in that we do not treat ( Y GLYPH<1> X GLYPH<0> θ ) as the universally superior channel direction for modeling in-context learning, rather arguing that depending on the end task, the causal direction ( X GLYPH<1> Y GLYPH<0> θ ) is sometimes better. This view is supported by our empirical results in Appendix B.

3 Method

Here we demonstrate how the proposed theory can be practically applied to select optimal demonstration examples. Since latent variable θ encodes both the task and format information, the whole distribution over Θ is too complex to model. Unlike traditional topic models, we will only focus on estimating an optimal value θ d corresponding to task d .

First, we perform latent concept learning , wherein the task latent θ d is learned as a set of new token embeddings using prompt tuning over the full demonstration candidate set. With this optimal task latent, we then perform demonstration selection , where a smaller set of demonstrations is chosen to maximize the likelihood of postfixing the latent concept tokens. We only need to use a small LLM to do the above steps to obtain an optimal set of demonstrations that can be directly transferred to other LLMs. Figure 1 is an overall illustration of our proposed method.

Algorithm 1 Latent concept learning

Input: Dataset D = { ( x i , y i , d i ) } i associated with a set of tasks S , LLM M , number of concept tokens per task c , learning rate α , and number of training steps N .

Output: LLM M ′ with fine-tuned concept tokens.

Add c |S| new tokens to the vocabulary. i.e. The concept tokens ˆ θ d for each task in S . Randomly initialize their embeddings E new . Freeze all parameters in M except E new ;

Sample a random batch B in D and initialize gradient g ← 0 ;

for each data point ( x, y, d ) in B do

end for

E

new

=

E

new

-

αg

;

end for

3.1 Latent Concept Learning

We want to first find the optimal value of the latent concept variable θ d corresponding to a task d ∈ T . As arg max y ∈Y P d M ( Y = y | θ d , X ) is the Bayes optimal classifier according to Proposition 2.2, θ d should be able to minimize -E X,Y,d [log P d M ( Y | θ d , X )] for the X GLYPH<1> Y GLYPH<0> θ direction. In practice, we try to align θ d to the token embedding space by adding new tokens to the vocabulary. After this alignment, we hope to be able to use the learned new tokens of θ d as regular tokens.

More specifically, building upon the methodology proposed by Lester et al. [17], for each specific task d , c new concept tokens (denoted as ˆ θ d ) are added to the original vocabulary of LLM M to represent the corresponding task concept θ d . Subsequently, the embedding of these new tokens E new ( ˆ θ d ) is fine-tuned while freezing the remaining parameters of LLM M . The variable c is treated as a hyperparameter. In practice, in order to condition on θ d , the corresponding c concept tokens are appended to the input X (or Y ) as shown in the example provided below, where c = 2 :

<sentiment_token_1><sentiment_token_2> Can't wait to see the second movie!

By giving the above input tokens, we ask the LLM to predict the correct label Positive for us. Note that <sentiment_token_1> here is just a label assigned to the newly added concept token. It can be anything as long as it does not overlap with the original vocabulary of LLM.

The fine-tuning objective would then be minimizing L ( ˆ θ d ) = E X,Y [ ℓ ( X,Y ; ˆ θ d )] , where

Theoretically, if we can minimize the above loss function, a Bayes optimal classifier can be obtained, and the concept tokens would be a reasonable delegate of the real latent concept variable:

Proposition 3.1. When L ( ˆ θ d ) is minimized, P d M ( Y | ˆ θ d , X ) = P ( Y | θ d , X ) for X GLYPH<1> Y GLYPH<0> θ . If the LLM M is invertible, then ˆ θ d = θ d . 3

We denote the LLM M with fine-tuned concept tokens by M ′ . Since we add the concept tokens into the regular token vocabulary, the raw LLM output probability P M ′ ( ˆ θ d | w 1: t ) ( w 1: t denote a given prompt) would be in the token sequence space X instead of the concept space Θ . Since learning all possible θ d ∈ Θ is infeasible, we propose to approximate the concept space Θ by sampling a diverse subset of tasks S ⊆ T . Then the estimated conditional probability of θ d would be:

To obtain the concept tokens for all tasks in S , we fine-tune all tasks together with the loss ∑ d ∈S L ( θ d ) . We summarize the proposed algorithm in Algorithm 1.

Algorithm 2 Demonstration selection

Input: dataset D d for a task d . LLM with fine-tuned concept tokens M ′ . The number of demonstrations k .

Output: A set of selected demonstrations.

Compute ˆ P d M ( ˆ θ d | X d , Y d ) ;

end for

Select top k examples with the largest ˆ P d M ( ˆ θ d | X d , Y d ) , denoted as ( X d 1 , Y d 1 ) , ..., ( X d k , Y d k ) ;

Note that the embedding matrix of a generative LLM is shared on both the input and output sides. So while we only see the concept tokens on the input side at the training time, they can be viewed as regular word tokens that can be generated on the output side.

3.2 Demonstration Selection

According to Theorem 2.3, for a task d , to make the in-context learning classifier closer to the Bayes optimal classifier, we need to select demonstrations ( X d 1 , Y d 1 ) , ..., ( X d k , Y d k ) that maximize P d M ( θ d | X d 1 , Y d 1 , ..., X d k , Y d k , X ) for all X ∈ X . Then our goal then becomes selecting demonstrations that can best infer the task concept for all test inputs on average:

As test examples are sampled independent of the demonstrations, and P M ( X ) = P ( X ) according to Assumption 2.1, we have

If we assume each demonstration is also sampled independently, we have:

Assuming that θ has a uniform prior, then our goal becomes finding the top k demonstrations that maximize ˆ P d M ′ ( ˆ θ d | X d i , Y d i ) . Note that the independence between demonstrations is a simplified assumption to reduce the combinatory search space of ( X d 1 , Y d 1 ) , ..., ( X d k , Y d k ) . In practice, selected demonstrations are likely correlated as some demonstrations may work well together but not necessarily work well by themselves. However, it would be too expensive to search the O ( |D d | k ) combinations over the candidate set D d . In practice, this simplification works reasonably well. We leave this combinatory search problem to future research.

Also, as we are using an LLM to approximate the data distribution, the order of the demonstrations might matter. We will show in the Experiment section that the order does not matter, so no reordering of the selected demonstrations is needed. The full selection algorithm is shown in Algorithm 2.

4 Experiments

Datasets. We conduct experiments on eight datasets from five different types of NLP classification tasks: sentiment analysis, linguistic analysis, topic classification, emotion classification, and hate speech detection. For sentiment analysis, we choose the Stanford Sentiment Treebank (SST2) dataset [35] from the GLUE benchmark [43] and the financial phrase bank (FPB) dataset [21]. SST2 is constructed based on movie reviews labeled 'positive" or 'negative", and FPB is based on financial news labeled 'positive", 'negative", or 'neutral". For linguistic analysis, we choose the Corpus of Linguistic Acceptability (COLA) dataset [46] from the GLUE benchmark, based on sentences collected from linguistic books, labeled with 'acceptable" or 'unacceptable". For topic classification, we choose the DBpedia ontology classification dataset [52], based on DBpedia 2014 [16], labeled with 14 different ontology classes. For emotion classification, we choose the dataset from Chatterjee et al. [6] and Saravia et al. [33], both of which are collected from Twitter. Chatterjee et al. [6] (EmoC)

Uniform

Similar

Ours

69.2

66.9

66.8

66.6

61.2

60.3

62.6

59.7

64.8

60

65.1

61.3

64.2

62.4

62.3

62.4

63.7

56.8

56.2

57.4

56.7

56.8

57.7

53.7

GPT2

GPT2-large

GPT2-xl

GPT3-ada

GPT3-babbage

GPT3-curie

GPT3-davinci

(124M)

(355M)

(774M)

(1.5B)

(350M)

(1.3B)

(6.7B)

(175B)

predict emotion given a three-turn contextual dialogue, while Saravia et al. [33] predict emotion given a Twitter message with clear emotion. For hate speech detection, we choose the online hate speech detection dataset (ETHOS) [27], collected from online social media platforms. Here we detect two types of hate speech: sexual orientation (ETHOS-SO) and religion (ETHOS-R). While in Section 2, we assume that all tasks share the same label space Y , here we relax such assumption and allow a different number of labels for different tasks. We use minimal formatting to process each example. A detailed description of the datasets and our data processing procedure can be found in Appendix B.

Experiment settings. To determine the causal direction for each task, we select the direction that can give higher accuracy when using random demonstrations 4 . We adopt the Y → X ← θ direction for sentiment analysis, topic classification, and emotion classification tasks, which is consistent with the intuition that people usually have some sentiment, topic, or emotion in mind before writing a piece of text. We adopt the X → Y ← θ direction for the linguistic analysis and hate speech detection type of tasks. While this is less intuitive, we can understand this as linguistic error and hate speech detection are more of a post hoc task in contrast to the previous tasks.

Without specification, we use k = 4 number of demonstrations and c = 10 number of concept tokens per dataset for our experiments, as the context length of GPT2 is 1024, and a larger number of demonstrations may not be able to completely fit into it. We use GPT2-large to learn the concept tokens and then compute the probability of each candidate demonstration example. We select our demonstrations from a randomly selected 100 example subset of the train set as the candidate set D d . We use the same set of demonstrations selected by GPT2-large for all other LLMs. We test the performance of the selected demonstrations using at most 1000 examples randomly sampled from the test set. Each experiment is repeated for five runs with different random seeds (the randomness comes from the sampling of the candidate set and the sampling of the test set). We adopt a large portion of the code from Min et al. [25], which is based on Huggingface [49].

Baselines. We consider the following baselines:

Uniform : We uniformly select k demonstrations from D for each test example.

Similar : According to Liu et al. [19], demonstrations that are semantically similar to the test example would hare more performant. Following their method, we use a pre-trained sentence Transformer [31] to calculate the cosine similarity between the demonstrations and test examples. We choose the top k similar demonstrations from D for each test example.

Main results. 5 Figure 2 shows our main results averaged over all eight datasets, using the firstgeneration GPT2s and GPT3s, without any instruction fine-tuning [28] or Reinforcement Learning from Human Feedback (RLHF) [36]. Our method significantly outperforms baselines on eight different LLMs, with 12.5% relative improvement to the uniform selection baseline on average, which shows the effectiveness of our method. The demonstrations selected by our method are exclusively based on GPT2-large, while the same set of demonstrations can be generalized to all other GPTs.

Results with non-GPT models. In Figure 3a, we test the demonstrations selected by our method using GPT2-large on more LLMs (GPT3 [4], GPT3-instruct [28, 36], GPT-J [44], OPT [51], and LLaMA [38]) with similar sizes (6-7B), and show that the selected demonstrations improve in-context learning performance of all of them. The fact that GPT3-curie obtains the largest performance improvement is likely because similar pre-training data distributions help the generalization of the

(a) Proposed method v.s. randomly selected demonstrations

(b) Proposed method v.s.

using

randomly selected tokens

Uniform

Ours

69.2

62.3

62.7

52.8

56

58.8

54.7

56.6

36.4

37.6

GPT3-curie

GPT3-curie-insturct

GPT-J (6B)

OPT-6.7B

LLaMA-7B

(6.7B)

(6.7B)

Uniform

Randomtokens

Ours

64.8

61.2

57.457.2

53.7 53.7

GPT2 (124M)

GPT2-large (774M)

selected demonstrations. Different-size GPT2 models share the same pre-training corpus [30], while GPT3s are pre-trained on a dataset expanded from the GPT2 pre-training corpus [4]. Thus the pre-training distribution of GPT3-curie and GPT2-large can be assumed to be similar.

Results on GSM8K. Since our primary goal is to connect the theory with real-world models and datasets, we did not try to include harder tasks in the main results in Figure 2. In practice, our proposed method is most effective with hard tasks that even parameter-efficient fine-tuning with smaller models cannot outperform in-context learning with the same or larger models. To showcase the usefulness of our proposed algorithm, We added a new dataset, GSM8K [9], which is a math word problem-solving dataset with chain-of-thoughts solutions. Table 1 shows the test accuracy of the final numerical answer with greedy generation. We randomly select a test set of 200 examples instead of using the full test set for computation efficiency. 6

As shown in the first row of Table 1, prompt tuning with ten new tokens can only obtain less than 4% accuracy on the GSM8K test set. The last four rows show the in-context learning results with different size Llama 2 models [39] and ChatGPT. Our proposed demonstration selection method (last two columns) significantly outperformed the Uniform and Similar baseline. We also find that the demonstrations selected with a larger model (7B) are more effective than those selected with a smaller model (1.5B). The results show that our demonstration selection method is a good choice under a low data setting, with a small computing budget and minimal inference latency. Our proposed method can also potentially be combined with other prompting techniques [8] to boost performance further.

Learned tokens v.s. Random tokens. To confirm the critical role of the latent concept variable in the proposed demonstration selection algorithm, we compare the performance of using the learned concept tokens versus using randomly selected tokens from the original vocabulary in Figure 3b. The demonstrations selected by random tokens only obtain the same performance as randomly selected demonstrations, showing that the performance gain of our method comes from the learned concept tokens containing the task and format information, not other elements of our algorithm.

k ablation study. While we use k = 4 demonstrations for all experiments, we also test the effectiveness of our method using different k . As shown in Figure 4a, our method significantly outperforms the random selection baseline with k = 2 , 4, 8, and 16. To fit in large k s, we use GPT3-ada with a longer context length (2048). Note that for real-world tasks, it is in general not true that more demonstrations guarantee higher performance [7]. We can see that the uniform baseline performance increases from k = 2 to k = 8 , then drops a little at k = 16 . Our method improves the uniform baseline by around 5% absolute for all k s, while k = 4 improves the most (6.6%). Our method appears to have a diminishing effect when k becomes larger, which is likely because the effect of more demonstrations overwhelms the effect of demonstration choices.

c ablation study. While we use c = 10 number of concept tokens for all experiments, we also investigate the effect of different c on our method. When c is small ( c = 5 ), the concept tokens cannot effectively capture the task and format information, thus cannot im-

(a)

k

ablation study.

(b)

c

ablation study.

Uniform

Ours

64.2

62.3

63.4

63

62.4

60.92

59.8

56.8

58.8

57.4

56.8

56.6

54.7

k=2

k=4

k=16

Uniform

c=5

c=10

c=15

c=20

prove the performance. When c increases from 10 to 20, we observe a drop in the performance. It is likely because the selectivity of the concept tokens decreases when c increases. The longer the concept token sequence is, the more likely it will contain meaningless tokens that do not contribute to demonstration selection.

Effect of demonstrations' order. We find that the demonstrations selected by our method are insensitive to their order in most cases. 7 An exception is the EmoC dataset, where our method has a high variance. On the contrary, Lu et al. [20] found that the order of the demonstration matters, and a good ordering cannot be transferred between different LLMs. We suspect that the ordering only matters when the demonstration selection method is not robust. Since Lu et al. [20] randomly selects one set of demonstrations for the whole test set, the variance in performance is high with different demonstrations, thus ordering matters. And since such ordering is not transferable while our selected demonstrations are highly transferable, we suspect the core task information is stored in the content of the demonstrations, while the ordering mainly captures model-specific artifacts.

Qualitative analysis. In Figure 5, we provide a t-SNE [40] projection of the learned concept token embeddings. The tokens corresponding to semantically similar tasks are close together. Note that this result only aims to provide a straightforward illustration of concept tokens. The effect of concept tokens should be understood by the previous quantitative results. 8

We also list the top 4 selected demonstrations in Table 14 in Appendix B. Compared to the examples with lower scores, the selected examples for GSM8K have more deductive reasoning (i.e. with the connecting words 'so', 'then', 'thus', etc.), instead of listing parallel conditions. For SST2, the selected examples are longer and more complex, sometimes including a 'but'. This can be understood as these

Webtechnology

numbers

Negativefinance

Positivefinance

Negativeemotions

Positive

characteristics

9

SST2

FPB

3

COLA

DBpedia

3

ETHOS-SO

Environment

Medical

1

2

0T6

610

10

8

10

EmoC

EmoS

ETHOS-R

Adverbs

1

8

5

10

2

40

Academic

3

-4

2

0

2

4

Emotionallyunstable/immature

harder examples can represent the task more comprehensively. This conclusion also aligns with the findings in [13] that hard examples in the pre-training data contribute to in-context learning the most. The label distribution of the selected demonstrations is usually balanced in class, which reduces the possible biases introduced by the demonstrations.

5 Related Work

Heuristic solutions, such as selecting demonstrations based on the similarity between the demonstrations and test input [19, 37, 32] have been proposed. [20] propose to reorder the demonstration based on the entropy of the predicted labels. In this paper, we use the similarity-based selection method

as a baseline while do not include the label entropy-based reordering method as we show that the ordering of the demonstrations does not matter for our method.

Previous research on the phenomenon of in-context learning in Transformers has identified a number of pre-training data distributions that can lead to the emergence of this capability, including a Hidden Markov Model distribution [50] and a skewed Zipfian distribution with high burstiness [5]. Other studies have sought to understand the underlying mechanisms of in-context learning by making connections with gradient descent [42, 10, 1], formalizing it as an algorithm learning problem [18], or proposing a latent variable theory similar as ours [14, 12, 50]. While providing valuable insights on how in-context learning works, these works are limited to synthetic datasets and toy Transformers, while it remains unclear if these results generalize to LLMs pre-trained on real-world text data and whether these results can help in-context learning performance. In contrast, we propose a Bayesian explanation of in-context learning that can be verified with real-world LLMs on various NLP datasets. Dai et al. [10] provide a practical algorithm based on the understanding that the Transformer has a dual form of gradient descent. However, their empirical results are smaller in scale, with six datasets and only one model (350M), and has less significant improvements (5.4% relative to baseline).

There are also works trying to understand in-context learning from an empirical perspective [2, 24]. Min et al. [26] found demonstrations' ground truth labels do not matter for in-context learning, which we find is not entirely accurate in Appendix B. On the other hand, chain-of-thoughts prompting [48, 53, 45] find that providing step-by-step explanations improves in-context learning performance.

6 Conclusion

In this work, we endeavor to comprehend large language models (LLMs) through a Bayesian lens and posit them as implicit topic models that infer a latent conceptual variable from prompts. Motivated by this understanding, we propose a two-step algorithm that first extracts latent conceptual tokens from a small LLM and then selects demonstrations that have the greatest probability of predicting the corresponding conceptual tokens. The selected demonstrations can then be directly generalized to other LLMs. The efficacy of our algorithm across various text classification datasets and GPT models validates our explanation of in-context learning.

References

E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661 , 2022.

H. Bansal, K. Gopalakrishnan, S. Dingliwal, S. Bodapati, K. Kirchhoff, and D. Roth. Rethinking the role of scale for in-context learning: An interpretability-based case study at 66 billion scale. arXiv preprint arXiv:2212.09095 , 2022.

D. M. Blei, A. Y. Ng, and M. I. Jordan. Latent dirichlet allocation. J. Mach. Learn. Res. , 3 (null):993-1022, mar 2003. ISSN 1532-4435.

T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

S. C. Chan, A. Santoro, A. K. Lampinen, J. X. Wang, A. Singh, P. H. Richemond, J. McClelland, and F. Hill. Data distributional properties drive emergent few-shot learning in transformers. arXiv preprint arXiv:2205.05055 , 2022.

A. Chatterjee, K. N. Narahari, M. Joshi, and P. Agrawal. Semeval-2019 task 3: Emocontext contextual emotion detection in text. In Proceedings of the 13th International Workshop on Semantic Evaluation , pages 39-48, Minneapolis, Minnesota, USA, 2019. Association for

Computational Linguistics. doi: 10.18653/v1/S19-2005. URL https://www.aclweb.org/ anthology/S19-2005 .

J. Chen, L. Chen, and T. Zhou. It takes one to tango but more make trouble? in-context training with different number of demonstrations. arXiv preprint arXiv:2303.08119 , 2023.

W. Chen, X. Ma, X. Wang, and W. W. Cohen. Program of thoughts prompting: Disentangling computation from reasoning for numerical reasoning tasks. arXiv preprint arXiv:2211.12588 , 2022.

K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168 , 2021.

D. Dai, Y. Sun, L. Dong, Y. Hao, Z. Sui, and F. Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta optimizers. arXiv preprint arXiv:2212.10559 , 2022.

L. Devroye, L. Györfi, and G. Lugosi. A probabilistic theory of pattern recognition. In Stochastic Modelling and Applied Probability , 1996.

M. Hahn and N. Goyal. A theory of emergent in-context learning as implicit structure induction. arXiv preprint arXiv:2303.07971 , 2023.

X. Han, D. Simig, T. Mihaylov, Y. Tsvetkov, A. Celikyilmaz, and T. Wang. Understanding in-context learning via supportive pretraining data. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 12660-12673, 2023.

H. Jiang. A latent space theory for emergent abilities in large language models. arXiv preprint arXiv:2304.09960 , 2023.

B. LeBrun, A. Sordoni, and T. J. O'Donnell. Evaluating distributional distortion in neural language modeling. In International Conference on Learning Representations , 2022.

J. Lehmann, R. Isele, M. Jakob, A. Jentzsch, D. Kontokostas, P. N. Mendes, S. Hellmann, M. Morsey, P. Van Kleef, S. Auer, et al. Dbpedia-a large-scale, multilingual knowledge base extracted from wikipedia. Semantic web , 6(2):167-195, 2015.

B. Lester, R. Al-Rfou, and N. Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 3045-3059, 2021.

Y. Li, M. E. Ildiz, D. Papailiopoulos, and S. Oymak. Transformers as algorithms: Generalization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067 , 2023.

J. Liu, D. Shen, Y. Zhang, B. Dolan, L. Carin, and W. Chen. What makes good in-context examples for GPT-3? In Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures , pages 100-114, Dublin, Ireland and Online, May 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.deelio-1.10. URL https://aclanthology.org/2022.deelio-1.10 .

Y. Lu, M. Bartolo, A. Moore, S. Riedel, and P. Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 8086-8098, 2022.

P. Malo, A. Sinha, P. Korhonen, J. Wallenius, and P. Takala. Good debt or bad debt: Detecting semantic orientations in economic texts. Journal of the Association for Information Science and Technology , 65, 2014.

Y. Miao, L. Yu, and P. Blunsom. Neural variational inference for text processing. In M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning , volume 48 of Proceedings of Machine Learning Research , pages 1727-1736, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https://proceedings.mlr. press/v48/miao16.html .

Y. Miao, E. Grefenstette, and P. Blunsom. Discovering discrete latent topics with neural variational inference. In International conference on machine learning , pages 2410-2419. PMLR, 2017.

S. Min, M. Lewis, H. Hajishirzi, and L. Zettlemoyer. Noisy channel language model prompting for few-shot text classification. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 5316-5330, 2022.

S. Min, M. Lewis, L. Zettlemoyer, and H. Hajishirzi. MetaICL: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 2791-2809, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022. naacl-main.201. URL https://aclanthology.org/2022.naacl-main.201 .

S. Min, X. Lyu, A. Holtzman, M. Artetxe, M. Lewis, H. Hajishirzi, and L. Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In EMNLP , 2022.

I. Mollas, Z. Chrysopoulou, S. Karlos, and G. Tsoumakas. Ethos: an online hate speech detection dataset, 2020.

L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. L. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. Training language models to follow instructions with human feedback. arXiv preprint arXiv:2203.02155 , 2022.

E. Perez, D. Kiela, and K. Cho. True few-shot learning with language models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems , 2021. URL https://openreview.net/forum?id=ShnM-rRh4T .

A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. 2019.

N. Reimers and I. Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing . Association for Computational Linguistics, 11 2019. URL https://arxiv.org/abs/1908. 10084 .

O. Rubin, J. Herzig, and J. Berant. Learning to retrieve prompts for in-context learning. arXiv preprint arXiv:2112.08633 , 2021.

E. Saravia, H.-C. T. Liu, Y.-H. Huang, J. Wu, and Y.-S. Chen. CARER: Contextualized affect representations for emotion recognition. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing , pages 3687-3697, Brussels, Belgium, Oct.-Nov. 2018. Association for Computational Linguistics. doi: 10.18653/v1/D18-1404. URL https://www.aclweb.org/anthology/D18-1404 .

N. Saunshi, S. Malladi, and S. Arora. A mathematical exploration of why language models help solve downstream tasks. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=vVjIW3sEc1s .

R. Socher, A. Perelygin, J. Wu, J. Chuang, C. D. Manning, A. Ng, and C. Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing , pages 16311642, Seattle, Washington, USA, Oct. 2013. Association for Computational Linguistics. URL https://aclanthology.org/D13-1170 .

N. Stiennon, L. Ouyang, J. Wu, D. Ziegler, R. Lowe, C. Voss, A. Radford, D. Amodei, and P. F. Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems , 33:3008-3021, 2020.

H. Su, J. Kasai, C. H. Wu, W. Shi, T. Wang, J. Xin, R. Zhang, M. Ostendorf, L. Zettlemoyer, N. A. Smith, et al. Selective annotation makes language models better few-shot learners. arXiv preprint arXiv:2209.01975 , 2022.

H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V . Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan, M. Kardas, V. Kerkez, M. Khabsa, I. Kloumann, A. Korenev, P. S. Koura, M.-A. Lachaux, T. Lavril, J. Lee, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov,

P. Mishra, I. Molybog, Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, P. Xu, Z. Yan, I. Zarov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, S. Edunov, and T. Scialom. Llama 2: Open foundation and fine-tuned chat models. 7 2023. URL http://arxiv.org/abs/2307.09288 .

L. van der Maaten and G. Hinton. Visualizing data using t-sne. Journal of Machine Learning Research , 9(86):2579-2605, 2008. URL http://jmlr.org/papers/v9/vandermaaten08a. html .

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

J. von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677 , 2022.

A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. EMNLP 2018 , page 353, 2018.

B. Wang and A. Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax , May 2021.

X. Wang, J. Wei, D. Schuurmans, Q. Le, E. Chi, and D. Zhou. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171 , 2022.

A. Warstadt, A. Singh, and S. R. Bowman. Neural network acceptability judgments. arXiv preprint arXiv:1805.12471 , 2018.

C. Wei, S. M. Xie, and T. Ma. Why do pretrained language models help in downstream tasks? an analysis of head and prompt tuning. Neural Information Processing Systems (NeurIPS) , 2021.

J. Wei, X. Wang, D. Schuurmans, M. Bosma, E. Chi, Q. Le, and D. Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903 , 2022.

T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771 , 2019.

S. M. Xie, A. Raghunathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=RdJVFCHjUMI .

S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 , 2022.

X. Zhang, J. Zhao, and Y. LeCun. Character-level convolutional networks for text classification. Advances in neural information processing systems , 28, 2015.

D. Zhou, N. Schärli, L. Hou, J. Wei, N. Scales, X. Wang, D. Schuurmans, O. Bousquet, Q. Le, and E. Chi. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205.10625 , 2022.

A Proofs

A.1 Direct direction

Assumption A.1. (Assumption 2.1) Assume that P M ( X ) = P ( X ) , and P d M ( Y | θ , X ) ∝ P ( Y | θ , X ) for X Y θ .

GLYPH<1> GLYPH<0> Proposition A.2. (Proposition 2.2) If task d follows the X GLYPH<1> Y GLYPH<0> θ direction, arg max y ∈Y P d M ( Y = y | θ d , X ) is the Bayes optimal classifier.

Proof. Since the data generation of the task d can be written as Y = f ( X,θ d , ϵ ) , we have

And by Assumption A.1, we have

Thus arg max y ∈Y P d M ( Y = y | θ d , X ) is the Bayes optimal classifier.

Theorem A.3. (Theorem 2.3) If task d follows the X GLYPH<1> Y GLYPH<0> θ direction, then the in-context learning classifier

always has a higher or equal probability of misclassification to the Bayes optimal classifier arg max y ∈Y P d M ( Y = y | θ d , X ) . Equality only takes when

Proof. Recall that in Equation (1), we have

By Proposition A.2, arg max y ∈Y P d M ( Y = y | θ d , X ) is the Bayes optimal classifier. Let C θ ( X ) = arg max y ∈Y P d M ( Y = y | θ , X ) , then the risk is defined as the probability of misclassification

̸

R ( C θ ) = P ( C θ ( X ) = Y ) = E XY [ ✶ C θ ( X ) = Y ] . Denote the in-context learning classifier arg max y ∈Y P d M ( Y = y | X d 1 , Y d 1 , ..., X d k , Y d k , X ) by C k ( X ) . We then have

̸

Such risk is minimized if and only if C k ( X ) = C θ d ( X ) , which only holds when P d M ( θ d | X d 1 , Y d 1 , ..., X d k , Y d k , X = x ) = 1 for all x ∈ X .

̸

A.2 Channel direction

Assumption A.4. Assume that P M ( X ) = P ( X ) , and P d M ( X | θ , Y ) ∝ P ( X | θ , Y ) for the Y GLYPH<1> X θ direction.

GLYPH<0> Proposition A.5. If task d follows the Y GLYPH<1> X GLYPH<0> θ causal direction, arg max y ∈Y P d M ( X | θ d , Y = y ) is the Bayes optimal classifier when the label assignment is balanced.

Proof. Since the data generation of the task d can be written as X = g ( Y, θ d , ϵ ) , we have

When the label is balanced, i.e. P d ( Y ) = 1 |Y| , we have

And by Assumption A.4, we have

Thus arg max y ∈Y P d M ( X | θ d , Y = y ) = arg max y ∈Y P d ( Y = y | X ) is the Bayes optimal classifier.

Theorem A.6. If task d follows the Y GLYPH<1> X GLYPH<0> θ direction, then the in-context learning classifier

always has a higher or equal probability of misclassification to the Bayes optimal classifier arg max y ∈Y P d M ( X | θ d , Y = y ) . Equality only takes when

Proof. This theorem can be proved similarly as Theorem A.3. Recall that in Equation (2), we have

By Proposition A.5, arg max y ∈Y P d M ( X | θ d , Y = y ) is the Bayes optimal classifier. Let C θ ( X ) = arg max y ∈Y P d M ( X | θ , Y = y ) , then the risk is defined as the probability of misclassification

̸

R ( C θ ) = P ( C θ ( X ) = Y ) = E XY [ ✶ C θ ( X ) = Y ] . Denote the in-context learning classifier arg max y ∈Y P d M ( X | Y d 1 , X d 1 , ..., Y d k , X d k , Y = y ) by C k ( X ) . We then have

̸

Such risk is minimized if and only if C k ( X ) = C θ d ( X ) , which only holds when P d M ( θ d | Y d 1 , X d 1 , ..., Y d k , X d k , Y = y ) = 1 for all y ∈ Y .

̸

A.3 Method

Proposition A.7. (Proposition 3.1) When L ( ˆ θ d ) is minimized, P d M ( Y | ˆ θ d , X ) = P ( Y | θ d , X ) for X GLYPH<1> Y GLYPH<0> θ , and P d M ( X | ˆ θ d , Y ) = P ( X | θ d , Y ) for Y GLYPH<1> X GLYPH<0> θ . If the LLM M is invertible, then ˆ θ d = θ d .

Proof. The proof of this proposition is straightforward.

Since

when L ( ˆ θ d ) is minimized, we have P d M ( Y | ˆ θ d , X ) = P ( Y | θ d , X ) for X GLYPH<1> Y GLYPH<0> θ , and P d M ( X | ˆ θ d , Y ) = P ( X | θ d , Y ) for Y GLYPH<1> X GLYPH<0> θ .

If M is invertible, since the embedding matrix is invertible with or without new concept tokens, P d M ( Y | ˆ θ, X ) = P d M ( Y | ˆ θ ′ , X ) implies that ˆ θ = ˆ θ ′ . Thus θ is identifiable, which means ˆ θ d = θ d .

B Experiments

Dateset. In Table 2, we show how we process the text classification datasets into prompts. For each dataset, we take at most 16384 examples from the training set for training, and uniformly sample at most 1000 examples from the test set to test the in-context learning performance. In Table 3, we show the train size and test size we used for each dataset. We also list the set of diverse tasks trained with each dataset, which are denoted by their name in Huggingface datasets. 9 The license for SST2, ETHOS-SO and ETHOS-R is GNU General Public License v3. FPB is under a Creative Commons Attribution-NonCommercial-ShareAlike 3.0 Unported License. Note that these two datasets are hate speech detection datasets for different kinds of hate speech and contain many offensive texts. COLA is excerpted from the published works available on the website, and the copyright (where applicable) remains with the original authors or publishers. DBpedia is under a Creative Commons Attribution-ShareAlike License and the GNU Free Documentation License. EmoC and EmoS should be used for educational and research purposes only.

Experiment details. We run our experiments on A100, V100, and A6000 GPUs. We adopt a large portion of the code from the MetaICL repository [25] 10 . The training takes around 20 to 40 hours on a single GPU. We use a learning rate of 1e-4 and a batch size of 16, and train for 10k steps in total.

Main results. In Table 4, we list the detailed results of our method and baselines with different LLMs on different datasets in Figure 2.

Causal direction results. The detailed results with anti-causal direction (the opposite direction to what we described in Section 4 are in Table 7) are shown in Table 7, corresponding to Figure 6 in the main text.

Other LLMs results. The detailed results with other LLMs are shown in Table 6, corresponding to Figure 3a in the main text.

Random token results. The detailed results with random tokens are shown in Table 5, corresponding to Figure 3b in the main text.

Causal

Anti-causal

76.3

70

73.6

60.8

53.4

50.8

56.7 53

57.7

54.1

55.6

38.837.4

31.5

26.6

14.8

SST2

FPB

COLA

Dbpedia

EmoC

EmoS

ETHOS-SOETHOS-R

k -ablation study results. The detailed results of k ablation study are shown in Table 10, corresponding to Figure 4a in the main text. In this experiment, we do not reorder the selected demonstrations according to Equation (3), as we need to use GPT2-large for the reordering, and it cannot fit in all the demonstrations. Instead, we order the selected demonstrations from the largest ˆ P d M ( θ d | X d , Y d ) to the smallest.

c -ablation study results. The detailed results of c ablation study are shown in Table 11, corresponding to Figure 4b in the main text.

Effect of using ground truth labels. According to [26], the ground truth label is not necessary for demonstrations to have a good in-context learning performance, which we found is not entirely true for all the tasks. We compare our method with the randomly selected demonstration baseline under three scenarios: (a) Original : demonstrations with the correct labels; (b) Random words : using a random label projection map τ d instead of a meaningful one. i.e., map each label to a fixed random word. In this case, the mapping from the input tokens X to the labels Y is still preserved; (c) Random labels : assign a random label to each demonstration, with the original label projection map τ d . As shown in Figure 7, by using a random label projection map or randomly assigning the labels, the performance of the randomly selected demonstration baseline drops considerably. And randomize the label assignment gives a larger performance drop than only using a random label projection map,

Uniform

Ours

64.8

57.4

50.5

47.5

35.9

35.3

Original

Randomwords

Randomlabels

Optimal

Ours

90.3 86.2

86.1

92.6

86.282.5

78.276.6

75

69.1

60.4

56.5

57.3

53.8

48.4

38.6

SST2

FPB

COLA

Dbpedia

EmoC

EmoS

ETHOS-SOETHOS-R

which shows that the mapping between X and Y in the demonstrations matters. This indicates that in-context learning infers the mapping between X and Y from the demonstrations instead of merely invoking some learned function stored in the LLM parameters based on the appearance of X and Y . We also show that the demonstrations selected by our method represent the X -Y mapping better, as under the Random words condition, our method performs better than the random selection baseline, while our method does not improve the random selection baseline under the Random labels condition. The detailed results with random words and random labels are shown in Table 8

Optimal performance As stated in Theorem 2.3, the optimal performance of an in-context learning classifier is the Bayes optimal classifier arg max y ∈Y P d M ( Y = y | θ d , X ) , which is approximated by using the learned concept tokens as prefixes. Note that this approximated Bayes optimal classifier cannot be transferred across different LLMs, as the learned concept tokens embeddings are aligned with a specific LLM. The advantage of in-context learning with our method is that the demonstrations can be transferred to any LLMs without training. Here we only compare the accuracy of in-context learning with our method and the approximated Bayes optimal classifier using GPT2-large, as it is the LLM that concept tokens are fine-tuned with. As shown in Figure 8, our method comes close to the optimal accuracy on many datasets, while there are some datasets that our method is lagging. This indicates that there are two ways to improve our method: the first is to improve the performance of the optimal classifier, by introducing a better latent concept learning algorithm. The other way is to reduce the performance gap between our method and the optimal classifier, by improving the demonstration selection algorithm. The detailed results using the learned concept tokens as prefixes are shown in Table 9.

Reordering results. We reorder the selected demonstrations to maximize the posterior of the concept tokens:

Where π (( X d 1 , Y d 1 ) , ..., ( X d k , Y d k )) is a permutation of ( X d 1 , Y d 1 ) , ..., ( X d k , Y d k ) . Π is the set of all possible permutations of the k demonstrations. The detailed results with and without reordering are shown in Table 12, corresponding to Figure 9.

Similar tokens. We show the top ten similar tokens to some learned concept tokens in Table 13, as summarized in Figure 5 in the main text.

Reorder

Notreorder

06

工

80

工

70

60

50

40

I

30

20

10

SST2

FPB

COLA

Dbpedia

EmoC

EmoS

ETHOS-SOETHOS-R

Likelihood histogram. We also show histograms of the probability of each example predicting corresponding concept tokens in different datasets. We can see that the probability of prediction concept tokens can well differentiate examples in a dataset.

Selected demonstrations. Table 14 shows the selected top 4 demonstration by our proposed algorithm.

C Limitations and Future Work

While the assumption that a large language model captures the true distribution of language is fairly common in the literature studying LLMs [50, 34], this assumption is not entirely accurate in practice. According to [15], LLMs systematically underestimate rare text sequences, which constitute a significant portion of the long-tail distribution of language. Although this assumption is adequate to achieve favorable empirical results, it is expected that more accurate language models will, in theory, lead to improved outcomes.

The selection of the accompanying diverse tasks S is currently left to the user's discretion. A better approach to constructing such a task set is needed to gain a deeper understanding of latent concept variables and to improve the latent concept learning algorithm.

Our algorithm currently only applies to classification tasks. More complex latent variables could be designed to improve the in-context learning performance of more complex tasks like math word questions and logical reasoning problems.

D Broader Impact

The utilization of language models (LLMs) for specific tasks is often hindered by the high cost associated with training or fine-tuning them. However, the in-context learning paradigm offers a cost-effective and convenient alternative for utilizing the power of pre-trained LLMs. Our work has demonstrated a significant improvement in the performance of in-context learning through a relatively low-cost and simple approach, thus making the use of LLMs more accessible for individuals with limited resources.

However, it is important to consider the broader implications of the increasing use of LLMs. As LLMs are not infallible and may make mistakes, it is crucial to explicitly warn users of the potential for misleading output and to regulate the distribution of LLMs in order to prevent any negative societal impact. Additionally, it is possible that LLMs could be intentionally misused, thus it is important to consider the ethical implications of their use and to take appropriate measures to mitigate

any potential negative effects. We posit that these regulations and measures should be put in place at the time of distributing LLMs to ensure the safe and responsible use of these models. Furthermore, as we publicly release our code, we will also provide clear warnings and guidelines to users to ensure that the potential risks associated with the use of our method are fully understood and addressed.

(a) SST2

(b) FBP

(c) COLA

(d) DBpedia

(e) EmoC

(f) EmoS

(g) ETHOS-SO

(h) RTHOS-R

Histogram of FPB

8

6

Percentage

4

0

0.020

0.025

0.030

0.035

0.040

P(0|X,Y)

Histogram of SST2

6

5

Percentage

4

3

2

1

0

0.24

0.26

0.28

0.30

0.32

0.34

0.36

P(0|X,Y)

Histogram of DBpedia

8

7

6

Percentage

5

4

3

2

1

0

0.015

0.020

0.025

0.030

0.035

0.040

P(0|X,Y)

Histogram of COLA

6

5

4

Percentage

3

2

1

0.09

0.10

0.11

0.12

0.13

0.14

0.15

0.16

P(0|X,Y)

Histogram of EmoC

8

/

6

Percentage

5

4

3

2

1

0

0.002

0.003

0.004

0.005

0.006

P(0|X,Y)

Histogram of EmoS

8

7

6

Percentage

5

4

3

2

1

0

0.0010

0.0015

0.0020

0.0025

0.0030

0.0035

P(0|X,Y)

Histogram of ETHOS-SO

10

8

Percentage

6

4

0

0.001

0.002

0.003

0.004

0.005

P(0|X,Y)

Histogram of ETHOS-R

10

8

Percentage

5

4

Q

0.001

0.002

0.003

0.004

0.005

P(0|X,Y)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a new perspective to understand the in-context learning behavior of large language models (LLMs) from the angle of latent concept learning resembling topic models. Based on the generative process defined by input data, latent concepts, and labels, the authors propose a two-stage algorithm to first learn the latent concept and then use it to select the best-performing demonstrations that boost in-context learning performance. Experiments on eight datasets show that the method is able to consistently outperform random selection and selection based on semantic similarity.\n\nSTRENGTHS:\n* Originality: Although the latent concept learning perspective is largely inspired by Xie et al., the goal and method proposed in this paper for selecting effective demonstrations is still sufficiently interesting and different from Xie et al., which seem novel to me. \n* Quality: The paper first states its generative assumptions defined by input data, latent concepts, and labels, based on which it then derives a two-stage method to first learn the concept and then select demonstrations. The theories and methods seem solid to me.\n* Clarity: The paper is overall clear and nicely written.\n* Significance: While it's nice to see that the method is able to consistently outperform simple baselines across the tasks selected, I generally feel that the setting considered in this paper is somewhat artificial, and that the evaluation is not comprehensive enough to testify the generalization ability of the method. These concerns weaken the significance of the paper. See weaknesses below for details.\n\nWEAKNESSES:\n* Problem setting: The paper assumes access to a (relatively) large training set from which a few demonstrations can be drawn. With this amount of training data (e.g., 100), there could be better alternative choices than in-context learning. For example, one can tune an LLM with parameter-efficient methods and easily outperform in-context learning (shown in Liu et al.) without even having to consider how to select the best demonstrations. Of course, one can argue that in-context learning is applicable to non-open-source LLMs while training-based methods are not. However, given the recent growth in the availability of open-source LLMs such as LLaMA and Falcon, I believe it would be generally better to consider parameter-efficient tuning than in-context learning if the authors assume ~100 training samples are available. In summary, I feel that the problem setting of \"carefully selecting good demonstrations from a larger training set\" is somewhat artificial on its own from the very beginning.\n* Evaluation tasks: The eight downstream tasks selected for evaluation appear to be too easy for nowadays LLMs, and it's unclear given the current evaluation how generalizable the method is to more challenging tasks like MMLU and reasoning, potentially when combined with chain-of-thought prompting.\n* Typo: Line 104: \"semantic analysis\" -> \"sentiment analysis\"\n\nReference:\nLiu et al. \"Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning.\" NeurIPS (2022).\n\nQUESTIONS:\n* What is the role of $\\boldsymbol{\\epsilon}$ in the generative process (Line 91)? It seems like it is not taken into account in the current modeling process. \n* How does the method work for instruction-tuned models (e.g., ChatGPT, Alpaca)?\n* Why are the results of LLaMA that bad (Figure 3a)? LLaMA is generally recognized as a better-performing model than OPT series under similar sizes.\n\nLIMITATIONS:\nPlease see the Weaknesses section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a novel demonstration selection method aimed to enhance performance of (few-shot) in-context learning. The approach is characterized by a two-stage process that includes latent concept learning and demonstration selection, each holding a unique significance. In the latent concept learning stage, the authors present a method for acquiring a task-specific token embedding set through prompt-tuning. Subsequently, the demonstration selection stage involves the process of selecting in-context samples. This process is based on maximizing the likelihood of post-fixing the previously acquired task latent. The efficacy of the demonstration selection has been evaluated across various language models, leading to improved in-context learning performance.\n\nSTRENGTHS:\nThis paper is clearly written with a clear definition of the task. \nThe proposed method is presented with great clarity and detail, which significantly aids in understanding the overall procedure. \nThe paper presented impressive performance results. The successful transfer of demonstrations selected from a smaller model (GPT-2) to other larger models is very impressive.\nThere are extensive ablation and additional experiments for deeper analysis of each component of the proposed method.\n\nWEAKNESSES:\nAn area of concern lies in the assumption that the task latent derived through prompt-tuning is considered the “optimal task latent” (line 146). This assumption may not hold universally, and could be re-considered as the use of manual prompts - which could offer a more intuitive understanding of the task.\n\nQUESTIONS:\n- Considering that the order of samples in the demonstration does not play a critical role in the proposed method, what motivated the decision to include a \"re-ordering\" (selecting the permutation) step in the demonstration selection phase? \n- How would the performance change if a larger model, as opposed to GPT-2, was employed for demonstration selection? Would this potentially enhance the overall performance?\n\nLIMITATIONS:\nLimitations are stated in the appendix C.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper describes a framework to select demonstrations for in-context learning by using bayseian formulation for the data generation process. Similar to topic models, the formulation uses a \"concept\" variable and words in this sequence are conditioned over this concept variable and are conditionally independent of the other tokens in the generated text. The concept variable models the prompt/task instruction and is modeled by learning concept tokens by prompt tuning a small LLM. Since the concept tokens are made part of the vocabulary, the selection of in-context demonstrations can be learnt by maximizing the probability for the concept tokens. Experiments are presented on multiple NLP tasks and they indicate that selecting in-context demonstrations using this formulation works better than using random in-context examples.\n\nSTRENGTHS:\n1. Simple formulation that aids better selection of in-context demonstrations\n2. Experiments have been presented on multiple NLP tasks\n\nWEAKNESSES:\nModeling: I'm not sure the topic-model like bag-of-words assumption is accurate in the way the modeling has been described. As Line 56 in the paper states, the generation of tokens would be independent of the previous tokens but to truly model this wouldn't you need to modify the current latent concept learning setup to be sequence-agnostic? The method works empirically so perhaps its okay but some clarification would help here. Perhaps the authors can explain further in the rebuttal phase if I misunderstood it.  \n\nExperiments: The experiment baselines refer to the \"Similar\" baseline but any discussion or analysis of its results is completely missing from the main paper. I also checked the appendix, and the results in Table 3 (appendix) do not correspond to the results reported in the main paper in the histogram plots.\n\nNotation: The notation is a bit hard to follow and could benefit from simplification.\n\nQUESTIONS:\nSee weakness\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper tries to study in-context learning through a Bayesian lens, namely treating LLMs as implicit topic models. It proposes an algorithm to select optimal demonstrations from a set of annotated data with a small LLM and then use the selected demonstrations with larger LLMs. 12.5% improvement compared to random selection is observed by adopting the proposed algorithm.

STRENGTHS:
The paper has extensive empirical results showing the proposed algorithm can effectively find helpful demonstration examples to boost the performance of in-context learning.

WEAKNESSES:
Although we see significant performance improvement, it's hard to conclude LLMs are topic models. It could be possible that topic (or distribution of input text) is one of the  important factors but there are other attributes affecting in-context learning. In some ways, this also conflicts with some previous work, e.g. "In-context Learning and Induction Heads" points out the pattern copying behavior is the key, "Robustness of Demonstration-based Learning Under Limited Data Scenario" finds out random tokens are also helpful.

QUESTIONS:
Have you tried to compared with better algorithms to find demonstrations, e.g. "Selective Annotation Makes Language Models Better Few-Shot Learners"?

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-17

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work aims at proposing a demonstration example selection algorithm in in-context learning, using a formulation of topic models for language models. Specifically, the proposed algorithm applies prompt tuning on some prefixing learnable tokens to obtain latent concepts, and select demo examples that maximize the probability of inferring the learned tokens. The authors test the algorithm on small LMs and find generalizability to larger LMs, on a set of classification tasks.\n\nSTRENGTHS:\nThe writing of the paper is overall clear. The theoretic formulation is well-structured and points out where assumptions and approximations are made. The proposed algorithm is tested on a set of tasks and models. The ablation studies are also reasonable.\n\nWEAKNESSES:\nThe motivation and takeaway of this work are rather vague. As a demonstration example selection algorithm, the setup is not realistic enough since it uses prompt tuning with labeled task data to obtain the latent concept tokens. Subsequently, the comparison with the baseline methods using uniform or similar demonstration examples is not fair, since they do not assume prompt tuning with labeled task data. \n\nSome natural and probably necessary questions here include: (1) how does the performance of the proposed method compare with the vanilla prompt tuning performance? and (2) since labeled task data are used, what is the performance of directly (and independently) select ICL examples based on their contribution to the predictions of the labeled task data? \n\nAdditionally, the current analysis is mostly on the performance of the algorithm rather than the selected demonstration examples themselves. What attributes do they share in common? (apart from the qualitative clustering shown in the tSNE figure) What label distribution do they have? (do they simply act as a *calibration* to the model's output distribution?) In other words, probably as the title hints, what are good demonstrations for in-context learning?\n\nQUESTIONS:\nPlease see the weakness section for my main questions to the authors. Additionally, is any generation task considered in this work apart from the classification tasks?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}