PaperlayerView exact revision on arXiv

Source-linked AI summary

Dual Operating Modes of In-Context Learning

Ziqian Lin, Kangwook Lee

arXiv:2402.18819v2cs.LG

TL;DR

Existing theory has not formally explained how ICL’s task-learning and task-retrieval modes arise together. This paper models clustered linear tasks and derives posterior updates, using them to explain early ascent and predict bounded efficacy for biased-label ICL.

  • Problem

    How to formally explain ICL’s dual operating modes and the conditions under which retrieval or learning dominates remains unresolved.

    Source: §1 · p. 1

  • Method

    The paper extends a pretraining model to clustered linear tasks with task-dependent input distributions and derives a closed-form task posterior for optimal prediction.

    Source: §3.3 · p. 4, §O.2 · p. 52

  • Results

    The analysis explains early ascent and predicts bounded efficacy for biased-label ICL, and the paper reports validating its findings and predictions with large language models.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The model offers a framework for analyzing both task retrieval and task learning and for interpreting the reported ICL phenomena.

    Source: §7 · p. 9

    The framework assumes linear-regression pretraining tasks and noiseless labels for in-context samples, unlike complex real-world LLM tasks.

    Source: §7 · p. 9

Abstract

from arXiv · show

In-context learning (ICL) exhibits dual operating modes: task learning, i.e., acquiring a new skill from in-context samples, and task retrieval, i.e., locating and activating a relevant pretrained skill. Recent theoretical work investigates various mathematical models to analyze ICL, but existing models explain only one operating mode at a time. We introduce a probabilistic model, with which one can explain the dual operating modes of ICL simultaneously. Focusing on in-context learning of linear functions, we extend existing models for pretraining data by introducing multiple task groups and task-dependent input distributions. We then analyze the behavior of the optimally pretrained model under the squared loss, i.e., the MMSE estimator of the label given in-context examples. Regarding pretraining task distribution as prior and in-context examples as the observation, we derive the closed-form expression of the task posterior distribution. With the closed-form expression, we obtain a quantitative understanding of the two operating modes of ICL. Furthermore, we shed light on an unexplained phenomenon observed in practice: under certain settings, the ICL risk initially increases and then decreases with more in-context examples. Our model offers a plausible explanation for this "early ascent" phenomenon: a limited number of in-context samples may lead to the retrieval of an incorrect skill, thereby increasing the risk, which will eventually diminish as task learning takes effect with more in-context samples. We also theoretically analyze ICL with biased labels, e.g., zero-shot ICL, where in-context examples are assigned random labels. Lastly, we validate our findings and predictions via experiments involving Transformers and large language models.

1. Introduction

The introduction frames ICL as both learning a skill from examples and retrieving a pretrained one, while noting that existing theoretical models do not formally explain both modes together. The paper proposes a clustered-task model with task-dependent inputs and analyzes how posterior updates account for these modes and observed phenomena.

1. Introduction

  • Existing theoretical work has not formally established how ICL’s task-learning and task-retrieval modes arise together or when either mode dominates.The introduction frames these as open questions about the link between pretraining and ICL.Source: §1 · p. 1

A New Model for Pretraining Data

  • The model extends prior pretraining-data models with multiple task groups and task-dependent input distributions, representing clustered structure in real-world data.Each Gaussian-mixture component models a task group, and each component can have its own input distribution.Source: A New Model for Pretraining Data · p. 2
  • The analysis treats the pretraining distribution as a prior and in-context examples as observations, deriving a closed-form task posterior for the optimal squared-loss predictor.It tracks example-driven updates to component means and mixture probabilities.Source: A New Model for Pretraining Data · p. 2
  • With few examples, component re-weighting can dominate and yield task retrieval; with many, component shifting can dominate and yield task learning.The paper presents these as consequences under mild assumptions.Source: A New Model for Pretraining Data · p. 2
  • A limited number of examples may retrieve an incorrect skill and raise ICL risk, which can later diminish as task learning takes effect.This is the paper’s plausible explanation for early ascent, observed under certain conditions.Source: A New Model for Pretraining Data · pp. 2–3
  • For biased-label ICL, the model predicts bounded efficacy: retrieval can work with sufficient information, but later learning may fit biased labels and raise test risk.The paper reports observing the predicted phenomenon with several large language models when using more examples.Source: A New Model for Pretraining Data · p. 3

2. Related Work

Related work studies ICL through Bayesian inference, task recognition, and gradient descent, using models with differing analytical tractability and assumptions. This paper positions its approach as tractable while representing clustered pretraining tasks.

