Understanding Emergent In-Context Learning from a Kernel Regression Perspective

TL;DR

The paper proposes that in-context learning in Transformer-based LLMs can be understood as kernel regression, deriving a kernel from HMM-inspired pre-training structure and showing that Bayesian inference on in-context prompts converges to a kernel-weighted predictor as demonstrations grow. Empirically, it shows that attention mechanisms allocate weights consistent with kernel regression, enabling reconstruction of ICL predictions and revealing where in the network such information is stored. The findings explain phenomena like the benefits of retrieving similar demonstrations, sensitivity to output formats, and the advantages of in-distribution representative samples, while highlighting remaining challenges such as sample-order effects and pre-training biases. This kernel-regression lens provides a concrete, testable framework to analyze and potentially improve ICL in large language models.

Abstract

Large language models (LLMs) have initiated a paradigm shift in transfer learning. In contrast to the classic pretraining-then-finetuning procedure, in order to use LLMs for downstream prediction tasks, one only needs to provide a few demonstrations, known as in-context examples, without adding more or updating existing model parameters. This in-context learning (ICL) capability of LLMs is intriguing, and it is not yet fully understood how pretrained LLMs acquire such capabilities. In this paper, we investigate the reason why a transformer-based language model can accomplish in-context learning after pre-training on a general language corpus by proposing a kernel-regression perspective of understanding LLMs' ICL bahaviors when faced with in-context examples. More concretely, we first prove that Bayesian inference on in-context prompts can be asymptotically understood as kernel regression $\hat y = \sum_i y_i K(x, x_i)/\sum_i K(x, x_i)$ as the number of in-context demonstrations grows. Then, we empirically investigate the in-context behaviors of language models. We find that during ICL, the attention and hidden features in LLMs match the behaviors of a kernel regression. Finally, our theory provides insights into multiple phenomena observed in the ICL field: why retrieving demonstrative samples similar to test samples can help, why ICL performance is sensitive to the output formats, and why ICL accuracy benefits from selecting in-distribution and representative samples. Code and resources are publicly available at https://github.com/Glaciohound/Explain-ICL-As-Kernel-Regression.

Figures (25)

• Figure 1: Our results suggests that LLMs might be conducting kernel regression on ICL prompts.
• Figure 2: Averaged attention map over GLUE-sst2 test set. A portion of attention on demonstrative samples is generally focused on label positions $y_i$. This conforms to the intuition in Theorem \ref{['thm:kernel']} that the inference on in-context learning prompts is a weighted average over sample labels.
• Figure 3: We use each head's attention weights on demonstrative samples to manually average sample labels $y_i$. These figures show that "reconstructed" outputs in some heads from layers 16$\sim$21 matches LLM prediction with as high as 89.2% accuracy, and matches ground truth with 86.4% accuracy.
• Figure 4: Pearson correlation between sample's attentions and prediction similarity$\text{sim}_\text{pred}(\mathbf{x}_\text{test}, \mathbf{x}_i)$ (Equation \ref{['eq:prediction_similarity']}). $x$-axis are layers and $y$-axis are heads in each layer. Note the resemblence between this heatmap and Figure \ref{['fig:attention_reconstruction']}.
• Figure 5: Key and value vectors encode label and LLM prediction information at high-attention position. Here y-axis denotes the relative position to the high-attention position in each sample. Each sphere is an attention head. The curve shows average accuracy within each layer.

Theorems & Definitions (2)

• Theorem 1
• proof