Large Language Models Are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning
Authors: Xinyi Wang
Wanrong Zhu
Michael Saxon
Mark Steyvers
William Yang Wang
Venue: NeurIPS 2023 poster (2023)
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
Xinyi Wang', Wanrong Zhu', Michael Saxon', Mark Steyvers", William Yang Wang'
*Department of Computer Science, University of California, Santa Barbara
*Department of Cognitive Sciences, University of California, Irvine
sxinyi_wang, wanrongzhu, saxon) edu, msteyverGuci. edu, william@ca ucsb. edu
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 Bayesian lens, viewing real-world LLMs as latent variable models. On this premise, propose algorithm to select optimal demonstrations from set annotated data with small LMM, and then directly generalize the selected demon- strations 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- aolld usefulness of our algorithm GSM8K, math word problem dataset. Our empirical findings support our hypothesis that LLMs implicitly infer latent variable containing task information.
Introduction
Transformer-basee [41] pre-trained large language models (LLMs) have demonstrated significant advancements in a variety of natural language processing (NLP) tasks. As the size 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 example input sequence for semantic analysis with in-context learning:
Great movie. Positive. An The worst movie ever. Negative. An Can 't wait to see the second movie!
The first two lines are two demonstrations, and the third line a test input. We expect an LLM to output the correct label Positive continuation.
In-context learning has been demonstrated to an effective technique for 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
'Code: https: /ttps://githhb
saseeepppt-baeedddeee demeeddddmmonttatioon selection.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

