Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression

Allan Raventós; Mansheej Paul; Feng Chen; Surya Ganguli

Source-linked AI summary

Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression

Allan Raventós, Mansheej Paul, Feng Chen, Surya Ganguli

arXiv:2306.15063v2Jun 26, 2023cs.LGcs.AIcs.CL

TL;DR

The paper asks whether in-context learning can solve tasks unlike those seen during pretraining and studies this question by varying task diversity in a controlled linear-regression setting. It finds a task-diversity threshold: below it, transformers follow the pretraining-task Bayesian estimator, while beyond it they approach ridge regression and solve new tasks.

Problem
The paper asks whether in-context learning can solve fundamentally new tasks that differ from those represented in pretraining.
Method
The authors vary the number of latent regression vectors in a finite pretraining distribution and evaluate transformers on seen tasks and Gaussian-distributed new tasks.
Results
Transformers exhibit a task-diversity threshold, shifting above it from pretraining-distribution Bayesian behavior toward ridge regression and improved performance on new tasks.
Takeaways & Limitations
Task diversity, alongside dataset and model scale, is critical for the emergence of in-context learning on tasks not seen during pretraining.
Takeaways & Limitations
The findings are established in linear regression, while translating task diversity and these observations to language remains unclear.

Abstract

from arXiv · show

Pretrained transformers exhibit the remarkable ability of in-context learning (ICL): they can learn tasks from just a few examples provided in the prompt without updating any weights. This raises a foundational question: can ICL solve fundamentally $\textit{new}$ tasks that are very different from those seen during pretraining? To probe this question, we examine ICL's performance on linear regression while varying the diversity of tasks in the pretraining dataset. We empirically demonstrate a $\textit{task diversity threshold}$ for the emergence of ICL. Below this threshold, the pretrained transformer cannot solve unseen regression tasks, instead behaving like a Bayesian estimator with the $\textit{non-diverse pretraining task distribution}$ as the prior. Beyond this threshold, the transformer significantly outperforms this estimator; its behavior aligns with that of ridge regression, corresponding to a Gaussian prior over $\textit{all tasks}$, including those not seen during pretraining. Thus, when pretrained on data with task diversity greater than the threshold, transformers $\textit{can}$ optimally solve fundamentally new tasks in-context. Importantly, this capability hinges on it deviating from the Bayes optimal estimator with the pretraining distribution as the prior. This study also explores the effect of regularization, model capacity and task structure and underscores, in a concrete example, the critical role of task diversity, alongside data and model scale, in the emergence of ICL. Code is available at https://github.com/mansheej/icl-task-diversity.

1 Introduction

The paper asks whether in-context learning can solve fundamentally new tasks and investigates how pretraining task diversity affects this ability using linear regression. It identifies a task-diversity threshold beyond which transformers depart from Bayesian behavior tied to the pretraining distribution.

Motivation: The study asks whether ICL can solve tasks fundamentally different from those seen during pretraining and which learning algorithm it implements.The question arises because ICL emerges from next-token pretraining rather than an explicitly specified learning algorithm.
Approach: The authors vary pretraining task diversity in a controlled linear-regression setting to analyze ICL on unseen tasks.Tasks are latent regression vectors, while evaluation uses a Gaussian distribution over all latent regression vectors.
Findings: With low task diversity, transformers behave like Bayesian estimators using the finite pretraining task distribution as their prior.They perform optimally on pretraining tasks but cannot learn new tasks in-context.
Findings: Above a task-diversity threshold, transformers deviate from the pretraining-distribution Bayesian estimator and significantly outperform it on new tasks.At sufficiently large but finite diversity, performance closely matches the optimal estimator for the ideal task distribution.
Implication: The results indicate that ICL emergence cannot be fully explained by Bayesian inference on the pretraining distribution.The paper also examines regularization, model capacity, and task structure.

2 Problem setup

The paper formulates ICL for linear regression as prediction from sequential examples generated by a latent regression task. It pretrains transformers on a finite task distribution, evaluates them on seen and idealized new tasks, and compares them with distribution-specific optimal estimators.