Dual Operating Modes of ICL

  • Pan et al. (2023) distinguish task recognition, called task retrieval here, from task learning; random-label examples preserve efficacy, while correct abstract labels improve with model size and example count.Correct labels benefit from both modes and yield the highest accuracy in the cited study.Source: §2 · p. 3

Explaining ICL via Bayesian Inference

  • Xie et al. (2022) model pretraining with an HMM that reflects long-range dependencies, but lacks a closed-form optimal predictor for detailed ICL analysis.Garg et al. (2022) and Raventos et al. (2023) instead study sequences of labeled pairs, with Raventos et al. proposing noisy linear regression.Source: §2 · p. 3

Explaining ICL via Gradient Descent

  • Prior work connects ICL to gradient descent through theoretical constructions and empirical similarities, including conditions under which trained Transformers implement gradient descent.Other studies examine training dynamics and prediction error relative to a best linear predictor.Source: §2 · p. 3

Others

  • This paper assumes an optimally trained model on infinite pretraining data and does not study a particular architecture or pretraining’s statistical aspects.Its model aims to retain tractable analysis while capturing clustered pretraining data, which prior noisy-linear-regression models do not model.Source: §2 · pp. 3–4

3. Pretraining and Data Generative Model

The paper formalizes next-token prediction on sequences of labeled samples drawn from diverse tasks, then specifies a Gaussian/linear task model to make optimal prediction and posterior analysis tractable.

3. Pretraining and Data Generative Model

  • At inference, a next-token predictor receives k labeled input-output pairs and a test input, and should predict the test label.The pretraining sequences contain labeled samples drawn from a common task.Source: §3 · p. 4

3.1. Data Generative Model

  • The generative model samples a task from a prior, then draws K i.i.d. labeled examples from that task to form a pretraining sequence.A task is a joint distribution over input-output pairs, and the prior is a distribution over tasks.Source: §3.1 · p. 4

3.2. Bayes-Optimal Next-Token Predictor

  • Under squared loss, the optimal predictor is an MMSE estimator: it predicts a label using the task posterior given the preceding examples and test input.A variable-length predictor decomposes into fixed-length predictors, and the paper reports that a pretrained Transformer can empirically approximate Bayesian inference.Source: §3.2 · p. 4

3.3. Gaussian/Linear Assumptions on Pretraining Data Generative Model

  • The pretraining model extends a shared-input, single-Gaussian task model to Gaussian-mixture task functions and task-varying input distributions.Its assumptions include task groups, varying input means with shared covariance, noisy linear-regression labels, and comparable mixture weights.Source: §3.3 · pp. 4–5

4. Inference and Dual Operating Modes

The inference analysis derives how in-context observations update task posteriors and the optimal prediction. It interprets the balance between mixture-component re-weighting and component shifting as the basis for task retrieval and task learning.

4. Inference and Dual Operating Modes

  • The inference setup treats the in-context sequence and test input as observations for posterior derivation, under assumptions about the task and function generating the examples.The paper uses Gaussian/linear assumptions for the in-context inputs and noiseless labels.Source: §4 · p. 5, §4.1 · p. 5

4.2. Closed-Form Expression of Posterior

  • The task posterior remains a Gaussian mixture, with in-context observations shifting component centers and re-weighting component probabilities.The analysis studies these updates under its stated assumptions to characterize the transition from prior to posterior.Source: §4.2 · p. 6

4.3. Closed-form Expression of ICL Prediction

  • The optimal prediction is a convex combination of predictions from the shifted posterior component centers.The component predictions are weighted by their posterior mixture probabilities.Source: §4.3 · pp. 6–7

4.4. Prior Task Noises, CS, CR, and ICL Prediction

  • In the numerical setting, low task noise and small k let re-weighting boost component 1; as k grows, shifting aligns component centers with the in-context task.The computations track shifting, re-weighting, and prediction risk as the number of examples increases.Source: §4.4 · p. 7

4.5. Dual Operating Modes

  • Task retrieval occurs when re-weighting outweighs shifting, leaving prediction shaped by priors and examples; task learning occurs when shifting dominates and prediction nearly depends on examples alone.These modes describe the relative influence of posterior weights and component centers.Source: §4.5 · p. 7

5. Early Ascent

The section analyzes how posterior reweighting and component shifting shape ICL risk, including conditions under which retrieval can raise risk before task learning lowers it.

5.1. Finegrained Upper Bound

  • The finegrained asymptotic risk bound depends on posterior mixture weights and λ1(A), corresponding to component re-weighting and component shifting.Ignoring re-weighting reduces the finegrained bound to the coarser bound.Source: §5.1 · p. 7

