PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding Emergent In-Context Learning from a Kernel Regression Perspective

Chi Han, Ziqi Wang, Han Zhao, Heng Ji

arXiv:2305.12766v3cs.CLcs.AIcs.LG

TL;DR

The paper addresses how pretrained language models can perform in-context learning despite the gap between pretraining and ICL prompts. It develops a kernel-regression account, showing convergence from Bayesian prompt inference and finding matching behavior in LLM attention and hidden features.

  • Problem

    The mechanism by which pretrained LLMs acquire ICL remains unclear, and the format shift between prompts and pretraining poses a challenge for optimization-based explanations.

    Source: §1 · p. 1

  • Method

    The paper analyzes Bayesian inference on in-context prompts as kernel regression and empirically examines LLM attention and hidden features during ICL.

    Source: §1 · p. 2

  • Results

    Over 80% accuracy is achieved when attention values are used to reconstruct model outputs, and attention and hidden features show kernel-regression behavior during ICL.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The theory explains reported benefits of similar demonstrations and in-distribution representative samples, as well as sensitivity to output formats.

    Source: §1 · p. 2

    The HMM-mixture assumptions may not capture natural language’s full complexity and unbounded lengths, and the many hidden states may be inefficient.

    Source: §7 · p. 11

Abstract

from arXiv · show

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.

1 Introduction

The paper investigates whether kernel regression can explain how pretrained transformers perform in-context learning on sequential data. It links this account to Bayesian inference and examines whether LLM attention and hidden features behave as the theory predicts.

1 Introduction

  • The kernel-regression perspective links predictions to similarities between the test input and demonstration inputs, with more similar examples contributing through their outputs.The kernel measures input similarity and the prediction is a weighted sum of the examples’ outputs.Source: §1 · p. 2
  • Bayesian predictions on in-context prompts converge to kernel regression as demonstrations grow, providing a tractable account of ICL inference.The kernel estimates a test output from a similarity-weighted sum of demonstration outputs.Source: §1 · p. 2
  • The theory offers explanations for benefits from retrieving similar demonstrations and selecting in-distribution, representative examples, as well as sensitivity to output formats.These phenomena are identified as insights from the theoretical analysis.Source: §1 · p. 2
  • During ICL, the last token’s attention allocation follows the proposed account, and plugging attention values into its equation reconstructs model outputs with over 80% accuracy.The studies also examine how intermediate layers compute information needed for kernel regression.Source: §1 · p. 2

2 Related Work

Prior work studies ICL empirically, theoretically, and as a possible algorithm implemented by transformers. This paper is situated among explanations connecting ICL to Bayesian inference and attention to kernel regression, while emphasizing emergent language-model behavior.

2.1 In-Context Learning

  • Prior ICL studies examine demonstration formats and effects, sample selection, behavioral characteristics, and links between ICL and pretraining data.The cited work includes Min et al. (2022b;a), Zhao et al. (2021), Liu et al. (2022), Lu et al. (2022), and others.Source: §2.1 · pp. 2–3
  • Bayesian accounts interpret ICL under assumptions about data generation, but prior work leaves the computational feasibility of growing-depth inference unresolved.Other studies construct transformers that approximate optimization algorithms, often in simplified linear or synthetic settings rather than emergent language pretraining.Source: §2.1 · p. 3

2.2 Emergent Ability of LLMs

  • Emergent abilities including ICL have been reported after pretraining on massive language data, while theoretical analysis of these capacities remains challenging.The paper frames its work as a preliminary step toward understanding ICL as a special case of language-model capacity emergence.Source: §2.2 · p. 3

2.3 Associating Attention with Kernel Regression

  • Related work connects self-attention with kernel ridge regression for uncertainty calibration, while other work relates attention to kernel regression using a different definition of ICL.The cited discussion also places these links alongside associative-memory accounts of recalling patterns.Source: §2.3 · p. 4

3 Formulation

The formulation models pretraining as a mixture of task-specific hidden Markov models and represents ICL prompts as demonstrations followed by a test input. It adopts regularity, delimiter, and task-distinguishability assumptions to support the analysis.

3.1 Preliminaries: Hidden Markov Models

  • The pretraining corpus is modeled as a mixture of task-specific HMMs, with a unified HMM defined through state transitions and observation emissions.The formulation specifies initial, transition, and emission distributions, and defines the transition and emission matrices.Source: §3.1 · p. 4

