SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Assuming that  $\theta$  has a uniform prior, then our goal becomes finding the top  $k$  demonstrations that maximize  $\hat{P}_{M'}^d (\hat{\theta}^d |X_i^d,Y_i^d)$ . Note that the independence between demonstrations is a simplified assumption to reduce the combinatory search space of  $(X_1^d,Y_1^d),\dots,(X_k^d,Y_k^d)$ . In practice, selected demonstrations are likely correlated as some demonstrations may work well together but not necessarily work well by themselves. However, it would be too expensive to search the  $O(|\mathcal{D}^d|^k)$  combinations over the candidate set  $\mathcal{D}^d$ . In practice, this simplification works reasonably well. We leave this combinatory search problem to future research.
Also, as we are using an LLM to approximate the data distribution, the order of the demonstrations might matter. We will show in the Experiment section that the order does not matter, so no reordering of the selected demonstrations is needed. The full selection algorithm is shown in Algorithm 2.
Datasets. We conduct experiments on eight datasets from five different types of NLP classification tasks: sentiment analysis, linguistic analysis, topic classification, emotion classification, and hate speech detection. For sentiment analysis, we choose the Stanford Sentiment Treebank (SST2) dataset [35] from the GLUE benchmark [43] and the financial phrase bank (FPB) dataset [21]. SST2 is constructed based on movie reviews labeled "positive" or "negative", and FPB is based on financial news labeled "positive", "negative", or "neutral". For linguistic analysis, we choose the Corpus of Linguistic Acceptability (COLA) dataset [46] from the GLUE benchmark, based on sentences collected from linguistic books, labeled with "acceptable" or "unacceptable". For topic classification, we choose the DBpedia ontology classification dataset [52], based on DBpedia 2014 [16], labeled with 14 different ontology classes. For emotion classification, we choose the dataset from Chatterjee et al. [6] and Saravia et al. [33], both of which are collected from Twitter. Chatterjee et al. [6] (EmoC)
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 a three-turn contextual dialogue, while Saravia et al. [33] predict emotion given a Twitter message with clear emotion. For hate speech detection, we choose the online hate speech detection dataset (ETHOS) [27], collected from online social media platforms. Here we detect two types of hate speech: sexual orientation (ETHOS-SO) and religion (ETHOS-R). While in Section 2, we assume that all tasks share the same label space  $\mathcal{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<sup>4</sup>. We adopt the  $Y \rightarrow X \leftarrow \theta$  direction for sentiment analysis, topic classification, and emotion classification tasks, which is consistent with the intuition that people usually have some sentiment, topic, or emotion in mind before writing a piece of text. We adopt the  $X \rightarrow Y \leftarrow \theta$  direction for the linguistic analysis and hate speech detection type of tasks. While this is less intuitive, we can understand this as linguistic error and hate speech detection are more of a post hoc task in contrast to the previous tasks.
Without specification, we use  $k = 4$  number of demonstrations and  $c = 10$  number of concept tokens per dataset for our experiments, as the context length of GPT2 is 1024, and a larger number of demonstrations may not be able to completely fit into it. We use GPT2-large to learn the concept tokens and then compute the probability of each candidate demonstration example. We select our demonstrations from a randomly selected 100 example subset of the train set as the candidate set  $\mathcal{D}^d$ . We use the same set of demonstrations selected by GPT2-large for all other LLMs. We test the performance of the selected demonstrations using at most 1000 examples randomly sampled from the test set. Each experiment is repeated for five runs with different random seeds (the randomness comes from the sampling of the candidate set and the sampling of the test set). We adopt a large portion of the code from Min et al. [25], which is based on Huggingface [49].
Baselines. We consider the following baselines:
- Uniform: We uniformly select  $k$  demonstrations from  $\mathcal{D}$  for each test example.
- Similar: According to Liu et al. [19], demonstrations that are semantically similar to the test example would hare more performant. Following their method, we use a pre-trained sentence Transformer [31] to calculate the cosine similarity between the demonstrations and test examples. We choose the top  $k$  similar demonstrations from  $\mathcal{D}$  for each test example.
Main results.5 Figure 2 shows our main results averaged over all eight datasets, using the first-generation GPT2s and GPT3s, without any instruction fine-tuning [28] or Reinforcement Learning from Human Feedback (RLHF) [36]. Our method significantly outperforms baselines on eight different LLMs, with  $12.5\%$  relative improvement to the uniform selection baseline on average, which shows the effectiveness of our method. The demonstrations selected by our method are exclusively based on GPT2-large, while the same set of demonstrations can be generalized to all other GPTs.
Results with non-GPT models. In Figure 3a, we test the demonstrations selected by our method using GPT2-large on more LLMs (GPT3 [4], GPT3-instruct [28, 36], GPT-J [44], OPT [51], and LLaMA [38]) with similar sizes (6-7B), and show that the selected demonstrations improve in-context learning performance of all of them. The fact that GPT3-curie obtains the largest performance improvement is likely because similar pre-training data distributions help the generalization of the
Detailed results see Figure 6 in Appendix B.
The complete results with standard deviations in this section can be found in Appendix B.
(a) Proposed method v.s. randomly selected demonstrations
(b) Proposed method v.s. using randomly selected tokens
Figure 3: In-context learning accuracy averaged over all eight datasets.
Uniform	Similar	Ours w/ Llama 2 (7B)	Ours w/ GPT2-XL (1.5B)
Prompt tuning	-	-	3.7	1.3
Llama 2 (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-turbo)	76.5	78.1	81.2	80.4
Table 1: Prompt tuning and 4-shot in-context learning accuracy on a 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 a dataset expanded from the GPT2 pre-training corpus [4]. Thus the pre-training distribution of GPT3-curie and GPT2-large can be assumed to be similar.
Results on GSM8K. Since our primary goal is to connect the theory with real-world models and datasets, we did not try to include harder tasks in the main results in Figure 2. In practice, our proposed method is most effective with hard tasks that even parameter-efficient fine-tuning with smaller models cannot outperform in-context learning with the same or larger models. To showcase the usefulness of our proposed algorithm, We added a new dataset, GSM8K [9], which is a math word problem-solving dataset with chain-of-thoughts solutions. Table 1 shows the test accuracy of the final numerical answer with greedy generation. We randomly select a test set of 200 examples instead of using the full test set for computation efficiency.
As shown in the first row of Table 1, prompt tuning with ten new tokens can only obtain less than  $4\%$  accuracy on the GSM8K test set. The last four rows show the in-context learning results with different size Llama 2 models [39] and ChatGPT. Our proposed demonstration selection method (last two columns) significantly outperformed the Uniform and Similar baseline. We also find that the demonstrations selected with a larger model (7B) are more effective than those selected with a smaller model (1.5B). The results show that our demonstration selection method is a good choice under a low data setting, with a small computing budget and minimal inference latency. Our proposed method can also potentially be combined with other prompting techniques [8] to boost performance further.
Learned tokens v.s. Random tokens. To confirm the critical role of the latent concept variable in the proposed demonstration selection algorithm, we compare the performance of using the learned concept tokens versus using randomly selected tokens from the original vocabulary in Figure 3b. The demonstrations selected by random tokens only obtain the same performance as randomly selected demonstrations, showing that the performance gain of our method comes from the learned concept tokens containing the task and format information, not other elements of our algorithm.
$k$  ablation study. While we use  $k = 4$  demonstrations for all experiments, we also test the effectiveness of our method using different  $k$ . As shown in Figure 4a, our method significantly outperforms the random selection baseline with  $k = 2, 4, 8,$  and 16. To fit in large  $ks$ , we use GPT3-ada with a longer context length (2048). Note that for real-world tasks, it is in general not true that more demonstrations guarantee higher performance [7]. We can see that the uniform baseline performance increases from  $k = 2$  to  $k = 8$ , then drops a little at  $k = 16$ . Our method improves the uniform baseline by around  $5\%$  absolute for all  $ks$ , while  $k = 4$  improves the most  $(6.6\%)$ . Our method appears to have a diminishing effect when  $k$  becomes larger, which is likely because the effect of more demonstrations overwhelms the effect of demonstration choices.
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.
$c$  ablation study. While we use  $c = 10$  number of concept tokens for all experiments, we also investigate the effect of different  $c$  on our method. When  $c$  is small  $(c = 5)$ , the concept tokens cannot effectively capture the task and format information, thus cannot im
(a)  $k$  ablation study.
(b)  $c$  ablation study.
Figure 4: In-context learning accuracy of our method versus random selection baseline averaged over all eight datasets with GPT3-ada.
prove the performance. When  $c$  increases from 10 to 20, we observe a drop in the performance. It is likely because the selectivity of the concept tokens decreases when  $c$  increases. The longer the concept token sequence is, the more likely it will contain meaningless tokens that do not contribute to demonstration selection.
Effect of demonstrations' order. We find that the demonstrations selected by our method are insensitive to their order in most cases.<sup>7</sup> An exception is the EmoC dataset, where our method has a high variance. On the contrary, Lu et al. [20] found that the order of the demonstration matters, and a good ordering cannot be transferred between different LLMs. We suspect that the ordering only matters when the demonstration selection method is not robust. Since Lu et al. [20] randomly selects one set of demonstrations for the whole test set, the variance in performance is high with different demonstrations, thus ordering matters. And since such ordering is not transferable while our selected demonstrations are highly transferable, we suspect the core task information is stored in the content of the demonstrations, while the ordering mainly captures model-specific artifacts.
Qualitative analysis. In Figure 5, we provide a t-SNE [40] projection of the learned concept token embeddings. The tokens corresponding to semantically similar tasks are close together. Note that this result only aims to provide a straightforward illustration of concept tokens. The effect of concept tokens should be understood by the previous quantitative results.
We also list the top 4 selected demonstrations in Table 14 in Appendix B. Compared to the examples with lower scores, the selected examples for GSM8K have more deductive reasoning (i.e. with the connecting words 'so', 'then', 'thus', etc.), instead of listing parallel conditions. For SST2, the selected examples are longer and more complex, sometimes including a 'but'. This can be understood as these
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.
Figure 5: t-SNE plot of the learned concept tokens for each task. Concept tokens that can be explained by similar tokens are summarized in the graph.
Heuristic solutions, such as selecting demonstrations based on the similarity between the demonstrations and test input [19, 37, 32] have been proposed. [20] propose to reorder the demonstration based on the entropy of the predicted labels. In this paper, we use the similarity-based selection method
Detailed results see Figure 9 in Appendix B.
The list of similar tokens for these concept tokens can be found in Table 13 in 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 a Hidden Markov Model distribution [50] and a skewed Zipfian distribution with high burstiness [5]. Other studies have sought to understand the underlying mechanisms of in-context learning by making connections with gradient descent [42, 10, 1], formalizing it as an algorithm learning problem [18], or proposing a latent variable theory similar as ours [14, 12, 50]. While providing valuable insights on how in-context learning works, these works are limited to synthetic datasets and toy Transformers, while it remains unclear if these results generalize to LLMs pre-trained on real-world text data and whether these results can help in-context learning performance. In contrast, we propose a Bayesian explanation of in-context learning that can be verified with real-world LLMs on various NLP datasets. Dai et al. [10] provide a practical algorithm based on the understanding that the Transformer has a dual form of gradient descent. However, their empirical results are smaller in scale, with six datasets and only one model (350M), and has less significant improvements (5.4% relative to baseline).
There are also works trying to understand in-context learning from an empirical perspective [2, 24]. Min et al. [26] found demonstrations' ground truth labels do not matter for in-context learning, which we find is not entirely accurate in Appendix B. On the other hand, chain-of-thoughts prompting [48, 53, 45] find that providing step-by-step explanations improves in-context learning performance.
In this work, we endeavor to comprehend large language models (LLMs) through a Bayesian lens and posit them as implicit topic models that infer a latent conceptual variable from prompts. Motivated by this understanding, we propose a two-step algorithm that first extracts latent conceptual tokens from a small LLM and then selects demonstrations that have the greatest probability of predicting the corresponding conceptual tokens. The selected demonstrations can then be directly generalized to other LLMs. The efficacy of our algorithm across various text classification datasets and GPT models validates our explanation of in-context learning.
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 or position of the US government. We thank Google for its generous gift to the University of California.
[1] E. Akyurek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661, 2022.
[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-based case study at 66 billion scale. arXiv preprint arXiv:2212.09095, 2022.
[3] D. M. Blei, A. Y. Ng, and M. I. Jordan. Latent dirichlet allocation. J. Mach. Learn. Res., 3 (null):993-1022, mar 2003. ISSN 1532-4435.
[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, 33: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 arXiv:2205.05055, 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
P. Mishra, I. Molybog, Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, P. Xu, Z. Yan, I. Zarov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, S. Edunov, and T. Scialom. Llama 2: Open foundation and fine-tuned chat models. 7 2023. URL http://arxiv.org/abs/2307.09288.
[40] L. van der Maaten and G. Hinton. Visualizing data using t-sne. Journal of Machine Learning Research, 9(86):2579-2605, 2008. URL http://jmlr.org/papers/v9/vandermaaten08a.html.
[41] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[42] J. von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vlademyrov. 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 S. R. 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. https://github.com/kingoflolz/mesh-transformer-jax, 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.11171, 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] T. Wolf, L. Debut, V. Sanh, J. Chaumont, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771, 2019.
[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://openreview.net/forum?id=RdJVFCHjUMI.
[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 language models. arXiv preprint arXiv:2205.01068, 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. Scharli, L. Hou, J. Wei, N. Scales, X. Wang, D. Schuurmans, O. Bousquet, Q. Le, and E. Chi. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205.10625, 2022.
A.1 Direct direction
Assumption A.1. (Assumption 2.1) Assume that  $P_{M}(X) = P(X)$ , and  $P_{M}^{d}(Y|\pmb{\theta}, X) \propto P(Y|\pmb{\theta}, X)$  for  $X \to Y \gets \pmb{\theta}$ .
Proposition A.2. (Proposition 2.2) If task  $d$  follows the  $X \to Y \gets \theta$  direction,  $\arg \max_{y \in \mathcal{Y}} P_M^d(Y = y|\theta^d, X)$  is the Bayes optimal classifier.
Proof. Since the data generation of the task  $d$  can be written as  $Y = f(X, \theta^d, \epsilon)$ , we have
$$
P ^ {d} (Y | X) = P (Y | \theta^ {d}, X).
$$
And by Assumption A.1, we have
$$
\operatorname * {a r g   m a x} _ {y \in \mathcal {Y}} P _ {M} ^ {d} (Y = y | \theta^ {d}, X) = \operatorname * {a r g   m a x} _ {y \in \mathcal {Y}} P (Y = y | \theta^ {d}, X).
$$
Thus  $\arg \max_{y\in \mathcal{Y}}P_M^d (Y = y|\theta^d,X)$  is the Bayes optimal classifier.
Theorem A.3. (Theorem 2.3) If task  $d$  follows the  $X \to Y \leftarrow \theta$  direction, then the in-context learning classifier
always has a higher or equal probability of misclassification to the Bayes optimal classifier  $\arg \max_{y\in \mathcal{Y}}P_M^d (Y = y|\theta^d,X)$ . Equality only takes when
$$
\forall x \in \mathcal {X}, P _ {M} ^ {d} \left(\theta^ {d} \mid X _ {1} ^ {d}, Y _ {1} ^ {d}, \dots , X _ {k} ^ {d}, Y _ {k} ^ {d}, X = x\right) = 1.
$$
Proof. Recall that in Equation (1), we have
$$
P _ {M} ^ {d} (Y | X _ {1} ^ {d}, Y _ {1} ^ {d}, \dots , X _ {k} ^ {d}, Y _ {k} ^ {d}, X) = \int_ {\Theta} P _ {M} ^ {d} (Y | \pmb {\theta}, X) P _ {M} ^ {d} (\pmb {\theta} | X _ {1} ^ {d}, Y _ {1} ^ {d}, \dots , X _ {k} ^ {d}, Y _ {k} ^ {d}, X) d \pmb {\theta}.
$$
By Proposition A.2,  $\arg \max_{y\in \mathcal{Y}}P_M^d (Y = y|\theta^d,X)$  is the Bayes optimal classifier. Let  $C_{\pmb{\theta}}(X) =$ $\arg \max_{y\in \mathcal{Y}}P_M^d (Y = y|\pmb {\theta},X)$  , then the risk is defined as the probability of misclassification
$$
R \left(C _ {\boldsymbol {\theta}}\right) = P \left(C _ {\boldsymbol {\theta}} (X) \neq Y\right) = \mathbb {E} _ {X Y} \left[ \mathbb {1} _ {C _ {\boldsymbol {\theta}} (X) \neq Y} \right].
$$
Denote the in-context learning classifier  $\arg \max_{y\in \mathcal{Y}}P_M^d (Y = y|X_1^d,Y_1^d,\dots,X_k^d,Y_k^d,X)$  by  $C_k(X)$ . We then have
$$
R (C _ {k}) = \mathbb {E} _ {X Y} [ \mathbb {1} _ {C _ {k} (X) \neq Y} ] = \mathbb {E} _ {X} [ \sum_ {y \in \mathcal {Y}} (1 - P _ {M} ^ {d} (Y = y | \theta^ {d}, X)) \mathbb {1} _ {C _ {k} (X) = y} ].
$$
Such risk is minimized if and only if  $C_k(X) = C_{\theta^d}(X)$ , which only holds when  $P_M^d(\theta^d | X_1^d, Y_1^d, \dots, X_k^d, Y_k^d, X = x) = 1$  for all  $x \in \mathcal{X}$ .
A.2 Channel direction
Assumption A.4. Assume that  $P_{M}(X) = P(X)$ , and  $P_{M}^{d}(X|\pmb{\theta}, Y) \propto P(X|\pmb{\theta}, Y)$  for the  $Y \to X \gets \pmb{\theta}$  direction.
Proposition A.5. If task  $d$  follows the  $Y \to X \gets \theta$  causal direction,  $\arg \max_{y \in \mathcal{Y}} P_M^d(X|\theta^d, Y = y)$  is the Bayes optimal classifier when the label assignment is balanced.
Proof. Since the data generation of the task  $d$  can be written as  $X = g(Y,\theta^d,\epsilon)$ , we have
$$
P ^ {d} (X | Y) = P (X | \theta^ {d}, Y)
$$
When the label is balanced, i.e.  $P^{d}(Y) = \frac{1}{|\mathcal{V}|}$ , we have
$$
P ^ {d} (Y | X) = \frac {P ^ {d} (X | Y) P ^ {d} (Y)}{P (X)} \propto P ^ {d} (X | Y)
$$
And by Assumption A.4, we have