5.2. The Effect of Dual Operating Modes on ICL Risk

  • When pretraining task noise is low, retrieval with few examples helps ICL if the target task is close to a prior center.Retrieval can quickly activate the nearest prior task.Source: §5.2 · p. 7

5.3. Early Ascent with Biased x Distribution

  • Earlier work reports lower GPT-3 one-shot than zero-shot LAMBADA accuracy, followed by higher few-shot accuracy.The reported accuracies are 72.5% one-shot, 76.2% zero-shot, and 86.4% few-shot.Source: §5.3 · p. 7
  • A one-shot translation prompt about an apple leads GPT-3.5 to answer that a banana is yellow, illustrating retrieval of an incorrect skill from misleading input.The paper identifies the activated skill as question answering.Source: §5.3 · p. 7
  • Among three target tasks, the task with the shortest distance to its closest prior is easier to learn than the more distant tasks.Figure 4 labels the closest task blue and the farthest red.Source: Figure 4 · p. 8
  • With few examples, retrieval can activate a misleading task and raise risk; with more examples, task learning can lower it.The reported analysis reproduces an initial increase in both risk and its upper bound when the misleading task is retrieved first.Source: §5.3 · p. 8
  • Early ascent occurs when the misleading task has higher risk than zero-shot risk and δµ and δw are sufficiently small.The analysis establishes this behavior under a stated assumption.Source: §5.3 · p. 8

6. Bounded Efficacy of Biased-Label ICL

The section analyzes biased-label ICL and predicts bounded efficacy: performance can improve initially but worsen as examples accumulate, including with random labels.

6. Bounded Efficacy of Biased-Label ICL

  • The biased-label setting assumes the target function differs from the retrieved prior task while its input distribution is closer to that task’s prior center than to other centers.This proximity supports retrieval of the task associated with that prior center.Source: §6 · pp. 8–9

6.1. Upper Bound for ICL Risk with Biased Labels

  • For sufficiently small δµ and δw, the risk bound predicts exponential decay at small k, followed by an increase when the first term dominates at large k.The theorem considers a next-token predictor attaining optimal pretraining risk.Source: §6.1 · p. 9

6.2. Bounded Efficacy of Biased-Label ICL in GPT-4

  • In GPT-4, the error rate for addition (+) and biased + decreases through k = 2, then increases afterward.The section describes this pattern as bounded efficacy and attributes the later behavior to retrieval of addition (+) before learning biased +.Source: Table 1 · p. 9, §6.2 · p. 9

6.3. Bounded Efficacy for Zero-Shot ICL

  • The model predicts bounded efficacy for zero-shot ICL with random labels, which provide no information.The prediction follows under the stated optimal-pretraining-risk assumption.Source: §6.3 · p. 9
  • In experiments extending beyond 32 examples, random-label ICL error increased at a smaller k than true-label ICL error.The experiments used different LLMs from Min et al. (2022) because context lengths constrained the number of examples.Source: §6.3 · p. 9

7. Conclusion

The paper presents a probabilistic model for task learning and retrieval, uses it to explain early ascent and predict bounded efficacy, and validates its findings with large language models.

7. Conclusion

  • The model explains ICL’s dual operating modes and early ascent, predicts bounded efficacy for biased-label ICL, and is validated with large language models.The conclusion describes this framework as groundwork for further ICL research.Source: §7 · p. 9
  • The framework assumes linear regression pretraining tasks and noiseless in-context labels, unlike complex real-world LLM tasks.The authors identify both as limitations of the current framework.Source: §7 · p. 9

B. Prior Examples

  • The prior examples arrange task centers as regular polyhedrons in 3D or extend configurations to d-dimensional spaces, with a separate early-ascent setup.The appendix presents tetrahedron, octahedron, hexahedron, icosahedron, and dodecahedron configurations.Source: B. Prior Examples · p. 13, §B.1 · p. 14

B.1. Regular Polyhedrons

  • The 3D regular-polyhedron settings use equally weighted mixture components, with µm = wm, and place the in-context task near selected prior centers.The default task-noise scales are σµ = σw = 0.25.Source: §B.1 · p. 14

B.3. Early Ascent Examples

  • The early-ascent setup makes the in-context task’s input distribution closer to a misleading prior task while its input-output mapping matches prior center 2.In the example, retrieval initially favors component 1 and risk rises; with more examples, learning shifts predictions toward component 2 and risk falls.Source: Table 2 · p. 15, Figure 8 · p. 16