3.2 In-Context Learning

  • Each task is associated with a distinct initial state, and the support of the pretraining distribution consists of those task-start states.The formulation gives a lower bound on each task’s transition rate.Source: §3.2 · p. 4
  • An ICL prompt contains sampled input-output demonstrations separated by delimiter tokens, followed by the test input.The demonstration pairs are i.i.d. samples from the test task; delimiter emissions have a stated lower bound.Source: §3.2 · p. 4
  • For an observation sequence, the formulation groups intermediate matrices into an operator and uses covariance and spectral-radius quantities to characterize task differences.It also bounds the Frobenius norm of the sequence operator.Source: §3.2 · pp. 4–5

3.3 Assumptions

  • The borrowed assumptions require delimiter tokens to indicate the start of new sequences, with line breaks offered as an example.The delimiter must be sufficiently indicative of sequence boundaries.Source: §3.3 · p. 5
  • Distinct tasks must have a lower-bounded KL divergence over their first l tokens, so increasing l can provide sufficient task information.This assumption requires tasks to be distinguishable.Source: §3.3 · p. 5

4 Theoretical Analysis

The analysis presents in-context inference as converging to kernel regression as demonstrations increase, then uses this account to explain several observed ICL behaviors.

4.1 Explaining ICL as Kernel Regression

  • As the number of demonstrations grows, inference on an ICL prompt converges to a kernel-regression form and approaches the prompt-conditioned likelihood.The prediction aggregates demonstration labels according to their similarity to the test input.Source: §4.1 · p. 5
  • The account loosely aligns kernel regression with attention: input representations correspond to keys and queries, while demonstration labels correspond to values.The comparison replaces the attention score with a kernel similarity; the authors assume preceding layers learn the sample representations.Source: §4.1 · pp. 5–6
  • The kernel compares semantic representations of inputs, assigning greater similarity when their follow-up conditional probabilities under the language model are similar.The paper interprets the flattened Tx representation as an HMM belief state determining conditional probabilities after a prefix.Source: §4.1 · p. 6

4.2 Insights Provided by the Explanation

  • Selecting demonstrations similar to the test input is explained as narrowing the kernel bandwidth to a local window, which reduces regression bias.The authors identify retrieval of samples with similar representations as a promising direction for improving ICL scores.Source: §4.2 · p. 6
  • Changing the label set can mislead predictions into a different output space, so demonstrations and test samples benefit from aligned label spaces and formats.The account treats the output as a weighted vote over demonstration labels.Source: §4.2 · p. 6
  • OOD input variants on SST2 more or less reduce ICL accuracy, with typos causing the largest drop, to near random.The variants include rare-word substitutions, more complex sentence structures, and typos.Source: §4.2 · pp. 6–7
  • The prediction depends on both demonstrations and pretraining priors; when a task distribution resembles pretraining, the stated error term is smaller.The authors connect this account to the finding that ICL cannot overcome prediction preferences from pretraining.Source: §4.2 · p. 7
  • The framework remains partial: it does not explain sensitivity to sample ordering, and the authors say implicit reasoning in text generation must also be considered.They also identify differing findings about robustness to perturbed or random labels as an unresolved phenomenon.Source: §4.2 · pp. 7–8

5 Empirical Analysis

The empirical analysis tests the kernel-regression account through attention patterns, prediction reconstruction, and intermediate features. Its findings align with that account, while the final reconstruction evaluation reports competitive task performance.

5 Empirical Analysis

  • The empirical study uses GPT-J 6B and SST-2 validation data as its case study, with results on additional tasks reported in the appendix.The authors cite academic-lab computation resources as a constraint and analyze the model on one Tesla V100.Source: §5 · p. 8

5.1 Where Are Attentions Distributed During ICL?

  • Attention concentrates on test-input endings, sequence starters, and demonstration labels; label-token attention accords with the predicted aggregation of label information.The analysis averages attention maps after rescaling sample lengths to 30 tokens across 28 layers and 16 heads.Source: §5.1 · p. 8

5.2 Can Attentions Be Interpreted as Kernel Functions?

  • Several heads in middle layers reconstruct the model’s predictions at 89.2% accuracy and ground-truth labels at 86.4% accuracy.The highest reported head is layer 17, head 10; the authors note that a few adjacent middle layers can perform the reconstruction.Source: §5.2 · p. 9

5.3 Which Samples Receive High Attention?

  • Across layers and heads, attention correlates with prediction similarity, and the attention-head pattern resembles Figure 3, although absolute correlations are not high.Prediction similarity is measured by applying the LLM to the texts outside the ICL prompt.Source: Figure 4 · p. 9, §5.3 · p. 9

