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.

Improving Systematic Generalization of Linear Transformer Using Normalization Layers and Orthogonality Loss Function

The Devil in Linear Transformer

Unified Normalization for Accelerating and Stabilizing Transformers

Linear attention is (maybe) all you need (to understand transformer optimization)

SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization

UN-η: an Offline Adaptive Normalization Method for Deploying Transformers

Transformers without Tears: Improving the Normalization of Self-Attention

Linear Log-Normal Attention with Unbiased Concentration

Your Transformer is Secretly Linear

Understanding the Difficulty of Training Transformers

Linear Transformers are Versatile In-Context Learners

How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?

Unveil Benign Overfitting for Transformer in Vision: Training Dynamics, Convergence, and Generalization

Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context

Parallelizing Linear Transformers with the Delta Rule over Sequence Length

Transformer Normalisation Layers and the Independence of Semantic Subspaces

Trained Transformers Learn Linear Models In-Context

Reality Only Happens Once: Single-Path Generalization Bounds for Transformers

Enhanced Transformer Architecture for Natural Language Processing

Mitigating Over-smoothing in Transformers via Regularized Nonlocal Functionals

Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis