Source-linked AI summary
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
TL;DR
In-context learning is sensitive to demonstration selection, and existing explanations do not fully connect its mechanisms to real-world LLMs. This paper models LLMs as latent-variable systems, learns task concepts with a small LM, and selects transferable demonstrations. The approach improves results across GPT models and classification datasets, with GSM8K providing an additional usefulness demonstration.
Problem
In-context learning is sensitive to demonstration choice, while its emergence from standard pretraining remains insufficiently understood for real-world LLMs.
Method
The paper uses a Bayesian latent-variable formulation and a two-step algorithm that learns task concept tokens with a small LM before selecting demonstrations for transfer to larger LLMs.
Results
Significant performance improvement is observed over baselines across eight GPT models and eight real-world text classification datasets, with additional usefulness demonstrated on GSM8K.
Takeaways & Limitations
The empirical findings support the hypothesis that LLMs implicitly infer a latent variable containing task information.
Takeaways & Limitations
The theoretical treatment focuses on the X → Y ← θ causal direction, while the alternative direction is discussed in the appendix.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper frames in-context learning as effective but highly demonstration-sensitive, while existing explanations remain disconnected from real-world LLMs. It proposes a Bayesian latent-variable account and a practical selection algorithm to bridge that gap.
- In-context learning achieves strong NLP performance by conditioning on a few demonstrations at inference time without updating model parameters.
- Demonstration choice, format, and order substantially affect in-context learning, making manual optimization difficult.
- Existing work leaves both effective-demonstration principles and the emergence of in-context learning from language-model pretraining insufficiently understood for real-world LLMs.
- The paper models LLMs through a Bayesian latent concept variable that captures task and format information.
- The proposed algorithm selects demonstrations with a small LM and transfers them directly to larger LLMs, improving performance across eight GPT models and eight real-world classification datasets.
- The authors present the latent-variable explanation as a Bayesian account intended to connect theoretical understanding with real-world LLM behavior.
2 Theoretical Analysis
The theoretical analysis represents tasks with a latent concept variable and considers causal directions linking inputs, labels, and concepts. Under stated distributional and independence assumptions, suitable demonstrations can make in-context prediction Bayes optimal.
- 2.1 Notations and Problem Setting: In-context learning conditions on labeled demonstrations and a test input, with generated continuation tokens representing the prediction.
- 2.1 Notations and Problem Setting: The task predicts a discrete target Y from token sequence X while introducing a potentially high-dimensional continuous latent variable θ.
- 2.1 Notations and Problem Setting: The analysis considers both X → Y ← θ and Y → X ← θ causal directions, focusing on the former in the main paper.
- 2.1 Notations and Problem Setting: For sentiment classification, the paper treats opinion as preceding the review, whereas helpfulness classification reverses the presumed causal direction.
- 2.2 Problem Analysis and Theoretical Results: Under the assumption that the LLM approximates the task distribution, the analysis establishes conditions connecting in-context prediction to the Bayes optimal classifier.
- 2.2 Problem Analysis and Theoretical Results: The framework does not treat one channel direction as universally superior; the preferred direction depends on the end task.
3 Method
The method learns task-specific latent concept tokens with a small LM, then selects demonstrations likely to predict those tokens and transfers the selected examples to other LLMs. Its tractability relies on simplifying independently sampled demonstrations, leaving combinatorial selection for future work.
- Latent concept learning: The method first estimates a task latent θd as new token embeddings through prompt tuning over the full candidate set.
- Demonstration selection: It then selects a smaller demonstration set maximizing the likelihood of generating the learned latent concept tokens.
- Demonstration selection: The selected demonstrations are obtained with a small LLM and directly transferred to other LLMs.
- Latent concept learning: Each task receives c newly added concept tokens whose embeddings are fine-tuned while the original LLM parameters remain frozen.
- Latent concept learning: The selection objective approximates the concept space using a diverse task subset because learning every possible θd is infeasible.
- Demonstration selection: Assuming demonstrations are independently sampled makes top-k individual selection tractable, but correlated demonstration combinations require an O(|Dd|^k) search left for future research.
- Demonstration selection: The authors report that demonstration order does not matter experimentally, so selected examples need not be reordered.
4 Experiments
Experiments evaluate latent-concept demonstration selection across diverse datasets and language models, finding improvements over baselines and useful behavior on GSM8K. Additional analyses examine token choice, demonstration count, concept-token count, ordering, and selected-example characteristics.
- Experiments cover eight datasets spanning sentiment, linguistic, topic, emotion, and hate-speech classification tasks.
- Main results: 12.5% relative improvement over uniform selection is achieved on average across eight GPT models and eight datasets.Demonstrations selected with GPT2-large generalize to the other GPT models.
- Results with non-GPT models: The GPT2-large-selected demonstrations improve in-context learning across GPT3, GPT3-instruct, GPT-J, OPT, and LLaMA models of similar sizes.GPT3-curie shows the largest improvement, likely because its pre-training distribution resembles GPT2-large’s.
- Results on GSM8K: On GSM8K, the proposed method significantly outperforms Uniform and Similar baselines, while demonstrations selected with 7B models outperform those selected with 1.5B models.The evaluation uses 200 test examples and greedy generation of final numerical answers.
- Ablation studies: Random concept tokens perform like random demonstrations, indicating that the gain comes from learned tokens containing task and format information.
- Ablation studies: The method outperforms random selection with k = 2, 4, 8, and 16, while uniform performance rises through k = 8 before declining slightly at k = 16.The study cautions that more demonstrations do not generally guarantee higher performance.
- Ablation studies: Performance drops when concept-token count increases from c = 10 to c = 20, while c = 5 cannot capture task and format information effectively.Longer token sequences may include meaningless tokens that reduce selectivity.
- Effect of demonstrations’ order: Selected demonstrations are usually insensitive to order, except for high variance on EmoC; their content appears to carry task information while ordering captures model-specific artifacts.The method’s ordering behavior contrasts with prior findings that good orderings may not transfer across models.
5 Related Work
Related work explains in-context learning through synthetic data distributions, optimization analogies, algorithm learning, and empirical studies. The paper positions its approach as a broader empirical and Bayesian explanation for real-world language models.
- Similarity-based demonstration selection is used as a baseline, while label-entropy reordering is omitted because the paper finds its method’s ordering does not matter.
- Prior studies connect in-context learning to Hidden Markov Models, skewed Zipfian distributions, gradient descent, and algorithm learning.
- Empirical work reports that demonstration labels may not matter and that chain-of-thought explanations can improve in-context learning performance.The paper states that the label-independence finding is not entirely accurate in its appendix.
6 Conclusion
The paper proposes understanding LLMs as implicit topic models that infer latent conceptual variables from prompts, motivating transferable demonstration selection across models.
- A two-step algorithm learns latent conceptual tokens with a small LLM, then selects demonstrations likely to predict those tokens for transfer to other LLMs.
A.1 Direct direction
The direct-direction analysis models task information as a latent variable and compares in-context classification with the Bayes-optimal classifier. The resulting theorem establishes that the in-context classifier cannot have lower misclassification probability than the Bayes-optimal rule.
- Under the direct X → Y ← θ direction, the task label is generated from the input and latent task variable.
- Theorem 2.3 states that the in-context learning classifier has higher or equal misclassification probability than the Bayes-optimal classifier.
- The Bayes-optimal classifier selects the label with maximum conditional probability given the latent task variable and input.
- For the reverse Y → X ← θ direction, the Bayes-optimal interpretation requires balanced label assignments.
- The corresponding reverse-direction theorem likewise states that the in-context classifier has higher or equal misclassification probability than the Bayes-optimal classifier.
A.3 Method
The method learns a latent task representation from demonstrations and uses it to model label or input distributions according to the task’s causal direction. Under an invertibility condition, the learned representation is identifiable with the true latent task variable.
- When the training loss is minimized, the learned representation matches the true conditional distributions in both causal directions.
- For X → Y ← θ tasks, the model matches P(Y|θd, X), while for Y → X ← θ tasks, it matches P(X|θd, Y).
- If the language model is invertible, the learned task representation θ̂d equals the true latent variable θd.
- The prompt templates and label mappings define how classification datasets are converted into model inputs.
B Experiments
Experiments evaluate the demonstration-selection method across diverse datasets, models, label mappings, causal directions, and ordering strategies. The method improves over random selection, transfers selected demonstrations across LLMs, and approaches a GPT2-large oracle on many datasets.
- Label mappings: With GPT2-large, the method is compared against random selection across original labels, random-word mappings, and random-label assignments.
- Label mappings: The method outperforms random selection when demonstrations use random words, indicating that demonstrations encode the mapping between inputs and labels.
- Optimal performance: The method comes close to the theoretical maximum accuracy on many datasets, although it lags behind on some datasets.
- Reordering: Reordering selected demonstrations is evaluated by maximizing the posterior of the concept tokens, with variability reported across five runs.
- Main results: Selected demonstrations from GPT2-large are applied unchanged to other LLMs, including models with up to 175B parameters.
C Limitations and Future Work
The method has scope and modeling limitations that motivate future work. Its current task coverage and task-set construction constrain broader application and understanding.
- Modeling assumptions: The method assumes that the language model captures the true language distribution, although this is inaccurate in practice.The paper notes that LLMs systematically underestimate rare text sequences in the long-tail distribution.
- Task-set construction: The accompanying diverse task set S is selected at the user’s discretion, limiting the current understanding of latent concept variables.The authors call for a better approach to constructing this task set.
- Task scope: The algorithm currently applies only to classification tasks, leaving more complex tasks such as math word questions and logical reasoning for future extensions.The paper suggests designing more complex latent variables for these tasks.
D Broader Impact
The paper presents a relatively low-cost approach for improving in-context learning, while emphasizing responsible use and documenting additional evaluations and analyses. It also highlights the need to address misleading outputs and intentional misuse as LLM use expands.
- Accessibility: The authors report significant in-context learning improvement through a relatively low-cost and simple approach.They frame this as making LLM use more accessible to individuals with limited resources.
- Responsible use: The broader impact discussion warns that LLMs may produce mistakes and misleading outputs, motivating user warnings and distribution regulation.The authors also identify intentional misuse as an ethical concern requiring mitigation.
- Implementation and safeguards: The authors plan to release code with clear warnings and guidelines to help users understand and address risks associated with the method.These measures are presented alongside broader regulation and safeguards for LLM distribution.