B.2. d-Dimensional Examples

  • The d-dimensional examples use d components with equal weights and set each µm to a standard-basis vector, with µm = wm.The tested dimensions are d ∈ {2, 4, 8, 16, 32}.Source: §B.2 · p. 15

C. Coarse Upper Bound for ICL Risk

The paper gives a coarse risk bound, tests whether Transformers approximate Bayesian prediction, and specifies a GPT-4 experiment for learning and retrieval tasks.

C. Coarse Upper Bound for ICL Risk

  • As k grows, the ICL risk upper bound decreases as the square of the inverse of k under noiseless in-context labels.The stated decay is faster than the standard 1/k rate for ridge regression; the experiment reports ICL performing better at small k and ridge regression better when k ≥ d.Source: C. Coarse Upper Bound for ICL Risk · pp. 16–17

E.1. Experimental Setting

  • The GPT-4 experiment design uses the biased-addition task as the in-context task and measures both learning that task and retrieving ordinary addition.Table 3 presents the experiment setting.Source: §E.1 · p. 17, Table 3 · p. 19

D. Transformer Performance in Approximating Bayesian Inference

E.2. Additional Results

Additional results describe biased-label ICL’s shift from task retrieval to learning, while the zero-shot analysis reports bounded efficacy as examples increase.

E.2. Additional Results

  • With biased labels, in-context learning initially retrieves a commonsense pretraining task due to task retrieval, then finally learns the in-context task because of task learning, across k = 0, 2, and 8.Source: §E.2 · p. 19

F. Bounded Efficacy in Zero-shot ICL

G. The Derivation of Posterior

The appendix derives the task posterior by showing how the prior-to-posterior expression yields a Gaussian mixture with componentwise updates.

G. The Derivation of Posterior

  • The derivation establishes that the posterior can remain a Gaussian mixture and explains how the expression in Eq. 2 yields the Gaussian mixture in Eq. 3.The derivation is organized into the prior-to-posterior step and the closed-form componentwise calculation.Source: G. The Derivation of Posterior · p. 21

G.1. Prior to Posterior

  • For fixed in-context observations, the first derivation step shows that the posterior can still be a Gaussian mixture.The appendix introduces this step in the prior-to-posterior subsection.Source: §G.1 · pp. 21–22

G.2. Closed-form Solution from Eq. 2 to Eq. 3

  • The posterior is expressed as Gaussian factors for µ and w, with component re-weighting; for CR, Ψµ, Ψw, and πm trends are examined as k increases under varying task-noise parameters.Source: §G.2 · p. 23, Figure 15 · p. 24

H.1. Analysis of Component Re-weighting

The analyses track how additional in-context examples change posterior component weights and component parameters. Experiments also examine early ascent and bounded efficacy under noisy labels and extend early-ascent tests to nonlinear regression and discrete-token prediction.

H.1. Analysis of Component Re-weighting

  • Smaller task-prior noise makes posterior component dominance easier as examples increase; both reweighting scores initially strengthen but eventually stabilize.For component means and weights, the analysis identifies conditions under which the corresponding score increases faster, while the mean score stabilizes as k grows.Source: §H.1 · pp. 24–25

H.2. Analysis of Component Shifting

  • As k increases, the shifted component parameters ˜µm and ˜wm converge to the in-context task parameters µ∗ and w∗.The analysis reports these limits and presents numerical distances to the target in Figure 16.Source: §H.2 · p. 26, Figure 16 · p. 26

I.1. Early Ascent and Bounded Efficacy under Noisy Labels

  • Early ascent and bounded efficacy persist to some extent under varied label noise, with the reported noise ranges differing between the two phenomena.With pretraining label-noise level σy = 1.0, early ascent persists for τy ∈ [0, 1.0], while bounded efficacy persists for τy ∈ [0, 0.1].Source: §I.1 · p. 28, Figure 17 · p. 28, Figure 18 · p. 29

I.2. Early Ascent under Non-Linear Regression and Discrete Token Prediction

  • Early ascent appears in experiments with both a 2-layer Tanh network for nonlinear regression and sequences of discrete tokens.The experiments use designed pretraining and in-context sample distributions.Source: §I.2 · pp. 28–29

J.1. A Specific Setting of Early Ascent

  • In the specific setting analyzed, the upper bound increases with the number of in-context examples when k is small.The analysis derives this behavior from the growth of r(1, 2) in that setting.Source: §J.1 · p. 32

J.2. Theorem of Early Ascent

The theorem formalizes early ascent under specified assumptions and sufficiently small prior-noise parameters. The appendix also bounds task-learning loss and documents how prediction error changes when in-context inputs span only a subspace.