5.4 Do Intermediate Features Store Information Useful for Kernel Regression?

  • Key and value features at the high-attention position generally achieve their best accuracy, encoding prediction information and label information, respectively.The analysis examines positions from −1 to 3 relative to the label; later positions can access label information, while distance introduces noise.Source: §5.4 · p. 10

5.5 How Well Can Our Explanation Reconstruct ICL Prediction?

  • Reconstructed ICL outputs achieve 68% to 80% accuracy except on MNLI, match model-level task performance, and perform similarly or better than sentence-encoder kernel regression.The evaluation uses the attention head with the highest correlation with model predictions and tests on 300 held-out examples per task.Source: §5.5 · p. 10

6 Conclusions and Future Work

The paper concludes that ICL in Transformer-based LLMs can be understood through kernel regression, with empirical attention and hidden-feature findings congruent with that view. Sample ordering and robustness to perturbed labels remain open challenges.

6 Conclusions and Future Work

  • The paper proposes kernel regression as a view of Transformer-based LLM in-context learning, supported by empirical congruence in attention and hidden features.The theory also addresses sample-similarity retrieval, output-format sensitivity, and selection of in-distribution, representative examples.Source: §6 · p. 10
  • The account offers explanations for why similar demonstrations, aligned output formats, and in-distribution representative samples benefit ICL.These are phenomena the authors say their theory explains.Source: §6 · p. 10
  • Understanding sample-order effects and robustness to perturbed labels remains future work.The authors identify these as remaining challenges in the topic.Source: §6 · p. 11

7 Limitations

The analysis depends on an HMM-mixture setting that may not elegantly represent natural language, and the paper notes challenges to the Bayesian-inference view of ICL.

7 Limitations

  • The theoretical analysis assumes an HMM mixture, which may not elegantly capture natural language’s complexity and unbounded lengths.The authors also note that representing it may require a large number of hidden states and may be inefficient.Source: §7 · p. 11
  • The paper presents its natural-language model as potentially inelegant rather than a full account of unbounded-length language.The limitation is tied to the model’s large hidden-state requirements and possible inefficiency.Source: §7 · p. 11
  • Some ICL phenomena are not explained as Bayesian inference, and prior work challenges that view.The authors say divergence from kernel-regression behavior can be partly explained by the first error term in Equation 7.Source: §7 · p. 11

A Proofs

The proof excerpts use linearity and concentration for independent samples, then develop a bridge from hidden-state inference to the Bayesian inference over the ICL prompt.

A Proofs

  • The proof denotes a kernel-regression function, then uses the linearity of vec(·) to replace vec(Tx_i)^T with Tx_i.The excerpt does not provide the full derivation.Source: A Proofs · pp. 14–15
  • Treating (x_i, e(y_i)) as independent samples from P(Y | x_i, p_θ⋆) lets the proof apply Hoeffding’s inequality for a probabilistic bound.The bound is stated with confidence parameter δ in the excerpt.Source: A Proofs · p. 15
  • The proof bridges hidden-state inference for the first token of x_test to Bayesian inference over the ICL prompt by bounding state likelihoods.It also considers why a test state is unlikely to correspond to another task’s state.Source: A Proofs · pp. 15–16

B Results on more tasks

The additional experiments extend beyond the SST2 case study to Rotten Tomatoes, three Tweet Eval subtasks, and MNLI.

B Results on more tasks

B.1 Rotten Tomatoes

The appendix captions describe attention-map and kernel-regression-oriented analyses across five tasks, alongside a synthetic HMM verification in which more demonstrations improve fitting metrics.

B.1 Rotten Tomatoes

B.2 Tweet Eval (Hate)

B.3 Tweet Eval (Irony)

B.4 Tweet Eval (Offensive)

B.5 MNLI

C Synthetic Verification

  • As demonstrations increase, synthetic-HMM kernel regression has decreasing loss and increasing accuracy when fitting the Bayesian posterior.The experiment uses randomly parameterized HMMs with 8 tasks, 80 states, and 100 observations; loss converges to a non-zero value.Source: C Synthetic Verification · p. 27, Table 3 · p. 27

D Results on Llama-2 and Llama-3

Llama-2 and Llama-3 8B visualizations generally show trends similar to Section 5. A single-datum visualization has a similar but slightly sparser pattern than Section 5.1.

D Results on Llama-2 and Llama-3

E Single Datum Visualization

Loading 2305.12766v3…