[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 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 task-specific data distribution 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 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 practical demonstration selection algorithm. Our explanation is inspired by the generation process of a topic model, i.e. a simple latent variable model:
P(wlr) / P(w.rr())
Where 0 O represents a potentially high dimensional topic/concept variable. θ the space of the topic/concept variable, and wirr 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:
T
P(WIT)
II P w;-1..... w1) 1=1
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:
Par (w w1:c) = -/o PM (We+1:7 (0) PM (0) w1 (e)
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 Wi: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( (). As LLMs do not explicitly learn 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 (wiz in the equations) that are most likely to infer the optimal latent variable value (those with the highest posterior probability P(O w )). We then directly use them as demonstrations for in-context learning with other larger LLMs (up to 175B parameters) and observed significant performance improvement. The general- ization of demonstrations between LLMs likely a result of similar pre-training data distributions.
While our work is inspired by that of Xic et al. [50]. our approach differs significantly in both theoretical analysis and experimental settings. Our main contributions
follows:
We assume a general data generation process specified by 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 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 em- pirically validated using real-world L.MSs 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 w1:t is composed of several demonstrations and a test input. The generated tokens Wr+1.7 represent the model's 's prediction for the test input.
2.1 Notations and Problem Setting
Suppose the objective of our task is to predict discrete target variable Y E V, given token sequence X E X. where X' is the space of all possible token sequences. E θ 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 underlying causal relation between X, Y, and 0. We examine two potential directions of this causal relation, namely - Y - 0 and Y X 0, which can be represented mathematically the following structural equations:
y= f(x
X F= 0,
Here E is an independent noise variable, f d' x y and g V θ 8 X' are two deterministic functions. Furthermore, we denote the joint data distribution by X, 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.
0 0 ~ and
We hypothesize that the causal direction depends on the nature of the task. For instance, in the task of predicting the sentiment (Y) of movie review (X), is reasonable 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 passage to express one's opinion about the movie" (0). Conversely, for the task of classifying whether product review (X) is helpful to other customers (Y), it is the quality of the review (X) that cause other customers to upvote (Y), along with the task concept of "rating the helpfulness of this review" (0). In the rest of the paper, we will focus on the X - Y + 0 direction and leave detailed discussion of the other direction in the Appendix.
Suppose we are interested in a task (e.g. semantic analysis) denoted by d E T, where T is the space of all possible tasks. We assume there is injective function between T and θ. i.e. for each task d, there is concept variable 0d, such that each data (Xd yd) sampled from task is generated by:
yd f(x 0d
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 Y (Xg, Y. Y:!),... (X Ytd) sampled from task d.
Following previous works [24, 26],
not using any instruction fine-tuned models, 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 7d > + X, which are typically defined by human understanding of the task d. e.g. for sentiment analysis, 7 map positive class the token "positive" and negative class to the token "negative". Additionally, a delimiter token und is defined, typically an empty space or a new line token, to separate the demonstrations when concatenated. We denote the I.LM output probability of X, Y, and O with the aforementioned preprocessing applied, by P
Par (Y) (MT)) w"....... _ 7 (YT) ,x Pd
Yd
X y.
2.2 Problem Analysis and Theoretical Results
Suppose a set of observed data sampled from task d, denoted as Ddd is available, allowing for the selection of the l most suitable demonstrations from it. For any incoming test example X, we have:
P (Y Yd !... Xg. yg, x) /
pd (Y () [0 yd,.... .... xg. Y X)d0
(I)
3

---

Latent Concept Learning
Demonstration Selection
context Learning
Yest
LIM
LLM
X
compui
compuux
Test
Lodote
Dross artrazy los
Dataset
Select par
Languege node
Inter
Figure I: An overview our proposed two-phased algorithm. Demonstration selection and latent concept learning share the same L.LM as demonstration selection needs to reuse the learned concept tokens, while at the in-context learning time, any other generative LLMs can be used. Here we only illustrate the X Y 0 direction. The Y X 0 direction can be illustrated similarly by exchanging X and Y in the above figure.
Here, we assume the sampling of the test example is independent of the sampling of the demonstra- tions, so Y is independent of the demonstrations given and X. We also assume that the pre-trained data distribution P is suitable approximation of the assumed data distribution P:
Assumption 2.1. Assume that (M = P(X), and Ph O, X) x P(Y for
0.
Note that the assumption that a large language model captures the true distribution of language is fairly common in the literature studying L.MSs (50, 34, 47]. With this assumption, we establish: Proposition 2.2. If task d follows the X + Y direction, then arg max PI (Y = y100, X) is
the Bayes optimal classifier.
.X) completely concentrate on 0 can the in- context learning classifier become the Bayes optimal classifier [11]:
In this case, only when Pdr (0|x ".. .x)
Theorem 2.3. If task d follows the X + Y + - direction, then the in-context learning classifier
arg max Pd (Y y|X Y .g. xd. yd X)
yEY
always has a higher or equal probability of misclassification to the Bayes optimal classifier arg max Pd (Y yy00, X). Equality only holds when
Vx E X, Pd (0 (x Yd ..., X2. Y X x)
A similar argument can be made for the Y + X - direction." Here, Equation (1) would become:
Ps (X AX Xy. Xq.... Y X Y) Pd (X
X Xf...., .... Yd. X Y de
(2)
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 do treat (Y - X - 0) as the universally superior channel direction for modeling in-context learning, rather arguing that depending on the end task, the causal direction (X Y 0) is sometimes better. This view is supported by our empirical results in Appendix B.
Method
Here we demonstrate how the proposed theory can be practically applied to select optimal demon- stration 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
First, we perform latent concept learning, wherein the task latent 6d 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 postfixing the latent concept tokens. We only need use small LLM to do the above steps to obtain an optimal set of demonstrations that can be directly transferred to other LLMs. Figure is an overall illustration of our proposed method.
The detailed argument of the Y X direction can be found in Appendix A.2.

---

Algorithm Latent concept learning
Input:
Dataset D = {11 yi, d:) associated with a set of tasks S, LLM M, number of concept tokens per task c, learning rate a, and number of training steps N.
Output: LLM M' with fine-tuned concept tokens.
Add c|8| new tokens to the vocabulary. i.e. The concept tokens 6 for each task in S. Randomly initialize their embeddings Enew- Freeze all parameters M except Enews for step to N do
Sample a random batch B in DD and initialize gradient t- 0; for each data point yz, v,d) in B do
- 9 + ae(X,y.0"), end for Enew Enew - agi end for
dEnew
3.1 Latent Concept Learning
We want to first find the optimal value the latent concept variable 6 corresponding to task d T. As arg max Pār (Y y/0", X) is the Bayes optimal classifier according to Proposition 2.2, dd should be able minimize - Ex,,y,a Da,y,a(log (Y X)] for the X Y A direction. In practice try align 8 to the token embedding space by adding new tokens the vocabulary. After this alignment, we hope to be able to use the learned new tokens of 1d as regular tokens.
More specifically, building upon the methodology proposed by Lester et al. [17], for each specific task d, new concept tokens (denoted 04) are added to the original vocabulary of LLM M to represent the corresponding task concept 0 Subsequently, the embedding of these new tokens Enew (Qd) is fine-tuned while freezing the remaining parameters L.MM M. The variable i treated as hyperparameter. In practice, in order to condition on 0 the corresponding e concept tokens are appended to the input X (or Y) as shown in the example provided below, where c 2:
<sentiment_ tokentttt_tt___eeeet westttttmetttttttm.mmmemmmmemmee teoeemeettttttttet 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 label assigned to the newly added concept token. It can be anything as long as it does not overlap with the original vocabulary of L.LMM
The fine-tuning objective would then be minimizing L(0") = Ex,y [e(X, Y; (d)], where
{
log Pd (Y 110 X) i Y - log Pd (X (dd Y) if + X 6.
(X,Y:
-1-
Theoretically, if we can minimize the above loss function, Bayes optimal classifier can obtained, and the concept tokens would be a reasonable delegate of the real latent concept variable:
Proposition 3.1. When L(Od) is minimized, Pd (Y)0d, X) P(Y X) X Y 0. If the LLM M is invertible, then 0 -
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 Par Por (#'|w1:) (w1:t denote given prompt) would be in the token sequence space X' instead of the concept space O. Since Icarning all possible 0d θ is infeasible,
propose approximate the concept space θ by sampling diverse subset of tasks S C T. Then the estimated conditional probability of 6d would be:
Pd (():)) pd (00) (1++) = ELes Pt (0' (p)))a))
To obtain the concept tokens for all tasks in S, fine-tune all tasks together with the loss Edes L(Od). We summarize the proposed algorithm in Algorithm
Mooe discussion be found in Appendix A.3.
5

---

Algorithm 2 Demonstration selection
Input: dataset Dd for a task d. L.M with fine-tuned concept tokens M'. The number of demonstrations k:
Output: A set of selected demonstrations.
for each (Xd,, Yd) in Dd do
Compute M (04) X yd);
end for Select top & examples with the largest Pdr (0d [X yd). denoted (Xq, Y (Xg.
Note that the embedding matrix of 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 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, Pd [04] Ydd , "...,,X, X YE X) for all X E X. Then our goal then becomes selecting demonstra- tions that can best infer the task concept for all test inputs on average:
need to select demonstrations (X Yi),..., (X² Y2 that maximize
Ex[PA (0 Xt, Y X YE X)]
arg max X X
As test examples are sampled independent of the demonstrations, and (x)(X) = P(X) according to Assumption 2.1, ,,1, we have
Ex [ (0 |X] YTd X yd. x)] = Ptt xx,, yI....,,;" Y,d)
we assume each demonstration is also sampled independently, we have:
( XT. yd 1 X X.YY) I Pit x y:d)
P (04)--1
Assuming that 0 has a uniform prior, then our goal becomes finding the top k demonstrations that maximize pM | Y!a) Note that the independence between demonstrations assumption to reduce the combinatory search space of (X YT'),... (Xg. Yg). In practice, selected demonstrations are likely correlated as some demonstrations may work well together but not nec- essarily work well by themselves. However, it would be expensive to search the O ({D combinations the candidate set Dd. In practice, this simplification works reasonably well. We leave this combinatory search problem to future research.
simplified
Also, 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 NL.P 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 different ontology classes. For emotion classification, we choose the dataset from Chatterjee 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.5
65.1
64.2
60.3
59.7
60
61.3
62.4
60.4
36.8
56.2
56.7
56.5
57.7
53.7
CPTS-babiaage
GFT3-davini [1758]
(I2eM)
(S55M)
(276M
(1.00)
BS50!
[1.50]
[6.70]
Figure 2: Accuracy of 4-shot in-context learning using demonstrations selected by our method and
other baselines, averaged over eight datasets. Our demonstrations are selected using GPT2-large, and
the same set of demonstrations is then applied to all other LLMs.
predict emotion given three-turn contextual dialogue, while Saravia [33] predict emotion given 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 V 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". We adopt the sentiment analysis, topic classification, and emotion classification tasks, which is consistent with the intuition that people usually have some sentiment, topic, or emotion mind before writing a piece of text. We adopt the X + Y <- 0 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 post hoc task in contrast to the previous tasks.
X direction for
Without specification,
use k number of demonstrations and e 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 randomly selected 100 example subset of the train set as the candidate set Dd 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 & demonstrations from D for each test example.
Similar: According to Liu et [19], demonstrations that are semantically similar to the test example would hare more performant. Following their method, we use treitraneed sentence Transformer [31] to calculate the cosine similarity between the demonstrations and test examples. We choose the top A similar demonstrations from D for each test example.
Main results." Figure shows our main results averaged over all eight datasets, using the first- generation GPT2s and GPT3s, without any instruction fine-tuning [28] 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 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. Figure 3a, we test the demonstrations selected by our method using GPT2-large more LLMs (GPT3 [4], GPT3-instruct [28, 36), GPT-J [44], OPT [51]. and ILaMA [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
4 Detailed results see Figure 6 Appendix B.
'The complete results with standard deviations in this section can be found in Appendix B.
7

---

Uniform
Random tokens
Ours
62.3
52.8
56
58.8
02.7
$6.6
61.2
64.8
57.4
57.2
53.7 53.7
36.4
GPT3-curie (6.78)
GPT3-cureeinsturct (6.78)
GPT-1 (68)
OPT-6,70
LLMM-78
GPT2 GPT2-arge (774M)
(b) Proposed method v.s. using randomly selected tokens
(a) Proposed method v.s. randomly selected demonstrations
Figure 3: In-context learning accuracy averaged over all eight datasets.
Uniform
Similar
Ours w/ Llama 2 (7B)
Ours GPT2-XX XL (1.5B)
Prompt tuning
3.7
1.3
Llama (7B)
11.4
13.1
19.3
15.9
Llama 2 (13B)
17.0
18.3
21.6
20.5
Llama 2 (70B)
50.2
53.5
54.3
52,9
ChatGPT (gpt-3.5-turg
76.5
78.1
81.2
80.4
Table Prompt tuning and 4-shot in-context learning accuracy on subset of GSM8K test set. Our demonstrations are selected with either 7B Llama 2 or GPT2-XL
selected demonstrations. Different-size GPT2 models share the same pre-training corpus [30], while GPT3s are pre-trained on 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 GSM8K. Since 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-efficiett 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 word problem-solving dataset with chain-of-thoughi solutions. Table shows the test accuracy of the final numerical answer with greedy generation. We randomly select test of 200 examples instead of using the full test set for computation efficiency.
dataset, GSM8K [9], which math
As shown in the first row of Table 1, prompt tuning with ten new tokens 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 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 larger model (7B) more effective than those selected with smaller model (1,5B). The results show that our demonstration selection method a good choice under a low data setting, with 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 method comes from the learned concept tokens containing the task and format information, not other elements of our algorithm.
A ablation study. While we use & 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 & 2, 4, 8, and 16. To fit in large ks, we use GPT3-ada with longer context length (2048). Note that for real-world tasks, that more demonstrations guarantee higher performance performance increases from & to = 8, then drops a little at = 16. Our method improves the uniform baseline by around 5% absolute for all ks, while h improves the most (6.6%). Our method appears to have diminishing effect when & becomes larger, which likely because the effect of more demonstrations overwhelms the effect of demonstration choices.
general not true [7]. We see that the uniform baseline
'Note that we did not use calculator to insert the correct result of each generated math equation during generation for time efficiency, which resulted in slightly lower scores.

---

ablation study. While we use 10 number of concept tokens for all ex- periments, we also investi- gate the effect of different c on our method. When с small (c 5), the concept tokens cannot effectively capture the task and format information, thus cannot im- prove the performance. When increases from 10 to 20, observe drop in the performance. It likely because the selectivity of the concept tokens decreases when c increases. The longer the concept token sequence is, the more likely will contain meaningless tokens that do not contribute demonstration selection.
Uniform
Ours
63.4
62.4
60.92
56.8
58.8
57.4
56.6
54.7
k-8
k-16
Unilorm
(a) ablation study.
(b) c ablation study.
Figure 4: In-context learning accuracy of our method versus random
selection baseline averaged over all eight datasets with GPT3-ada.
Effect of demonstrations' order. We find that the demonstrations selected by our method are insensitive to their order most cases." An exception 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 performance is high with different demonstrations, thus ordering matters. And since such ordering is not transferable while 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 t-SNE [40] projection of the learned concept token embeddings. The tokens corresponding to semanti- cally similar tasks are close together. Note that this result only aims pro- vide a straightforward illustration of concept tokens. The effect of concept tokens should be understood by the previous quantitative results.
Web technology
numbers
Negetive finance
Negative emotions
Positive finance
Positive
cherscteristics
HrS
SS12
COLA
Dlieea
ETHOS so
Medical
ETTIOS
EmoC
Environment
0-
also list the top selected demon- strations in Table 14 in Appendix B. Compared to the examples with lower scores, the selected examples for GSM8K have more deductive reason- ing (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 "but'. This can be understood as these 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.
40
Adverbs
Emotionally unstablle/mmaaturee
Figure 5: t-SNE plot the learned concept tokens for each task. Concept tokens that can be explained by similar tokens are summarized in the graph.
Related Work
Heuristic solutions, such as selecting demonstrations based on the similarity between the demonstra- tions 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
7 Detailed results see Figure Appendix B.
'The list of similar tokens for these concept tokens can be found in Table 13 Appendix B.

---

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 Hidden Markov Model distribution [50] and 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, IJ. formalizing it as an algorithm learning problem [18], or proposing 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 these results generalize LLMs pre- pre-trained on well-world text data and whether these results can help in-contex: learning performance. In contrast, we propose a Bayesian explanation of in-context learning that can be verified with real-world L.L.M on various NLP datasets. Dai et al. [10] provide a practical algorithm based on the understanding that the Transformer has 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, thain-of-thouulhe 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 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.
Acknowledgements
This work was supported by the National Science Foundation award #2048122. The views expressed are those of the author and do not reflect the official policy position of the US government. We thank Google for its generous gift to the University of California.
References
[1] 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.1500 2022.
[2] H. Bansal, K. Gopalakrishnan, S. Dingliwal, S. Bodapati, K. Kirchhoff, and D. Roth. Rethinking the role of scale for in-context learning: An interpretability-basse case study at 66 billion scale. arXiv preprini 22Xi::22/2.00000 2022.
[3] D. M. Blei, A. Y. Ng, and M. I. Jordan. Latent dirichlet allocation. J. Mach. Learn. Res., 3 9022 mar 2003. ISSN 1532-4435.
[4] 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, 13:1877-1901, 2020.
[5] 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 22X.:22055.00555 2022.
[6] 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
10

---

Computational Linguistics. doi: 10. 10.18653/v1/$19-2000 URL https://www..cllet org/ anthology/S13 2005.
[7] J. Chen, L. Chen, and T. Zhou. It takes one tango but more make trouble? in-context training with different number of demonstrations. arXiv preprint arXiv:2303.081 08,1,, 2023.
[8] 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.122 12588, 2022.
[9] 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:21 10. |4168, 2021.
[10] D. Dai, Y. Sun, I.. Dong, Y. Hao, Z. Sui, and F. Wei. Why can gpt learn in-context? language models secretly perform gradient descent meta optimizers. arXiv preprint arXiv:2212.10559, 2022.
[11] L.. Devroye, L. Györfi, and G. Lugosi. probabilistic theory of pattern recognition. Stochastic Modelling and Applied Probability, 1996.
[12] M. Hahn and N. Goyal. A theory of emergent in-context learning as implicit structure induction. arXiv preprint arXiv:2303.079// 2023.
[13] 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.
[14] H. Jiang. A latent space theory for emergent abilities in large language models. arXiv preprint arXiv:2304.0990 2023.
[15] B. LeBrun, A. Sordoni, and T. J. O'Doneel. Evaluating distributional distortion in neural language modeling. In International Conference on Learning Representations, 2022.
[16] Lehmann, Isele, M. Jakob, A. Jentzsch, D. Kontokostas, N. Mendes, Hellmann, M. Morsey, P. Van Kleef, S. Auer, et al. Dbpedia-a large-scale, multilingual knowledge base extracted from wikipedia. Semantic web, 672):167-195, 2015.
[17] B. Lester, R. Al-Rfou, and N. Constant. The power of scale for parameter-effiiient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 3045-3059, 2021.
[18] Y. Li. M. E. Ildiz, D. Papailiopoulos, and S. Oymak. Transformers algorithms: Generalization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067 2023.
[19] 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: 1/1665//1120222deeelioo 1.10. URL https: : //aaaaaaanooogy......o org/ 2022/022..2 deelio-1.11.. 1.100.
[20] 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 I: Long Papers), pages
2022.
[21] 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.
[22] Y. Miao, L. Yu, and P. Blunsom. Neural variational inference for text processing. In M. F. Balcan and K. Q. Weinberger, editors, Proceedings 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: https/////ccccooonsssooo mlr. toass/v48/mias html.
[23] Y. Miao, F. Grefenstette, and P. Blunsom. Discovering discrete latent topics with neural variational inference. In International conference on machine learning, pages 2410-2419. PMLR, 2017.
11

---

[24] S. Min, M. Lewis, H. Hajishirzi, and L. Zettlemoyer. Noisy channel language model prompting for few-shot text classification. In Proceedings of the 60th Anmual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 5336,,30, 2022.
[25] S. Min, M. Lewis, L.. Zeutlemoyer, and H. Hajishirzi. MetalCL: 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. 10.1866.. 18653/ naacl-maing mai.maainn...11 URL https: htt////aaaaanlloo/oo/o org/2022. naacl-main 201.
[26] 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.
[27] I. Mollas, Z. Chrysopoulou, S. Karlos, and G. Tsoumakas. Ethos: an online hate speech detection dataset, 2020.
[28] L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. L. Wainwright, P. Mishkin, C. Zhang, S. Agarwal,
K. Slama, A. Ray, al. Training language models to follow instructions with human feedback. arXiv preprint arXiv:2203.. 02550055,,,,,. 2022.
[29] E. Perez, D. Kiela, and K. Cho. True few-shot learning with language models. In A. Beygelzimer, Y. Dauphin, Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https: httpes:///peeereeeeeeeeee net. htt///orrmmmddmmhbee
[30] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. 2019.
[31] 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: wttp://ww...r/a////////////// httpg//b///19//// 10084.
[32] O. Rubin, J. Herzig, and Berant. Learning retrieve prompts for in-context learning. arXiv preprint
2021.
[33] 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 I. Linguistics. doi: 10. 1/15///11D1811400/// URL aclweb. org/ htt///nhhooooyyD-----aa
https://www. http://wwww.www..b../.bm
[34] N. Saunshi, S. Malladi, and S. Arora. A mathematical exploration of why language models help solve downstream tasks. In International Conference Learning Representations, 2021. URL
https://openreeve:: wwt.feeefffeoooee wtntt/forumiidee oidvV jWBBrdd= IW3sEc1s.
[35] R. Socher, A. Perelygin, Wu, J. Chuang, C. D. Manning, A. Ng, and C. Potts. Recursive deep models for semantic compositionality over sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pages 1631- 1642, Seattle, Washington, USA, Oct. 2013, Association for Computational Linguistics. URL https://aclaathoolggy org/ Dorg/D13--111 1170.
[36] N. Stiennon, Ouyang, 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- -0201, 2020.
[37] 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 N.xi:22000.0195.. 2022.
[38] 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: 13971, 2023.
[39] H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, Bhargava, 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, Kloumann, A. Korenev, P. Koura,
M.-. Lachaux, T. Lavril, J. Lee, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov,
12

---

P. Mishra, I. Molybog, Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, X. Kuan, P. Xu, Z. Yan, I. 7arov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, Edunov, and Scialom. Ilaama Open foundation and fine-tuned chat models. 7 2023. URL http: /tttp://arxee /atpivrxvvvvvv.///////ee//aaaava org/ haarg//b///aae 2b30 09288.
[40] L. van der Maaten and G. Hinton. Visualizing data using t-sne, Journal of Machine Learinng Research, ((86))2579-2605, 2008. URL http://jmlr org/ hatgg/apps/////e Saapers/v9/aadeeee aaaaeeemaattee0.... html.
[41] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, t. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[42] 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.
[43] A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. EMNLP 2018, page 353, 2018.
[44] B. Wang and A. Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. attps://githuue com http://oooomm/iiggoo/////////eee
May 2021. [45] 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./1177,, 2022.
[46] A. Warstadt, A. Singh, and S. R. Bowman. Neural network acceptability judgments. arXiv
preprint arXiv: 1805. 12471, 2018.
[47] 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.
[48] 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.
[49] Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, Cistac, T. Rault, R. Louf, M. Funtowicz, et al. Huggingface's 's transformers: State-of-the-art natural language processing. arXiv preprint 110007111, 2019.
[50] 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: //opnrreee//ereeeeeeeee wwt meet/EoruuniccccRdVFFAHH
[51] 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 languaae models, arXiv preprinl 2rXiv:2205.01063 2022.
[52] X. Zhang, J. Zhao, and Y. LeCun. Character-level convolutional networks for text classification. Advances in neural information processing systems, 28, 2015.
[53] D. Zhou, N. Schärli, L. Hou, J. Wei, N. Scales, X. Wang, D. Schuurmans, Bousquet, Q. Le, and E. Chi. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205. 10625, 2022.
13

---

A Proofs
A.1 Direct direction
Assumption A.1. (Assumption 2.1) Assume that (x) P(X). and P (Y x P(Y x))
for +
Proposition A.2. (Proposition 2.2) tatk follows the - - 0 direction, arg max PM (Y y|06, X) is the Bayes optimal classifier.
Proof. Since the data generation of the task can be written as Y f(X, E), we have
( P
And by Assumption A.I, we have
arg max Pdl (Y y004 X) arg max P(Y y10 X).
NEY
VEY
Thus arg yax PM (Y y10 X) is the Bayes optimal classifier.
Theorem A.3. (Theorem 2.3) If task d follows the - direction, then the in-context learning classifier
arg max Pd (Y = x/X Y X Yd x)
PM NEY
always has a higher or equal probability of misclassification to the Bayes optimal classifier arg max Pite (Y = -y00, X). Equality only takes when
Vz e X, Pd, (0 (0 Y Xx. Y X =x)==
Proof. Recall that in Equation (1), we have
(Y Y X yd X) / Por 0, (0 yp.... X Yg X)d0.
By Proposition A.2, arg yay P (Y y|0 X) is the Bayes optimal classifier. Let Cx(X)= =
arg max pd (Y 1/0, X). then the risk is defined the probability of misclassification
R(Co) (X) # Y) = Exx cx(X)
Denote the in-context learning classifier arg max Eyy Pd (Y y|X Y .... X Y x) by Cx(X). We then have
R(C+) = Exy [l = Ex
P (Y yl x))
yEY
Such risk minimized if and only Cx(X) = Cx,(X), which only holds when
Pd (64) [X Yd
Y x= x= for all x E X.
A.2 Channel direction
Assumption A.4. Assume that (M P(X), and (XIO,Y & P(X Y) for the Y X direction.
Proposition A.S. If task d follows the + - 0 causal direction, arg max max yey Pd (X y y) the Bayes optimal classifier when the label assignment is balanced.
Proof. Since the data generation of the task can be written as X "==((YY0' 04 E), we have
nd (d(X) = P(X Y)
14

---

When the label is balanced, i.e. Pd =
1 13jp we have
PY(Y[X)= (X(X))))'( x (X())
P(X)
And by Assumption A.4, we have
arg max PM (X = = arg max P(x 10 = 1)- yEy
NEY
Thus arg ya PM (X Y= =) arg maxyey Pd = xy is the Bayes optimal classifier.
Theorem A.6. If jask follows the Y X O direction, then the in-context learning classifier
M (XIV X X =
arg max
VEY
always has a higher or equal probability of misclassification to the Bayes optimal classifier arg yax pd (X Y v). Equality only takes when
Vy E E PM (0 X Y X y) =1
Proof. This theorem can be proved similarly as Theorem A.3. Recall that in Equation (2), we have
(X X - Ye. x Y) / PM ( (1
x YE x Y)dO. Y)dO
By Proposition A.5, arg maxyey P (X (X104 Y = y) is the Bayes optimal classifier. Let Co(X)= arg max pel O y), then the risk is defined as the probability misclassification
R(C,) (X) Y) F co())4l)
Denote the in-context learning classifier arg max Pd (X X1. .... Yd, X Y 9) v) Cx(X). We then have
R(Cx) Exy (A.())x))) = Ex
1
Pd PM Y = v))
yEy
Such risk minimized if and only if Cx(X) P (8dl)d Xf...., Y X X Y for all e.
Cos(X), which only holds when O
A.3 Method
Proposition A.7. (Proposition 3.1) When C(8d) is minimized, PA (Y ( x) P(Y .X) for X Y 0, and Pdl (X]0d Y) = P Rd Y) for X O. If the LLM M is invertible, then -d=00 0d.
Proof. The proof of this proposition is straightforward.
Since
C(O9) - H x)) KL(P(Y (XXIII (YI0s, x))
when C(0d) is minimized, we have Pd (Y x) P(Y x) for X Pdl (x Y) = P(X for Y X 0.
and
If M is invertible, since the embedding matrix is invertible with or without new concept tokens, pd (Y pd ( X) implies that Ô 0.. Thus O is identifiable, which means 9d gd. D
15

---

Table 2: Prompt template and label mapping for the datasets we use. Since almost all sentences from ETHOS contain offensive content, we mask out the key offensive words in the examples below.
Dataset
Prompt
Label Mapping
SST-2
senience: well worth revisiting as many times
negativefposiitve
positive
The company anticipates its tumoover for whole surpass that of the previous year when it was EUR 67.1 million. posilive
НPB
pegativooneutraalpooiivv
COLA
It this hat that know the boy who is wearing. unacceptable
acceptablelumcceptable
Athleto Athiltt/Buildingg Commpingg/Compnyyy/
The Noest River tribatary the Chiojdcanca River Romania. NaturalPlace
Eduuationallnsitutoo/lii
DBPedia
MeanO/frrnsportuiinn
PaturalPlacoOOttieeolll
Plant/ Village/ WritenWWoott
fast i mcan fastingis a way of skipping meals mona you move fast others
EmoC
angry/happyloohenssad
feel this place was tragic sadness
angen/Feurjoy//oww/
Emo'
sadnosalaurptise
[Makeed] should be removed from the face of the earth true
ETHOS-SO
false/true
hate being [Masked], wish I was [Masked]
ETHOS-R
and no [Masked] on earth existed
falise/true
false
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." The license for SST2, ETHOS-SS SO and ETHOS-R R is GNU General Public License v3. FPB is under a Creative Commons 3.0 Unported License. Note that these two datasets
Attribution-None SaenCoommmcillSaaaaeet 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-ShareAliin 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 large portion of the code from the MctaICL repository [25] The training takes around 20 to 40 hours on single GPU. We use a learning rate of le-4 and 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 LLMss 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.
16

---

datset d
urain size
lesi sizee
Lask set S
glue colafglue bue mali/gue qap
SST2 (glue-sst2)
16384
muuc-sst2/guummmliimaae otiioo/oaaiiiaatogoo/ss/////////
social L qu/wwino het/wiooggrandeeeee grande//gue--qup
IPB (financial_ phrasehank)
1811
453
ay newne/inaaaaaal phrasebunk/
poem connimentanii/qur//qquua
medical__qestionn questions_ pairs/ pwwssddcdidiiii.......
glue-cola/glae hatleeelilllggogggleeeeeeee glue qnli/glue rte/glue ssi2/glue wnli
COLA (colarsat2)
8551
1000
glue sstllluue mali/math ga/scia/ social i gatwino grande/glue-gapt ag sgg_mwss/inaacid phrasebank/ poem wwnniinen/anl//quueeeeeo medical, questions puirs pain/pawwwlbgelia. 14
DBpodia (dbpedia_14)
16384
1000
glue sst2/amazon, polarity/
EmoC (emo)
16384
1000
financial_ phrascbank/ppee httpo//n//oooeaaaaiiieieee
scatiment
dbppeiia Alethhos/em//mmmooii
glue sat2amazon, polarit//
fiaancial_ pbrascbank/poem scatiment/
EmoS (emotion)
16000
1000
yelp_ http///llllloolbliip//oo/ee
dlppedia, 14/ethon/emo/emmllii
glue sat2amazonn polarit//
financial phrasebank/ppeem sentiment/
ETHOS-SO schoo-scxuul_oriiaaatioo
346
87
yelp baapliyyglluocolla//iiiiee
dbpeedia_ settho//emo/mmottinn
glue sa12amazonn polarit//
financial phraschank/hoem sentiment/ yelp haw.attygguo--olalliiii........
ETHOS-R (ethos-religion)
346
87
dbpedia Wiethhowemmo/mmotion
Table Dataset details
Causal
Anti-causal
76.3
73.6
53.4
56.7
53
57.7
54.1
55.6
38.8
37.4
31.5 26.6
14.6
SST2
FPG
COLA
Dbpedia
ErncC
EmoS
ETHOS-SO ETHOS-R
Figure 6: Accuracy of randomly selected demonstrations averaged over seven different LLMs except for GPT3-davinci, using the adopted causal direction and the anti-causal direction.
k-ablation study results. The detailed results of i ablation study are shown in Table 10, correspond- ing 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 Pd (6d) pd yd) to the smallest.
M
c-ablation study results. The detailed results of 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 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 7d instead of a meaningful one. i.e., map each label 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 7 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 larger performance drop than only using random label projection map,
17

---

Uniform Ours
64.8
57.4
50.5
47.5
35.9
35.3
Original
Random words
Random labels
Figure 7: In-context learning accuracy of our method versus random selection baseline, with (a) ground truth labels (original), (b) random label mapping (random words), or random label assignments (random label), averaged over all eight datasets. Numbers are obtained with GPT2-large.
Optimal
90.3 86.2
86.1
92.6
86.2
76.6
69.1
60.4
56.5
57.3
48.4
38.6
SST2
FPB
COLA
Dbpedia
EmoC
EmoS
ETHOS-SO ETHOS-R
Figure Accuracy of in-context learning using our method versus the theoretical maximum accuracy obtained using the learned concept tokens prefixes. Numbers are obtained with GPT2-large.
which shows that the mapping between and 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 L..M parameters based on the appearance of X and Y. We also show that the demonstrations selected by 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 Table
Optimal performance As stated in Theorem 2.3, the optimal performance of an in in-context learning classifier the Bayes optimal classifier arg max Pd (Y = X)), which is approximated by using the learned concept tokens prefixes. Note that this approximated Bayes optimal classifier cannot be transferred across different L.Mss, as the learned concept tokens embeddings are aligned with specific L.M.. 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, it is the LLM that concept tokens are fine-tuned with. As shown Figure 8, our method comes close to the optimal accuracy many datasets, while there 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 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 Table 9.
Reordering results. We reorder the selected demonstrations to maximize the posterior of the concept tokens:
arg max Pd (1000//(((X))))) VP) (Xg. Y())
nell
Where ((Xf, YJ), (Xg. y()) is permutation of (X Y (Xg. x Yg). II is the set of all possible permutations of the demonstrations. The detailed results with and without reordering are shown in Table 12, corresponding to Figure
Similar tokens. We show the top ten similar tokens some learned concept tokens in Table 13, as summarized in Figure 5 in the main text.
18

---

Rearder
Not reorder
90
80
70
50
40
30
20
10
SST2
FPB
COLA
Dbpedia
EmoC
Emos
ETHOS-SO ETHOS-R
Figure 9: In-context learning accuracy our method versus random selection baseline, with and without reordering. The red error bars represent the standard deviation across five runs. Numbers obtained with GPT2-large.
Table 4: Accuracy of selected demonstration. Our demonstrations are selected using GPT2-large, and the same set of demonstrations applied to all different LLMss. All L.LMs are pre-trained only with the language modeling objective, while the pre-training data size of GPT2s is much smaller than GPT3s.
LLM
Method
SST2
FPB
COLA
Dippedia
EmaoC
Emoo'
ETHOS SO
ETHOS R
Avg
GPT2 (124M)
Uniform
69.7 1:
$2.9 d: 2.3
61.9 1.4
48.0 1
35.3 4:
26.4 1.0
64.1 4:88
710 188
53.7
Samilar
69.5 0.6
55.9 +
63.2 122
44.7 +
36.4 + 20
26.6
77.7 2:27
80.0 3.7
Ours
76.8 2:29
69.1
53.5
37.2
30.6
80.9
76.8
GPT2-m (355M)
Uniform
70.8 1.3
$2.0 1.7
$7.8 1.3
49.3 2.0
34.2 + 18
34.2 + 1.8
76.3 + 4.9
74.7 2:2
56.2
Similar
37.2
Ours
$1.2.113
4.3
69.0
0.02
$2,9 23
40.4 21.5
37.2 17.2224
83.7 1.7711
76.8
62.6
GPT2-1 (774M)
Uniform
77.1 1.2
2.
62.7 0:
54.4
38.7
34.5 1.2
67.6 1..1..4
72.9 1:28
Similar
80.7
54.8 3.38
50,9 1.4
51.1 +
39.9 4
35.1 2:2.1
80,9 28
84.4 2:26
59.7
Ours
86.2 +:
60.4 4:2555
69.1 1 0.
56.5 +
48.4 17.0
38.6 2.28
82.5 1.5
76.6 122
64.8
GPT2-xl (1.5B)
Uniform
Similar
80.6 1.3
53.0 2.5
$5.0 2.
$1.6 = 5.9
39.9 + 2.0
32.9 2:2.1
82.8 2.X4-22
83.9 +
60
164
GPT3-a (350M)
Uniform
76.9 0.7
56.6 1.1
$3.1 18
62.1 =
38.6
27.7 1.
65.5 5.7
74.0 3.0
56.8
Similar
78.7 1: 1.0
52.2
53.1 1.8
54.6 +
42.4
37.2 1.1
84.1
87.8 1:35
Ours
85.4
61.9
58.2 7.70
64.0
43.0
37.9 4:
84.4
78.9
64.2
GPT3-b
Uniform
80.8 06
55.2 3.3
46.8 2.
66.5 = 1/4
42.0 0.
27.0 + 1.2
71.0 4: 46
72.6 3.1
57.7
Similar
Ours
87.3 2.
64.3
67.2 09
70.2 = 3.2
43.6 + 13.0
38.9 5.0
8A.6 +
78.9 1.2
66.9
1.8
59.1
77.5
(6.7B)
Similar
85.7 1:7::0
(2.2::000
$8.0 1:00:00 1.7
62.2 2:
47.4 1.44433
39.8 1.7
89.2
89.7 1.99
66.8
Ours
88.8
64.1
69.0 0:03
73.6
50.3 11.9
43.1
86.2 0.0
78.2 0:00
GPT3- (175B)
Uniform
86.5
59.2
45.5
73.6
39.4 +
40.6
76.8
62.4
Similar
88.5 = 0.8
$5.4 9.3
45.4 15.411
67.2 =
37.6
39.8
86.9 2.4
89.0 89.0 3.8
63.7
Ours
Avg
Uniform
77.6
54.1
55.3
59.7
38.8
32.6
70.9
74.0
$7.9
Similar
55.9
54.9
40.5
84,4
85.7
Ours
84.6
62.4
66.1
63.1
43.5
39.1
83.8
77.9
65.0
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 dataset.
Selected demonstrations. Table 14 shows the selected top 4 demonstration by our proposed algorithm.
Task
Selected demonstrations
19

---

GSM8K
Question: It takes Bryan 5 minutes to walk from his house the bus station. Then rides the bus for 20 minutes. After that, he walks 5 minutes from the bus station to his job. takes the same amount of time
the morning and the evening. How many hours per year does Bryan spend traveling to and from work, if he works every day? Bryan spends 5+20+5 =<5+20+5=30x30 minutes traveling to work. He does this twice a day, so he spends 3042=<30x2=60x60 minutes traveling per day. Thus, he spends 60/60=<60/60-111 hour traveling to and from work every day. Thus, he spends hour+365 days =c14365-3653365 hours traveling per year. The answer is: 365
GSM8K
Question: Cherry put up a delivery service. She charges $2.50 for a 3-5 kilograms cargo and $4 for 6-8 kilograms cargo. If she delivers four 5 kilograms cargo and two kilograms cargo per day, how much money will she earn in a week? Cherry earns $2.50 x 4- 5x2.5x4=00>10 after delivering four 5 kilograms of cargo per day. She earns $4 2= $<4+2=8>8 after delivering two 8 kilograms of cargo per day. So, her total earning per day is $8+$10= $x8+10=18>18. Therefore, she will earn $18 7= $x1847=126x126 in week. The answer is: 126
GSM8K
Question: Bill is laying power cable for a new neighborhood. There are going to 18 east-west streets that are miles long and 10 north-south streets that are four miles long. takes 5 miles of cable to electrify mile of street.
cable costs $2000/mile, what is the total cost of cable for the neighborhood? First find the total distance of the east-west streets: 18 streets * 2 miles/street <18*2=36x36 miles. Then find the total distance of the north-south streets: 10 streets * miles/street = <10+4=40>40 miles. Then add the number of miles from each type of street to find the total distance: 36 miles + 40 miles = <36+40=76>76 miles. Then multiply that number by 5 to find the number of miles of cable needed: 76 miles street 5 miles cable/mile street <76*5=380x380 miles of cable. Then multiply that number by the cost of one mile of cable to find the total cost: 380 miles $2000/mile
$<380+200--7600000566, 000. The answer is: 760000
GSM8K
decides to replace the video card in it. He sells the old card for $300 and buys a new one for $500. How much money did he spend on his computer, counting the savings from selling the old card? He spent an extra 500-30-8550-30-0002200 the video card. That means the total cost was 1200+00-$41100+00001000000000
Question: John buys a gaming PC for $1200.
The answer is: 1400
SST2
sentence: faced and spindly attempt at playing an ingenue makes her nomination as best actress even more of positive
SST2
sentence: holofcener's film offers just enough insight to keep it from being simpleminded, and positive
SST2
i'm not a fan of the phrase life affirming' because it usually means schmaltzy," but real women have curves truly is life affirming negative
sentence:
20

---

SST2
sentence: the script is about as interesting as recording conversations at the wal-mart checkout line negative
DBpedia
OfficeHolder Lucie Papin (born September 7 1936) is a former Canadian politician who served in both the House of Commons and Senate.
DBpedia
Village Kunkalamarru is very renowned village under Karamchedu Mandal which is located about 15 km from the busy commercial town of Chirala in Prakasam district in the state Andhra Pradesh India. neighbouring villages are Karamchedu Veerannapalem.
DBpedia
EducatioonalIIstitttaaa The Pontifical Catholic University of Puerto Rico at Mayagez is a university located in the city of Mayagez Puerto Rico. It is part of the Pontifical Catholic University of Puerto Rico. The university began an extension the Catholic University of Puerto Rico in the early 1960s. In 1982 it was awarded the official title of Center and later became the Mayagez Campus of the Pontifical Catholic University of Puerto Rico at Mayagez 1996.
DBpedia
Artist Choi Dong-wook [citation needed]; born November 1984) better known by his stage name Se7en is South Korean singer from YG Entertainment. He has also advanced into Japan China and the United States.
Table 14: Selected demonstrations by our method.
C Limitations and Future Work
While the assumption that 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 significant portion of the long-tail 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 currently left to the user's discretion. A better approach to constructing such a task set is needed to gain deeper understanding of latent concept variables and 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 (LLMss 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 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 potential for misleading output and to regulate the distribution of L.MMs in order to prevent any negative societal impact. Additionally, it is possible that L.LMs could be intentionally misused, thus important to consider the ethical implications of their use and to take appropriate measures to mitigate
21

---

Table 5: Accuracy of selected demonstration. Our demonstrations are selected using GPT2-large, and the same set of demonstrations is applied to all different LLMs. All LLMs are pre-trained only with the language modeling objective, while the pre-training data size of GPT2s is much smaller than Gs.3s.
LLM
Method
SST2
COLA
DBpedia
EmoC
EmoS
ETHOS-SO
ETHOS-R
Avg
Uniform
+ 1.8
61.9
+ 1.7
448
71.0
(124M)
Random
69.8 +
51.1
69.0 0.1
49.0
33.7 155
24.2
66.4 +
66.2 162
53.7
Ours
76.8 + 2.9
64.5 1.2
69.1 0.22
53.5 2 2.95
37.2 37.2 11.1
30.6 488
80.9 =
76.8 4
61.2
GPT2-1
Uniform
77.1
51.3
0.
54.4
38.7
34,5
67.6.:33
72.9
57.4
(774M)
Random
81.9 4.5
46.5 + 4.7
64.9 7.8
50.3 + 4,3
42.5 à: 16.7
36.1 6.5
67.6 +: 20.4
67.8 15.0
57.2
Table 6: We test our method on other similar sizes (6-7B) LLMS..
LLM GPT2-1
Method
SST2
FPB
COLA
Dilpedia
EmoC
EmoS
ETHOOS SO
ETHOS R
Avg
Random
77.1 1.2
51.3 24
62.7 1 088
54.4 + 09
38.7 2.1
34.5 +
67.6
72,9 2.8
57.4
Ours
86.2
60.4 1:25
6..1::0.
48.4
38.6
76.6
64.8
GPT3-c
Random
84.2 1.4
52.6 t: 1.8
59.1 1
70.6 + 0.8
44.3 + 2.5
32.3 +.3410 19
47.5
17.5 0.6
62.3
78.2
GPTJ
Random
78.5 = 1.0
53.1 + 1.7
58.3
55.6 = 1.2
18.5 2.0
31.3
76.6
76.6 1666110 1.4
58.8
1.9
56.7 4.3
69.1 0.2
+
32.5 16.1
85.3
17.0
OPT
Random
32.8 0.
34.8
29.4
67.1
36.9
86
78.2
54.7
Ours
74.2 74.2
34.1 1.1
35.7
28.8 2.1
76.7 4.1
39.0
8600
78.2 78.2.200..
56.6
LLaMA
Random
57.7
3:
5.58
$7.2
36.4
Ours
60.3 4 4.7
19.1 + 1.9
30.8 + 0.2
16.9 + 1.3
4.3 = 0.7
35.3 0.66
77.2 13.6
56.3 + 10.8
37.6
any potential negative effects. We posit that these regulations and measures should be put place the time of distributing LLMs to ensure the safe and responsible use of these models. Furthermore, publicly release code, will also provide clear warnings and guidelines users ensure that the potential risks associated with the use of our method are fully understood and addressed.
22

---

Table 7: We test random selection baseline with anti-causal direction.
LIM
SST2
FPB
COLA
DBpedia
EmoC
Emo8
ETHOS so
ETHOS R
GPT2
57.4 +
$6.6 + 2.1
55.9 + 1.
11.3 + LB
24.6 24
22.1
1.1
64.1 4.88
58.6 5.5
GPT2.m
56.7 +
48.7 =
55.3
13,9 +
22.4 19
24.9
+ 2.3
44.8 199
45.5 3.
GPT2-1
58.7 0.77
33.7 1.3
50.8 =
136
28.2 16
26.2 27
48.7
336255
GPT2-xl
0.55
50.
1.
314
2.
65.3
61.8
GPT3-a
55.8 09
58.9
11
14
54
27.7 17.7-11
49.2
54.9 -0
GPT3-b
644
58.9 26
534
14.6
52.0
27.0 1.3
48
51.0
GPT3-c
78.2 1:
52.3 1:
53.7 1:
23.0 000::2
49.1 19.1::2
32.2 1.9
57.9 2:7
64.1 #.5.0
Avg
60.8
50.8
14.8
37.4
26.6
54.1
55.6
Table 8:
We test our method with random words and random labels using GPT2-large.
Method
SST2
FPB
COLA
DBpodia
EmoC
EmoS
ETHOS-SO
ETHOS-R
Avg
R words
Random
54.1 422
43.4 1.9
62.2 4.9
11.2 + 09
32.4 5.2
19.1 = 18
80.7 = 488
77.0 16
47.5
R labels
Random
51.5 0.9
32.5 + 1.2
49.3 = 3.
6.7 1.0
25.1 206
17.2
48.0 2.
$6.8
35.9
0.9
6.6
006
51.0
Table 9: Accuracy using concept tokens prefixes.
SST2 90.3 + 0.0
FPB
COLA
DRpedia
EmaC.
EmoS
ETHOS-SO
ETHOS.R 78.2 # 0.0
86.1 + 9.0
75.0 # 0.1
92.6 + as
57.3 188
53.8 + 0.7
86.2 t.00
Table 10: l: ablation study using GPre-large, without reordering.
Method
SST2
FPB
COLA
DBpedia
EmaC
Emo8
ETHOS-SO
ETHOS-R
Avg
À
Random
74.4 + 1.D
48.5
48.9 + 166
$2.9 = 2.0
42.8 + 0.6
37.1 + 1.2
66.9 4.7
66.4 =
54.7
78.1 4.5
50.1 2.9
41.1 998
2.6
84.6
Random
76.9
56.6 +
$3.1 1.8
62.1 =
1.
27.7
65.5
74.0
56.8
Ours
86.2 1.22111 1.4
59.7 2.
60.1
56.5 3.2
38.2 218
37.7 225
83.0 13
766612
63.4
Random
37.1
1.
66.5
15
36
68.5133
72.9
58.8
Ours
87.0 2.4
39.9 33
95.3
67.0 + 0.9
39,9 4:3
38.8 4: 2.6
77.0 11.1
78.9 $.99:00 0.9
63
33.72222
Ours
14.6
60.4 6.64
62.0
71.0 1.9
37.2
37.1 22.
7.44766
74.7 4.7
62.4
Table 11: ablation study using GPT2-large
SST2
FPB
COLA
DBpedia
EimoC
EmoS
ETHOS SO
ETHOS
Avg
34.3
62.9
62.1
85.4 1:
61.9 10.5
58.2 7.0
64.0 4.44
43.0 1:
37.9 2:23
84.4
78.9 0.9
64.2
80.1
64.3 2:77
63.1 9.8
58.7 4:32
36.4 11.5
38.6
80.9 3.9
76.3 5.9
62.3
20
78.5
80
6..5
58.0
36.3
41.8
80.7
73.8 5.4
60.92
Table 12: Reorder versus not reorder using our method, with GPT2-large.
SST2
FPB
COLA
DBpeeiaa
EmaCC
EmsoS
ETHOS-SE So
ETHOS R
Avg
reorder
86.2 1.4
60.4 2: 2.5
69.1 1 0.2
=
48.4 1 17.0
38.6 1:
82.5
76.6 =
64.8
rearder
86.2
59.7
69.1 4
56.5
38.2 + 21.8
37.7
83.0 +
76.6
63.4
23

---

Table 13:
We list the top 10 similar words (tokens) to some of the learned concept tokens.
concept
similar words
FPB-2
milo coondinate notify rendering beneliting routing Eniity liem routed Messages Plot
FPB.3
unlocked updating deleting dropping damage updates drops Gained taken dropped FX Safari Fixes advertisers I Links Coins Operator marketers Guidelines
FPB-4
FPB-5
674 592 693 696 696-498 593 793 504 691 683 fragmented clong ieerator initialized
COLA-2
Sp spa contributed cerv borrower paper liger Erica USII Schwartz Democrats
demacr
WH ophohic
Democrats
DBpedia-4
olten impede blockade incarcerated LEASE pollutants pesticides uphe lawmakers fosails
DBpedia-5
categorized closes therapies antidepressant refrospective elinically physicians therapists randomized ciselllinn clinicians
DBpedia-7
IS provided Killed richness Compet Nevertheless Probably Proceedings horizontally
ETHOS-S0-3
Revolution Spread itu Million Pascal stabil Indy Georgian Figure resy
ETHOS-R-2
council Chocobo uyomi aditional cumbers
Thubbnillmaee seems outlines emitted grin outline circuitry sized flips emits flipped
Pharooh
ETHOS R.
ETHOS-R-9
Cynus Sith Scorpion Snape
Morty
EmoC 6
behavi checkpoints unintention crih eleph louph mosquit blat piome
EmoC-8
depressed ballied choked stricken devastated unsuccesffll cheated distraught troubled failing
EmoS-1
frightened rebellious depressed eareless bullied restless reluctant distraught clumsy disgruntled obsessive crappy demonic delusions psychosis psychotic childish stupddiy reckless insanity
EmoS-5
benevolent charismatic perfected
unintention pione innocuvux learless glamonous ruthless whispers pundits Sadly horribly curiously noticeably Sadly gaping painfully shockingly
Fmo8-9
24

---

Histogram of 5572
Histogrom
9.24
0,200
(a) SST2
(b) FBP
Histegram of COLA
Histogram DBpedia
0.02
0.10
0.1
A12
6.3
011
635
616
6415
0.003
DADS
0.000
0.25
0.10
(c) COLA
(d) DBpedia
Histogram of EmoC
Histogram tmos
Panentogg
0.0000
0.5020
0.0005
2.0355
(c) EmoC
(f) EmoS
Histogram
Histogram
5:0c
5.000
(g) ETHOS-SO
(h) RTHOS-R
Figure 10: Historgrams the probability of train examples in each dataset predicting corresponding concept tokens.
25


Reviews:
[{"review_id": "pJhalR2alm", "reviewer": "Reviewer_waNU", "summary": "This 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.", "questions": "Please 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?", "limitations": "N/A", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "The 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. Below is our response to your comments:\n\n1. **Takeaway of our paper**: We would like to clarify that our study's focus is on elucidating the underlying mechanisms of in-context learning and presenting a novel approach to demonstration selection that empirically verifies our theory in a real-world setting. By employing prompt tuning with labeled data, we can tease apart latent variables and relationships that otherwise might be obscured. \n\n2. **Real-world use case and generation task**: In short, our proposed method is most useful with hard tasks that smaller LLMs perform significantly worse than larger models. Please refer to point 2 of our general response for detailed results on GSM8K. In this case, even parameter-efficient fine-tuning with smaller models cannot obtain reasonable performance (less than 4% accuracy), while in-context learning combined with our method can obtain significantly higher performance with both small and large models. Our method can obtain more than 80% accuracy with ChatGPT, compared with 76.5% accuracy using random selection. Note that our method only requires tuning a small model on limited data, so it is data and computation efficient. We want to reiterate that improving and connecting a previously synthetic-only theory of in-context learning [6] with real-world models and data is already non-trivial and is relatively rare in the current theories of LLMs, which are usually disconnected from the real world. The value of our paper does not only lie in its real-world application.\n\n3. **Baselines**: To the best of our knowledge, none of the existing demonstration selection methods involves tuning a smaller model. So we are not able to involve a similar baseline. About the proposed contribution-based baseline, we are not quite sure what the reviewer means by ‘select ICL examples based on their contribution to the predictions of the labeled task data’. Do you mean select demonstrations by attributing the prompt-tuned model to the labeled training data? If that is the case, we agree this is an interesting idea, which to the best of our knowledge, has not been used on any existing demonstration selection method yet. However, we want to argue that this baseline is too exploratory, which could be a new project by itself, and does not seem to be natural enough to be incorporated in the current paper. We think this is out of the scope of our project, but we are happy to discuss it with the reviewer if they can clarify what they mean here. \n\n4. **Compared to prompt tuning**: For simpler tasks like sentiment analysis, prompt tuning performs better than in-context learning, even when combined with our method (detailed results see Figure 8 in the Appendix). For harder tasks like GSM8K, simple prompt tuning with small models cannot obtain meaningful results (less than 4% accuracy), while in-context learning can obtain significantly higher performance with both large models (more than 80% accuracy with ChatGPT combined with our method) and small models (19.3% accuracy with Llama2 7B combined with our method). For detailed results, please refer to point 2 of our general response.\n\n5. **Analysis of the selected demonstrations**: We didn’t include an analysis of the selected demonstration examples in the paper because the common features shared between the selected examples are a bit hard to detect. However, we agree that it is still necessary to include such an analysis. Because of the space limitation, we only list the top demonstrations of GSM8K and SST2 in the supplemental 1-page pdf.\n\nCompared 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 harder examples can represent the task more comprehensively. This conclusion also aligns with the finds in [1] 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. We will add the analysis of all selected demonstrations in the revision.\n\n[1] Han, Xiaochuang, et al. Understanding In-Context Learning via Supportive Pretraining Data. ACL 2023."}, {"review_id": "Kosa7jXhdV", "reviewer": "Reviewer_XxTq", "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.", "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", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Below is our response to your comments:\n\n1. **About our claims**: We would like to clarify that we are not claiming that LLMs are topic models. Our claim is that LLMs are IMPLICITLY topic models, that learn and infer a latent task/topic variable for solving each query. We would be very happy to revise and clarify this.\n\nThere are many plausible ways of understanding in-context learning. The topic model/latent variable view is indeed just one important piece of it. We in no way claim ours is the only correct one. We would like to clarify that many of these are complimentary views of LLMs that do not necessarily conflict with either other. For example, the in-context learning as gradient descent view [6,7,8] can be understood as a specific way of inferring and utilizing the task latent at in-context learning time, as in the Bayesian interpretation framework [9,10,11].\n\n2. **About conflicting findings**: The pattern information of a task can be viewed as also included in the inferred latent variable, so our explanation is not conflicting [1]. There is a similar finding in [3] as the mentioned random tokens finding in [2], saying that the demonstrations with random labels also work. In contrast, [4] explicitly argues that ground truth labels matter, while the sensitivity to label correctness varies across different tasks. Interestingly, [3] and [4] are published at the same conference. The experiment in Table 7 of our appendix shows that the correct label does matter on our tasks, with or without our proposed demonstration selection method, which agrees with [4]. As the results in [2] are mostly from the NER task, it is likely that NER has a high tolerance for randomness in the demonstrations.\n\n3. **About baselines**: We are aware of the mentioned paper [5] and have cited it in our related work section. We did not include it as a baseline because its setting is too different from ours, which makes it incomparable to our method. They propose a two-phase algorithm, where the first phase selectively annotates a relatively small number (18 to 100) of data from a large set of unlabeled data (3K), and then the second phase retrieves demonstrations from the annotated data based on their similarity to the query. We are, on the other hand, exclusively focusing on selecting demonstrations from a small set of labeled data (100). So our proposed method is only comparable with the second phase of their algorithm, which is the same as the ‘Similar’ baseline (defined in line 238) used in our experiments. It is possible to combine our method with the first phase of [5] to obtain a better performance under the selective data annotation setting with a large set of unannotated data, but this setting is out of the scope of our current paper, so we leave this to the future work.\n\n[1] Olsson, Catherine, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895 (2022).\n\n[2] Hongxin Zhang, Yanzhe Zhang, Ruiyi Zhang and Diyi Yang. Robustness of Demonstration-based Learning Under Limited Data Scenario. EMNLP 2022\n\n[3] Min, Sewon, et al. Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?. EMNLP 2022.\n\n[4] Yoo, Kang Min, et al. Ground-Truth Labels Matter: A Deeper Look into Input-Label Demonstrations. EMNLP 2022.\n\n[5] Su, Hongjin, et al. Selective annotation makes language models better few-shot learners. ICLR 2022.\n\n[6] E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. NeurIPS 2022.\n\n[7] 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. ACL 2023.\n\n[8] J. von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. ICML 2023.\n\n[9] M. Hahn and N. Goyal. A theory of emergent in-context learning as implicit structure induction. arXiv preprint, 2023.\n\n[10] H. Jiang. A latent space theory for emergent abilities in large language models. arXiv preprint 2023.\n\n[11] S. M. Xie, A. Raghunathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference. ICLR 2022."}, {"review_id": "xR0X64w6MD", "reviewer": "Reviewer_sfNW", "summary": "This 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.", "questions": "See weakness", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. Simple formulation that aids better selection of in-context demonstrations\n2. Experiments have been presented on multiple NLP tasks", "weaknesses": "Modeling: 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Below is our response to your comments:\n\n1. **About topic model assumption**: We want to clarify that the topic model here is in a more general sense, which is not equivalent to LDA [1], but similar to the modern neural topic models proposed in [2,3]. In this more general definition of the topic model, the tokens/words are not required to be conditionally independent given the topic variable. We are aware that this definition of a topic model is basically a simple latent variable model of language with a single latent. Thanks for pointing this out, to avoid further confusion, we decided to change our title to Large Language Models as Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning. We will also revise our paper to make this point clearer.\n\n2. **About Similar baseline**: We define the \"Similar\" baseline at line 238,  which means using the most similar examples to the current query as the demonstrations for in-context learning, and it was first proposed in [4]. Since it is relatively straightforward and is a default demonstration selection method that has been used in many papers, we did not include much analysis of it in the main paper. Thanks for pointing this out. In general, our experiment results show that using demonstrations similar to the testing query will improve the in-context performance compared to random selection, which implies that similar examples contain useful information for LLMs to infer about some parts of the task latent that is relevant to the current query. However, it is still not as effective as directly selecting the examples that can best infer the whole task latent, as used in our proposed method, which can cover more aspects of the task. We will add this analysis to the revision.\n\n3. **About the main results**: Figure 2 in the main paper corresponds to the last column of Table 3 in the appendix. You probably looked at the last three rows (which correspond to the average over all the models) instead of the last column (which corresponds to the average over all the datasets). We will revise the table to make it clearer.\n\n[1] David M. Blei, Andrew Y. Ng, and Michael I. Jordan. Latent dirichlet allocation. J. Mach. Learn. Res. 2003.\n\n[2] Miao, Yishu, Edward Grefenstette, and Phil Blunsom. Discovering discrete latent topics with neural variational inference. ICML 2017.\n\n[3] Miao, Y., Yu, L. & Blunsom, P. Neural Variational Inference for Text Processing. ICML 2016.\n\n[4] 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\n"}, {"review_id": "ELnRwsmISi", "reviewer": "Reviewer_8JqR", "summary": "The 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.", "questions": "- 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", "limitations": "Limitations are stated in the appendix C.\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This 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.", "weaknesses": "An 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review. Below is our response to your comments:\n\n1. **About optimal latent**: Yes, in practice, the learned latent is not (even never) optimal, but approximate. The assumed optimality is for deriving the upper bound theory that the in-context learning classifier can be as good as the Bayes optimal classifier. And the manual prompt can also be viewed as an approximate of the optimal latent. However, by using a manual prompt, we are not able to prove our theory of in-context learning that LLMs infer a latent variable at the inference time. The prompt tuning approach allows us to reveal the unobservable latent task variable in LLMs.\n\n2. **About reordering**: The reordering step in the proposed algorithm is indeed redundant. It’s a leftover from the original algorithm design before we perform the ablation study. We will remove this step in the paper.\n\n3. **About better model for demonstration selection**: In point 2 of our general response, we show that a better/larger model can select better demonstrations with the GSM8K dataset. More specifically, our method with Llama 2 (7B) outperforms our method with GPT2-XL (1.3B) by 0.8-3.4% absolute accuracy when performing in-context learning with different models.\n"}, {"review_id": "dDkTJY4mMs", "reviewer": "Reviewer_uvg2", "summary": "The 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. ", "questions": "* 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.", "limitations": "Please see the Weaknesses section.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* 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.", "weaknesses": "* 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Below is our response to your comments:\n\n1. **About how realistic is the problem setting**: We want first to reiterate that our problem setting is chosen to empirically prove the correctness of our proposed theory of in-context learning. We want to reiterate that improving and connecting a previously synthetic-only theory of in-context learning [6] with real-world models and data is already non-trivial and is relatively rare in the current theories of LLMs, which are usually disconnected from the real world. The value of our paper does not only lie in its real-world application. However, we do agree that we need to show real-world use cases of our proposed algorithm. In short, our proposed method is most useful with hard tasks that smaller LLMs perform significantly worse than larger models. Please refer to point 2 of our general response for results on GSM8K. In this case, even parameter-efficient fine-tuning with smaller models cannot obtain a reasonable performance (less than 4%), but in-context learning with both small and large models combined with our method can achieve much higher performance (19.3% using Llama2 7B and more than 80% using ChatGPT).\n\n2. **About evaluation tasks**: Since our primary goal is to connect the theory with real-world models and datasets, we did not try to include harder tasks. However, to show that our method also works on more challenging tasks, we show new results on GSM8K in point 2 of our general response, which involved chain-of-thoughts reasoning and generation, and shows the 4-7.9% increase of using our method with Llama2 (7B) compared with the Unifrom baseline. We want to reiterate that improving and connecting a previously synthetic-only theory of in-context learning [6] with real-world models and data is already non-trivial and is relatively rare in the current theories of LLMs, which are usually disconnected from the real world. The value of our paper does not only lie in its real-world application.\n\n3. **About $\\epsilon$**: It represents a noise variable, with any zero-mean distribution. Without this variable, there is no randomness between $X$, $Y$, and $\\theta$, as the functions $f$ and $g$ are all deterministic. The $\\epsilon$ here is just to introduce some noises such that the conditional distribution $P(Y|X, \\theta)$ is meaningful. Thanks for pointing this out. We will clarify this in the revision.\n\n4. **About instruction-tuned models**: In point 2 of our general response, we include results with ChatGPT on GSM8K, which shows our proposed method also works on instruction-tuned models (Accuracy of random selection baseline: 76.5%; Ours: 81.2%).\n\n5. **About Llama's performance**: The bad performance of Llama is also surprising to us. We wrote an email to the authors of Llama to ask about this, and they replied that they had never tested under such a scenario, so they don’t know. The code for testing Llama is completely the same as for testing other models, though. Our understanding is that Llama (first generation) is good at generation tasks, but not so good at in-context learning with simple classification tasks."}]

Author rebuttal:
We want to first thank all the reviewers for taking the time to review our paper. We provide some clarification and additional results below:

1. **About topic modeling assumptions**: We want to clarify that the topic model here is in a more general sense, which is not equivalent to LDA [1], but similar to the modern neural topic models proposed in [2,3]. In this definition, the tokens/words are not required to be conditionally independent given the topic variable. We know that this topic model definition is essentially a simple latent variable model of language with a single latent. To avoid further confusion, we will change our title to **Large Language Models Can be Viewed as Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning**. We do not intend to claim that this latent variable explanation of in-context learning is the only correct one. As described in the related work section, multiple plausible and complementary ways exist to understand and interpret in-context learning. 

2. **About the realisticness of our setting**: Since our primary goal is to connect the theory with real-world models and datasets, we did not try to include harder tasks. 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 improve the performance under a low data setting, with a small computing budget, and minimal inference latency, our demonstration selection method is a reasonable choice. Our demonstration selection method can also potentially be combined with other prompting techniques to boost performance further. 

**We added a new dataset, GSM8K** [4], which is a math word problem-solving dataset with chain-of-thoughts solutions. The table below shows the test accuracy of the final numerical answer using greedy generation. Note that we did not use a calculator to insert the correct result of each generated math equation during generation for time efficiency, which resulted in slightly lower scores.

As shown in the first row of the table, prompt tuning with ten new tokens can only obtain less than 4% accuracy on the GSM8K test set. While it is possible that a better-designed efficient tuning algorithm with more trainable parameters will be able to get better results, it is still unlikely that the efficient tuning results with small models will be able to outperform the few-shot performance of larger/better models,  as [8] show that a fully fine-tuned Llama 7B model on the whole GSM8K training set can only get 35.9% accuracy on GSM8K, and can be improved to 49.3% if combined with data augmentation, which is still significantly lower than the over 80% accuracy obtained by ChatGPT combined with our method. 

The last 4 rows show the in-context learning results with different size Llama 2 models [5] and ChatGPT. Our proposed demonstration selection method (last two columns) significantly outperformed the Uniform and Similar baseline as defined in line 236 to line 247 in our paper. Also, note that the demonstrations selected with a larger model (7B) are more effective than those selected with a smaller model (1.5B).

|  | Uniform | Similar | Ours w/ Llama 2 (7B) | Ours w/ GPT2-XL (1.5B) |
| --- | --- | --- | --- | --- |
| Prompt tuning | N/A | N/A | 3.7 | 1.3 |
| Llama 2 (7B) / 4-shot | 11.4 | 13.1 | **19.3** | 15.9 |
| Llama 2 (13B) / 4-shot | 17 | 18.3 | **21.6** | 20.5 |
| Llama 2 (70B) / 4-shot | 50.2 | 53.5 | **54.3** | 52.9 |
| ChatGPT (gpt-3.5-turbo) / 4-shot | 76.5 | 78.1 | **81.2** | 80.4 |

*We want to reiterate that improving and connecting a previously synthetic-only theory of in-context learning [6] with real-world models and data is already non-trivial and is relatively rare in the current theories of LLMs, which are usually disconnected from the real world. The value of our paper does not only lie in its real-world application.*

3. **Analysis of the selected demonstrations**: We didn’t include an analysis of the selected demonstration examples in the paper because the common features shared between the selected examples are a bit hard to detect. However, we agree that it is still necessary to include such an analysis. Because of the space limitation, we only list the top demonstrations of GSM8K and SST2 in the supplemental 1-page pdf.

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 harder examples can represent the task more comprehensively. This conclusion also aligns with the finds in [7] 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. 

[1] David M. Blei, Andrew Y. Ng, and Michael I. Jordan. Latent dirichlet allocation. J. Mach. Learn. Res. 2003.

[2] Miao, Yishu, Edward Grefenstette, and Phil Blunsom. Discovering discrete latent topics with neural variational inference. ICML 2017.

[3] Miao, Y., Yu, L. & Blunsom, P. Neural Variational Inference for Text Processing. ICML 2016.

[4] Cobbe, Karl, et al. Training verifiers to solve math word problems. arXiv preprint 2021.

[5] Touvron, Hugo, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint 2023.

[6] Xie, Sang Michael, et al. An Explanation of In-context Learning as Implicit Bayesian Inference. ICLR 2021.

[7] Han, Xiaochuang, et al. Understanding In-Context Learning via Supportive Pretraining Data. ACL 2023.

[8] Yuan, Zheng, et al. Scaling Relationship on Learning Mathematical Reasoning with Large Language Models. arXiv preprint 2023.

Meta-review: This paper explores the in-context learning capability of large language models (LLMs) through a Bayesian lens, treating LLMs as implicit topic models. The authors propose a two-stage algorithm that first learns the latent concepts using prompt tuning and then selects optimal demonstrations from labeled data with a smaller language model. The method yields improvements compared to random selection, across several models and datasets.

The reviewers are generally positive about this work. They praise that the paper is clear and well-written, the proposed method shows improvements over baselines, and experiments performed on multiple datasets. 
However, there are some concerns regarding the manuscript:
Some reviewers felt the paper lacks a complete analysis, either of the algorithm's performance (waNU) or of the experiment baselines (sfNW). XxTq, sfNW, and 8JqR raise concerns about the methodological and theoretical assumptions made in the paper.
XxTq pointed out conflicts with previous research. sfNW questioned the bag-of-words assumption. 8JqR expressed concerns over the "optimal task latent" assumption.Reviewer uvg2 raised multiple issues including the problem setting being unrealistic due to the amount of training data assumed, and the simplicity of the evaluation tasks. While some of these issues are addressed in rebuttal, it seems like most reviewers remained only slightly positive about the work.
Decision: Accept (poster)