A Framework for Quantifying How Pre-Training and Context Benefit In-Context Learning

TL;DR

The paper introduces a framework to quantify how pre-training and context jointly shape in-context learning (ICL) in large language models. By modeling data generation via latent concepts and by formalizing context construction through a stacked prompt fed to a one-layer Transformer, the authors show that context can shift the output distribution toward the query-task distribution when the pre-training and query distributions differ. They establish a generalized Bayesian view linking ICL performance to context length and the KL divergence between pre-training and query distributions, and prove that, under certain conditions, the in-context predictor can approximate the Bayes-optimal predictor. Empirical validation on synthetic data and GPT-2 experiments demonstrates that context and task similarity between pre-training data and the target task materially improve predictive accuracy and Macro-F1, supporting the theoretical claims and highlighting practical guidance for prompt design and data curation.

Abstract

Pre-trained large language models have demonstrated a strong ability to learn from context, known as in-context learning (ICL). Despite a surge of recent applications that leverage such capabilities, it is by no means clear, at least theoretically, how the ICL capabilities arise, and in particular, what is the precise role played by key factors such as pre-training procedure as well as context construction. In this work, we propose a new framework to analyze the ICL performance, for a class of realistic settings, which includes network architectures, data encoding, data generation, and prompt construction process. As a first step, we construct a simple example with a one-layer transformer, and show an interesting result, namely when the pre-train data distribution is different from the query task distribution, a properly constructed context can shift the output distribution towards the query task distribution, in a quantifiable manner, leading to accurate prediction on the query topic. We then extend the findings in the previous step to a more general case, and derive the precise relationship between ICL performance, context length and the KL divergence between pre-train and query task distribution. Finally, we provide experiments to validate our theoretical results.

Figures (7)

• Figure 1: Summarization of the steps of ICL with latent concept generation, where each word consists of two attributes, i.e, topic and class. It follows standard ICL procedure with $3$ steps: pre-training, prompt construction, and in-context inference. Following the setting in xie2021explanation, we specialize the pre-training data and prompt generation process, so that they are conditioned on concepts. Intuitively, the concept defines a distribution of generated sequence, which is specified in Section \ref{['subsec: concept']}.
• Figure 2: Each word has two attributes: topic and class. For simplicity, we assume each topic consists of $K$ different classes, each with only one word.
• Figure 3: The distribution of topics in LDA example. Left: Prediction without ICL. Right: Prediction with ICL. Given $10$ different topics, denoted as '0' to '9'. The topic of the query task is '2'. When the prediction is made by $f_{\bf w}(\widetilde{X}_{\text{q}})$, the topic distribution does not lean toward '2'. However, when the prediction is made by $f_{\bf w}(\textbf{Z}_{\text{stacked}})$, it is more likely that the query topic '2' is predicted.
• Figure 4: Distribution of concept token visualized by t-SNE plot. Intuitively, the first four tasks are similar. The visualization of the learned concept token indeed shows the representation of these four tasks are close each other compared to other tasks.
• Figure 5: Example of encoded sequence.

Theorems & Definitions (2)

• proof
• proof