J.2. Theorem of Early Ascent

  • When δµ and δw are sufficiently small and the theorem’s assumptions hold, the risk exhibits early ascent.The theorem provides a formal result for the phenomenon; the supplied excerpt does not include its full conditions.Source: §J.2 · p. 32

L. ICL to Learn the In-Context Function

  • The proof of the in-context task-learning bound assumes Gaussian inputs around µ∗ and clean linear labels generated by w∗.The target is prediction of ⟨xk+1, w∗⟩ under the Bayes-optimal next-token predictor.Source: L. ICL to Learn the In-Context Function · p. 35

L.1. Case When In-context Input Variable Spans in Subspace

  • In the subspace-input case, the distance between predicted ˜w and w∗ decreases reciprocally with k, while expected squared loss decreases quadratically.The appendix refines the fine-grained risk bound for inputs spanning a subspace of Rd.Source: §L.1 · pp. 36–37

M. ICL with Biased Labels to Retrieve A Task

M.1. Non-Asymptotic Bound for the Bounded Efficacy Phenomenon

The biased-label analysis proves a bounded-efficacy risk bound over a particular interval when prior-noise parameters are sufficiently small. It also provides an asymptotic bound as the number of examples tends to infinity.

M.1. Non-Asymptotic Bound for the Bounded Efficacy Phenomenon

  • For sufficiently small δµ and δw, the non-asymptotic theorem upper-bounds biased-label ICL risk over a particular interval of k.The interval is specified through conditions involving Iµ, Iw, and IU.Source: §M.1 · p. 38, §M.1.5 · p. 44
  • The risk upper bound captures bounded efficacy: its first two terms decay exponentially at small k, while its third term grows at large k.The result concerns an upper bound on squared risk, with a constant depending on the prior setting, τx, and (µ∗, w∗).Source: §M.1 · p. 38

M.2. Asymptotic Bound

  • The asymptotic theorem gives an upper bound on biased-label ICL risk as k tends to infinity.The displayed bound includes constants depending on the prior setting, τx, and (µ∗, w∗).Source: §M.2 · p. 44

N. Proof of Lemma 6.2

The appendix bounds zero-shot ICL risk in a symmetric two-component, zero-label setting and characterizes its bounded-efficacy behavior. Its toy example then casts prediction as posterior averaging and tracks how Gaussian-mixture components shift and reweight with demonstrations.

N. Proof of Lemma 6.2

  • With two symmetric prior components and zero labels, the proof establishes an ICL risk upper bound for a particular interval of k.The setting assumes the in-context inputs are centered at a unit-norm µ* and that the labels are yi = 0.Source: N. Proof of Lemma 6.2 · p. 48
  • For zero-shot ICL, the bound has bounded efficacy: risk decays exponentially at small k but increases when k is large.The third term dominates and increases at large k; the stated interval for the bounded-efficacy bound is the union of Iµ and IU.Source: N. Proof of Lemma 6.2 · p. 48, §N.5 · p. 51
  • The proof splits expected squared loss into three parts, analyzing the A′′, B′′, and C′′ contributions separately.The proof summary combines the bounds for these three parts.Source: N. Proof of Lemma 6.2 · p. 48, §N.1 · p. 49, §N.2 · p. 50

O.1. Toy Example: Pretraing Data Generative Modela

  • The demo generates a task from a prior, samples a sequence conditional on that task, and uses varied tasks and noisy task-specific text as real-world analogies.Each xi is sampled from Dx(µ) after µ is drawn from Dµ.Source: §O.1 · p. 51

O.2. Toy Example: Bayes-Optimal Next-Token Predictor

  • The optimal predictor is an MMSE estimator whose prediction averages task-conditional expectations under the task posterior given the observed sequence.Figure 23 contrasts prediction from the prior without in-context examples with prediction from the posterior when examples are observed.Source: §O.2 · p. 52, Figure 23 · p. 52

O.3. Toy Example: Gaussian Assumptions on Pretraining Data Generative Model

  • The toy model assumes a Gaussian-mixture task prior with shared component covariance and Gaussian token inputs conditional on the task.The mixture components have centers µm and weights πm; the conditional token distribution is centered at µm.Source: §O.3 · p. 52

O.4. Toy Example: Posterior Analysis

  • The closed-form posterior changes both mixture-component centers and weights after observing in-context examples.The paper names these effects Component Shifting and Component Re-weighting.Source: §O.4 · pp. 52–53

Related papers

Loading 2402.18819v2…