Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Closer Look at In-Context Learning under Distribution Shifts

Kartik Ahuja, David Lopez-Paz

TL;DR

This work analyzes in-context learning (ICL) under distribution shifts by contrasting transformers with a permutation-invariant set-based MLP baseline within a linear-regression data-generation framework. It derives conditions where the optimal ICL predictor aligns with OLS or ridge regression and shows that, given sufficient data, transformers can emulate OLS on in-distribution prompts, while performance degrades under distribution shifts—more gracefully for transformers than for the set-based baseline under mild shifts. Under severe distribution shifts, both architectures lose ICL capabilities, highlighting limits of current approaches. The findings emphasize the role of inductive biases in ICL and motivate theoretical and empirical exploration to improve OOD-ICL across architectures and tasks.

Abstract

In-context learning, a capability that enables a model to learn from input examples on the fly without necessitating weight updates, is a defining characteristic of large language models. In this work, we follow the setting proposed in (Garg et al., 2022) to better understand the generality and limitations of in-context learning from the lens of the simple yet fundamental task of linear regression. The key question we aim to address is: Are transformers more adept than some natural and simpler architectures at performing in-context learning under varying distribution shifts? To compare transformers, we propose to use a simple architecture based on set-based Multi-Layer Perceptrons (MLPs). We find that both transformers and set-based MLPs exhibit in-context learning under in-distribution evaluations, but transformers more closely emulate the performance of ordinary least squares (OLS). Transformers also display better resilience to mild distribution shifts, where set-based MLPs falter. However, under severe distribution shifts, both models' in-context learning abilities diminish.

A Closer Look at In-Context Learning under Distribution Shifts

TL;DR

This work analyzes in-context learning (ICL) under distribution shifts by contrasting transformers with a permutation-invariant set-based MLP baseline within a linear-regression data-generation framework. It derives conditions where the optimal ICL predictor aligns with OLS or ridge regression and shows that, given sufficient data, transformers can emulate OLS on in-distribution prompts, while performance degrades under distribution shifts—more gracefully for transformers than for the set-based baseline under mild shifts. Under severe distribution shifts, both architectures lose ICL capabilities, highlighting limits of current approaches. The findings emphasize the role of inductive biases in ICL and motivate theoretical and empirical exploration to improve OOD-ICL across architectures and tasks.

Abstract

In-context learning, a capability that enables a model to learn from input examples on the fly without necessitating weight updates, is a defining characteristic of large language models. In this work, we follow the setting proposed in (Garg et al., 2022) to better understand the generality and limitations of in-context learning from the lens of the simple yet fundamental task of linear regression. The key question we aim to address is: Are transformers more adept than some natural and simpler architectures at performing in-context learning under varying distribution shifts? To compare transformers, we propose to use a simple architecture based on set-based Multi-Layer Perceptrons (MLPs). We find that both transformers and set-based MLPs exhibit in-context learning under in-distribution evaluations, but transformers more closely emulate the performance of ordinary least squares (OLS). Transformers also display better resilience to mild distribution shifts, where set-based MLPs falter. However, under severe distribution shifts, both models' in-context learning abilities diminish.
Paper Structure (9 sections, 9 theorems, 22 equations, 2 figures)

This paper contains 9 sections, 9 theorems, 22 equations, 2 figures.

Table of Contents

  1. Introduction
  2. Related Works
  3. In-context Learning under Distribution Shifts
  4. A natural baseline for the above task
  5. Distribution shifts for ICL.
  6. Experiments
  7. Discussion
  8. Appendix
  9. Implications for Theorem \ref{['thm2']} when $\Sigma$ is not identity

Key Result

Lemma 1

If $\ell$ is the square loss, then the solution to equation eqn2 satisfies, $M^{*}(p_j) = \mathbb{E}[y_j|p_j], \text{almost everywhere in } \mathcal{P}_j$, $\forall j \in \{1, \cdots, k\}$.

Figures (2)

  • Figure 1: Comparison of MLP-set and transformers for noiseless setting, i.e., $\sigma=0$. a) ID-ICL ($\mu=0$), b) OOD-ICL (Mild distribution shift with $\mu = 2 \cdot \boldsymbol{1}$), c) OOD-ICL (Severe distribution shift with $\mu = 4\cdot \boldsymbol{1}$).
  • Figure 2: Comparison of MLP-set and transformers for noisy setting, i.e., $\sigma=1$. a) ID-ICL ($\mu=0$), b) OOD-ICL (Mild distribution shift with $\mu = 2 \cdot \boldsymbol{1}$), c) OOD-ICL (Severe distribution shift with $\mu = 4\cdot \boldsymbol{1}$).

Theorems & Definitions (13)

  • Lemma 1
  • Theorem 1
  • Theorem 2
  • Theorem 3
  • Theorem
  • Lemma
  • proof
  • Theorem
  • proof
  • Theorem
  • ...and 3 more