A Closer Look at In-Context Learning under Distribution Shifts

1 Introduction

• The introduction frames ICL as learning from test-time examples without weight updates and asks whether simpler models can reproduce it.It also asks what drives transformer ICL and how resilient ICL is to distribution shifts.Source: §1 · p. 1
• The authors derive conditions under which the optimal prompt-based predictor coincides with OLS or ridge regression.These results provide a theoretical basis for comparing models’ predictions with regression algorithms.Source: §1 · p. 2
• Transformers show stronger ID-ICL than set-based MLPs, which are designed for the task’s permutation-invariant structure.The paper also reports that transformers more closely emulate OLS.Source: §1 · pp. 1–2
• Under mild shifts, transformers degrade more gracefully; under severe shifts, neither architecture exhibits ICL.The paper focuses on test prompts centered on the tail of the training prompt distribution when evaluating OOD-ICL.Source: §1 · p. 2
• ID-ICL performance does not predict OOD-ICL performance for either architecture.The paper therefore treats in-distribution performance and robustness to distribution shifts as distinct evaluation questions.Source: §1 · p. 2
• The authors leave open why transformers outperform set-based MLPs and how either architecture’s OOD-ICL abilities might be improved.Source: §1 · p. 2

2 Related Works

• Proposed accounts of ICL include induction heads that copy in-context information, implicit gradient descent during inference, and implicit Bayesian inference under hidden Markov model data.Other work analyzes learnability bounds and conditions for ICL when prompts have low training-distribution probability.Source: §2 · p. 2
• Prior studies connect ICL generalization to pretraining, state size, and model size, including large models’ ability to override semantic priors.The passage contrasts rule-based generalization in pretrained transformers with exemplar-based generalization in models trained from scratch.Source: §2 · p. 2

Optimal predictors and scope of guarantees

• Under square loss, the optimal model predicts E[y_j|p_j], and an equivalent statement holds for cross-entropy loss.Source: §3 · p. 3
• When Σ is identity, the noisy-setting solution coincides with ridge regression using σ² as the ridge penalty.Source: §3 · p. 3
• With Σ as identity, the noiseless limit yields the optimal predictor involving X_j’s Moore-Penrose pseudoinverse.The section obtains this limit by taking σ to zero.Source: §3 · p. 4
• For arbitrary β and input distributions with finite mean and invertible covariance, the stated solution requires Constraint 1, which restricts linearity in the query and observed labels.Without that constraint, the optimal model may not coincide with OLS or ridge regression.Source: §3 · p. 4
• Outside the training-prompt support, performance is not guaranteed and depends on inductive biases; the experiments examine the architecture’s role.On the training support, the section characterizes conditions under which the optimal model emulates OLS or ridge regression.Source: §3 · p. 4

Set-based MLP architecture

• The set-based architecture estimates β from unordered examples with MLPs ρ and ϕ, then uses a final MLP ψ to approximate a linear map for variable-length prompts.The parameter estimate is permutation-invariant with respect to the input order.Source: §3 · pp. 4–5

Distribution shifts

• Covariate shift over prompts changes the prompt distribution while leaving the conditional response invariant; the analysis emphasizes shifts toward low-probability training sequences.Changing the label-generation process so that the conditional response changes does not qualify as covariate shift over prompts.Source: §3 · p. 5

4 Experiments

• Test inputs shifted from N(0, I_d) training to N(µ, I), with µ=2·1 or 4·1; experiments used d=10 and compared MLP-set depths with GPT-2-family transformers.The experiments included noiseless labels and noisy labels with σ²=1.Source: §4 · pp. 5–6, Figure 1 · p. 6, Figure 2 · p. 6
• Across noiseless and noisy tests, MLP-set shows ID-ICL but trails transformers; transformers degrade more gracefully under mild shifts, while neither model shows OOD-ICL under severe shifts.Figure 1 shows the noiseless setting and Figure 2 the noisy setting.Source: Figure 1 · p. 6, Figure 2 · p. 6, §4 · p. 6
• ID-ICL performance rankings for either architecture do not predict their OOD-ICL abilities.Source: §4 · p. 6

5 Discussion

• Transformers outperform natural baselines in approximating OLS and ridge regression under mild distribution shifts.Source: §5 · p. 7
• The discussion leaves open why transformers outperform set-based MLPs and how their OOD-ICL abilities might be improved.Source: §5 · p. 7
• Whether the comparisons extend beyond OLS and set-based MLPs remains to be explored.Source: §5 · p. 7

• Under square loss, the risk-minimizing predictor is E[y_j|p_j] almost everywhere.Source: A Appendix · p. 9
• For Gaussian linear-regression prompts, the appendix derives the optimal predictor by computing the posterior distribution of β conditioned on prefix p_j.Source: A Appendix · pp. 9–10
• The zero-noise theorem assumes square loss and identity covariance, and identifies X_j^+ as the Moore-Penrose pseudo-inverse.Source: A Appendix · p. 10
• If X_jᵀX_j is noninvertible, the ridge-regression limit may not converge to the Moore-Penrose pseudo-inverse.Source: A Appendix · p. 10
• Under Constraint 1, square-loss risk depends on the prior over β through its mean and covariance, so a Gaussian prior with matching moments yields the same loss.Source: A Appendix · pp. 11–12