Trained Transformers Learn Linear Models In-Context (2306.09927v3)

Abstract: Attention-based neural networks such as transformers have demonstrated a remarkable ability to exhibit in-context learning (ICL): Given a short prompt sequence of tokens from an unseen task, they can formulate relevant per-token and next-token predictions without any parameter updates. By embedding a sequence of labeled training data and unlabeled test data as a prompt, this allows for transformers to behave like supervised learning algorithms. Indeed, recent work has shown that when training transformer architectures over random instances of linear regression problems, these models' predictions mimic those of ordinary least squares. Towards understanding the mechanisms underlying this phenomenon, we investigate the dynamics of ICL in transformers with a single linear self-attention layer trained by gradient flow on linear regression tasks. We show that despite non-convexity, gradient flow with a suitable random initialization finds a global minimum of the objective function. At this global minimum, when given a test prompt of labeled examples from a new prediction task, the transformer achieves prediction error competitive with the best linear predictor over the test prompt distribution. We additionally characterize the robustness of the trained transformer to a variety of distribution shifts and show that although a number of shifts are tolerated, shifts in the covariate distribution of the prompts are not. Motivated by this, we consider a generalized ICL setting where the covariate distributions can vary across prompts. We show that although gradient flow succeeds at finding a global minimum in this setting, the trained transformer is still brittle under mild covariate shifts. We complement this finding with experiments on large, nonlinear transformer architectures which we show are more robust under covariate shifts.

• The paper proves that LSAs, when properly initialized, globally converge to approximate optimal linear predictors for regression tasks.
• It demonstrates that both training and test prompt lengths critically affect ICL performance, with errors diminishing as prompt sizes increase.
• The study reveals transformers' resilience to task and query distribution shifts while exposing brittleness under covariate shifts, guiding future AI improvements.

An Analysis of In-Context Learning Abilities in Transformers with Linear Self-Attention Layers

The paper "Trained Transformers Learn Linear Models In-Context" by Zhang, Frei, and Bartlett provides a detailed paper of the in-context learning (ICL) capabilities of transformer architectures equipped with linear self-attention (LSA) layers. Through this analysis, the authors seek to uncover the mechanisms by which transformers achieve ICL, particularly the ability to form predictions on new tasks by leveraging training examples without parameter updates.

The paper focuses on transformers trained for linear regression tasks. Training involves gradient flow optimization on a population loss over Gaussian-distributed linear models. The authors demonstrate that despite the inherent non-convexity of this setting, gradient flow with specific initial conditions converges to models that effectively mimic ordinary least squares predictions. A key finding is the robust performance against task and query distribution shifts, with ICL performance strongly reflecting the best linear predictor's error.

The main contributions outlined in the paper are:

1. Convergence to Global Optima: The authors prove that for LSAs initialized appropriately, gradient flow converges globally. The trained transformer subsequently achieves prediction errors competitive with the best linear predictor under Gaussian marginals.
2. Impact of Prompt Lengths on Learning and Predictive Performance: A detailed analysis reveals that learning efficacy depends heavily on both training (N) and testing (M) prompt lengths. While convergence improves as N increases, the prediction error behaves as $O(1/M + 1/N^2)$, indicating greater sensitivity to training prompt length.
3. Interaction with Distribution Shifts: The paper examines the impact of various distribution shifts on ICL. Transformers exhibit resilience to task and query shifts, aligning model behavior with prior empirical findings. However, covariate shifts expose brittleness in model predictions, as performance metaphorically collapses when training and testing distributions diverge.
4. Training with Diverse Covariate Distributions: To overcome limitations posed by fixed training covariate distributions, researchers explore models trained on random covariate distributions. While theoretical results imply limitations for LSAs, empirical evaluations of complex transformer variants (e.g., GPT2) indicate enhanced robustness but acknowledge notable gaps in matching traditional least squares' adaptability.

Empirical comparisons with more extensive transformer architectures, such as GPT2, underscore an essential observation: architectural complexity plays a significant role in accommodating covariate shifts, albeit with trade-offs, particularly when evaluated on untrained sequence lengths.

This research has several implications for future AI development. Primarily, it exposes areas where even highly sophisticated models may not align with ideal algorithms, such as ordinary least squares, in robustness. The findings stress the need to enhance models' in-context learning capabilities further, possibly through novel architectures or initialization schemes that afford greater robustness across diverse contextual scenarios.

In conclusion, this in-depth exploration of transformers' in-context learning abilities paves the way for new methodologies to enhance their capacity to handle diverse tasks robustly. These insights provide valuable frameworks for subsequent inquiries into ICL's theoretical underpinnings and potential enhancements for practical applications in artificial intelligence.