Task formulation: Each ICL task is a latent D-dimensional regression vector, and the transformer predicts targets from sequential data-target examples.At each position, causal attention exposes only the preceding examples and the current input.
Pretraining: Pretraining samples one latent task from a finite uniform set of M regression vectors, so increasing M increases task diversity.The transformer minimizes next-token prediction mean squared error over sequences generated from these tasks.
Evaluation: Evaluation measures performance on new instances of pretraining tasks and on new tasks drawn from TTrue = N(0, I_D).The latter distribution covers all latent regression vectors in the experimental setting.
Optimal estimators: For any task distribution T, the loss-minimizing prediction is the Bayesian posterior mean conditioned on the available context.This provides the optimal-estimator reference for both seen and unseen-task evaluations.
Optimal estimators: For the finite pretraining distribution, the optimal estimator is dMMSE, a likelihood-weighted combination of the pretraining regression vectors.A transformer minimizing pretraining loss is expected to behave like this estimator.
Optimal estimators: For TTrue = N(0, I_D), ridge regression supplies the optimal estimator used to assess performance on new tasks.The comparison distinguishes behavior optimized for the finite pretraining distribution from behavior optimized for the ideal task distribution.

3 Experiments and results

Across linear-regression experiments, pretrained transformers exhibit a task-diversity threshold: below it they behave like Bayesian estimators tied to the pretraining tasks, while above it they approach Ridge and solve new tasks in-context. The threshold is supported by training dynamics, interpolation tests, and scaling analyses, and is affected by dimension, regularization, and model capacity.

Task diversity threshold: At low task diversity, the PT matches dMMSE on pretraining tasks but significantly underperforms Ridge on new tasks.In this regime, it behaves like the Bayesian estimator with prior TPretrain.
Task diversity threshold: Above the task-diversity threshold, the PT deviates from dMMSE, approaches Ridge, and significantly improves on unseen tasks.This improvement occurs alongside a gap from dMMSE on pretraining tasks, indicating that new-task performance does not require minimizing the pretraining objective exactly.
Training and learning dynamics: Increasing sequences per task sharpens the crossover between dMMSE and Ridge, with the PT moving toward Ridge above and away from it below the threshold.The same threshold appears when increasing dataset size through batch size or training steps.
Training and learning dynamics: The threshold is associated with a break in learning dynamics: below it, t* scales approximately as M^0.47, whereas above it the distance to Ridge decreases through 2M steps.The authors interpret this break as evidence for differing learning dynamics rather than under-fitting alone, while noting finite training duration and a decayed learning-rate schedule.
Dimension and model effects: At M = 220, PT performance remains near Ridge across dimensions, while dMMSE increasingly departs from Ridge as dimension grows.The task-diversity threshold itself increases approximately linearly with regression dimension.

4 Related work

The paper extends Bayesian accounts of ICL to controlled linear regression, finding that finite task diversity can produce new-task learning beyond Bayesian inference over the pretraining distribution.

At low task diversity, the transformer’s behavior is consistent with Bayesian inference over tasks learned during pretraining.
Unlike prior linear-regression studies using unlimited task diversity, this work studies new-task ICL with limited, finite pretraining task sets.
At finite but sufficiently large task diversity, ICL emerges because the transformer departs from the estimator optimal for the pretraining distribution.

5 Discussion

The experiments reveal a task-diversity phase transition: beyond an intermediate threshold, transformers solve unseen regression tasks and approach estimators suited to the underlying task distribution. The findings also show that data scale alone is insufficient and that extrapolation to language remains uncertain.

An intermediate task-diversity threshold marks a transition from the dMMSE estimator for TPretrain to Ridge for TTrue, enabling fundamentally new-task ICL.
Beyond the threshold, transformers learn the optimal estimator for the underlying generative task distribution, including Gaussian and Laplace priors.
The implications for language are uncertain because the relevant task categories and task diversity of language corpora are difficult to define.
Increasing dataset size without increasing task diversity can hurt ICL below the threshold, so both data scale and diversity matter.
The Bayesian estimator with task distribution T minimizes expected prediction loss by taking the posterior mean conditioned on the context.
For discrete TPretrain, the optimal estimator is dMMSE, whereas Gaussian TTrue yields the Ridge regression estimator.

B Experimental details

The experiments use transformer models trained on synthetic linear-regression sequences under varied dimensions, optimization settings, model sizes, and compute resources.

Most experiments use D = 8, up to K = 16 in-context examples, noise variance σ^2 = 0.25, and an 8-layer GPT2-style transformer.
The study varies training steps, batch size, model capacity, task dimension, learning rate, and noise scaling across experiments.
The larger-model experiments use 12 layers, 256-dimensional embeddings, and 4 attention heads for dimensions D = 8, 16, 24, 32.
Experiments were implemented in JAX and run on TPU v2-8 and v3-8 hardware, with each training run taking approximately 4 v2-8 TPU hours.

C Support Figure 2

Below the task diversity threshold, increasing batch size aligns the transformer with the optimal estimator for the pretraining distribution, which cannot learn new tasks in-context.

Below the threshold—less than 2^14 tasks—increasing batch size aligns the PT with dMMSE, the optimal estimator for TPretrain.This supports the conclusion that more data does not overcome limited task diversity below the threshold.
Below the task diversity threshold, the optimal estimator approaches the Bayesian estimator with prior TPretrain and cannot learn new tasks.
The visualization extends the x-axis to the task diversity threshold and uses a linear y-axis to distinguish batch sizes at large task counts.

D Dependence on number of sequences per task

The crossover from dMMSE-like to Ridge-like behavior depends primarily on the number of sequences per task, rather than batch size or training steps.

The experiments vary batch size to increase sequences per task while probing the PT’s crossover from dMMSE-like to Ridge-like behavior.
Keeping the number of sequences per task constant makes final PT behavior invariant to batch size or training steps within a reasonable range.The comparison uses batch sizes 256, 512, and 128 with adjusted training steps.
The additional experiments identify sequences per task, rather than batch size or training steps, as the primary factor driving final model behavior.

E Small PT

A smaller transformer reproduces the task diversity threshold, transitioning toward Ridge with more data, but has a lower threshold and lower overall performance than the base model.

The small PT transitions from becoming less like Ridge to becoming more like Ridge when trained on more data.
The small PT has a lower threshold, approximately 2^11.5 versus approximately 2^14.5 for the base PT.
Figure 10 evaluates distance to Ridge on new tasks across increased data and compares finite-diversity models with infinite-task-diversity models.
Figure 11 measures loss barriers and weight-space distances between finite-task and infinite-task models, using infinite-infinite pairs as a baseline.

F Examining the loss landscape along paths interpolating between finite and infinite task models

The loss-landscape analysis compares finite-task models with infinite-task models using interpolation barriers and weight-space distances, with infinite-task pairs as the baseline.

For each finite task count M and M = ∞, four small PTs are trained with different task-sampling seeds while other randomness is shared.
The analysis compares finite-task models with infinite-task models using training-loss barriers under the pretraining objective.
The loss barrier is the midpoint model’s loss minus the average losses of the two endpoint models.
Weight-space distance is the L2 norm of the difference between model weights, averaged across finite-infinite pairs and compared with infinite-infinite baselines.

G Smoothed dMMSE

The Smoothed dMMSE estimator replaces discrete task priors with Gaussian mixtures and has an optimal smoothing level for new-task performance. Beyond 2^13 pretraining tasks, the PT outperforms this optimally smoothed estimator.

The Smoothed dMMSE estimator replaces the discrete point prior over M pretraining tasks with a mixture of M isotropic Gaussians of variance ϵ^2.
Its smoothing parameter ϵ is selected by simulation because both ϵ → 0 and ϵ → ∞ perform poorly on TTrue.
2^13 tasks onward, the PT outperforms the optimally smoothed sMMSE estimator on TTrue.
The sMMSE estimator is derived by computing the posterior mean ẇ_k = E[w | S_k] and applying Eq. (6).

H Supporting figures

Supporting experiments examine regularization, embedding dimension, and task-prior structure. Weight decay lowers the task diversity threshold, whereas increasing embedding dimension has no discernible threshold effect in the base PT.

Regularization: Increasing weight decay consistently lowers the task diversity threshold of the base PT.
Regularization: More weight decay makes the PT’s predictions less similar to Ridge, as measured by ∆TTrue.
Model capacity: Increasing embedding dimension while scaling the number of heads by the same factor has no discernible effect on the threshold.
Task structure: The task-diversity-threshold experiment is also conducted when regression-vector elements are sampled from a Laplace prior instead of a Gaussian prior.