REVIEW 11 cited by
Trained Transformers Learn Linear Models In-Context
Not yet reviewed by Pith; the record is open.
This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.
SPECIMEN: schema-true, not a live event
T0 review · schema-true
One-sentence machine reading of the paper's core claim.
pith:XXXXXXXX · record.json · timestamp
read the original 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.
Forward citations
Cited by 11 Pith papers
-
Training with (Swap) Regret Loss in a Single-Layer Self-Attention Model: A Case Study on the Probability Simplex
Training single-layer attention with squared regret loss has stationary points that implement smoothed fictitious play (external regret) and, via a new swap-regret loss, the Blum–Mansour no-swap-regret algorithm.
-
Partition, Prompt, Aggregate: Statistical Self-Consistency in Language Models
LLM probability estimates violate the law of total probability across partitions, and subgroup-aggregated estimates often beat direct population-level estimates (the macro fallacy).
-
Learning Compositional Functions with Transformers from Easy-to-Hard Data
A transformer with O(log k) layers provably learns the k-fold permutation composition task in poly(N,k) samples with curriculum or mixed easy-to-hard data, despite an SQ lower bound requiring N^{Omega(k)} samples on h...
-
Token Sample Complexity of Attention
Attention outputs converge to their infinite-token limit at sub-parametric rates n^−β (with β<1/2) governed by token covariance and attention matrices, and only logarithmically in the hardmax limit.
-
Transformers with RL or SFT Provably Learn Sparse Boolean Functions, But Differently
Under hand-designed masks and task-specific activations, RL fine-tuning learns a k-sparse Boolean reasoning chain in one gradient update while SFT learns it one CoT step per update.
-
How Can Mamba Learn In Context with Outliers and Generalize Provably?
A simplified one-layer Mamba provably learns in-context binary classification tolerating outlier fractions approaching 1, whereas a linear Transformer can only tolerate α < 1/2.
-
Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regression
A one-layer linear-attention transformer implementing noisy gradient descent gives provable bounds showing linear noise plus ensembling avoids label-noise overfitting, and majority voting beats greedy decoding in spar...
-
Solving Empirical Bayes via Transformers
A transformer pre-trained on synthetic Poisson data can beat the classical NPMLE estimator on several empirical Bayes tasks and run about 100x faster.
-
Transformers versus the EM Algorithm in Multi-class Clustering
A pretrained transformer can approximate Lloyd's EM algorithm for multi-class Gaussian clustering and can achieve the minimax optimal clustering error with enough pretraining data.
-
Bigger Is Safer: Provable Robustness in In-Context Learning Scales with Capacity
For linear self-attention Transformers, the paper claims worst-case risk under Wasserstein adversarial shifts is bounded by L0 + C1 ρ√(d/m) + C2 ρ²/√N, giving ρmax∝√m and Nρ−N0∝ρ², but the m-dependence is asserted rat...
-
Only Large Weights (And Not Skip Connections) Can Prevent the Perils of Rank Collapse
A residual self-attention network with all weight entries bounded by a small η can be approximated by one layer to error O(η)‖X‖∞, so skip connections do not prevent layer collapse.
Discussion (0). Continue with ORCID to comment.