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.

How do we get there? Evaluating transformer neural networks as cognitive models for English past tense inflection

A Comprehensive Comparison of Neural Networks as Cognitive Models of Inflection

Frequency matters: Modeling irregular morphological patterns in Spanish with Transformers

Inflecting when there's no majority: Limitations of encoder-decoder neural networks as cognitive models for German plurals

Different kinds of cognitive plausibility: why are transformers better than RNNs at predicting N400 amplitude?

Can a Transformer Pass the Wug Test? Tuning Copying Bias in Neural Morphological Inflection Models

Neural Networks as Cognitive Models of the Processing of Syntactic Constraints

Evaluating Transformer Models and Human Behaviors on Chinese Character Naming

Cognitive Modeling of Semantic Fluency Using Transformers

The Antecedents of Transformer Models

Why transformers are obviously good models of language

Can Transformers Learn to Solve Problems Recursively?

Transformer Working Memory Enables Regular Language Reasoning and Natural Language Length Extrapolation

Attending to Mathematical Language with Transformers

Revenge of the Fallen? Recurrent Models Match Transformers at Predicting Human Language Comprehension Metrics

Trained Transformers Learn Linear Models In-Context

On the Ability and Limitations of Transformers to Recognize Formal Languages

Learning the Past Tense of English Verbs: The Symbolic Pattern Associator vs. Connectionist Models

On Neurons Invariant to Sentence Structural Changes in Neural Machine Translation

Injecting structural hints: Using language models to study inductive biases in language learning

Emergence of Abstractions: Concept Encoding and Decoding Mechanism for In-Context Learning in Transformers