Research map / Optimization / Implicit bias and generalization

In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization

NeurIPS 2024 · Optimization › Implicit bias and generalization · approach: In-context learning as gradient descent · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows that a linear transformer block (linear attention plus linear MLP) achieves nearly Bayes-optimal in-context linear regression under a non-zero-mean Gaussian task prior, while linear attention alone has an irreducible approximation gap, and that LTB is equivalent to one-step GD with learnable initialization trainable by gradient flow.

Problem

In-context learning of linear regression with a Gaussian task prior having non-zero mean (shared signal); prior theory only covered zero-mean priors where linear self-attention suffices.

Approach

Introduce a linear transformer block (LSA followed by a linear MLP) and show it implements one-step GD with learnable initialization (GD-beta); LSA alone incurs an approximation gap proportional to the squared norm of the task mean, while the MLP removes it, and gradient flow on the GD-beta class converges globally despite non-convexity.

Key results

• LSA ICL risk exceeds the LTB optimum by an additive gap of at least max{2/(3(M+1)), (tr(HPsi)+sigma^2)^2/((M+1)^2 tr((HPsi)^2))} * ||beta*||_H^2, an irreducible gap when the task mean is large.
• The optimal GD-beta and LTB achieve ICL risk matching Bayes-optimal up to constant factors when tr(HPsi) is at most of order sigma^2.
• On GPT2 with beta*=(10,...,10), loss was 0.013 with MLP versus 8.871 without MLP; with beta*=0 losses were 0.003 vs 0.024.

Measured effects reported in the paper

Surprising findings

• GPT2 without MLP reaches loss 8.871 when the task mean is (10,...,10), versus 0.013 with MLP.
• Every optimal LTB is effectively a GD-beta model, so the richer LTB class adds no approximation error beyond GD-beta.

Limitations

• Analysis restricted to linear transformer block (no softmax, no ReLU) and population risk with infinite tasks.
• Statistical complexity of pretraining LTB or GD-beta is not studied.
• Optimization of the full LTB parameterization is not addressed.

Experimental setup

Datasets: synthetic Gaussian linear regression ICL tasks

Benchmarks: synthetic linear regression ICL with d=20, M=40

Metrics: ICL risk (mean squared error) approximation gap loss

Baselines: linear self-attention (LSA) GD-zero models GPT2-noMLP Bayes optimal estimator

Builds on: linear self-attention ICL theory one-step GD zero initialization (GD-zero) Bayes optimal ICL for linear regression

Closest papers by meaning

• Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.91)
• Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context (NeurIPS 2024 · similarity 0.91)
• Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.90)
• Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond (NeurIPS 2024 · similarity 0.89)
• Transformers are almost optimal metalearners for linear classification (NeurIPS 2025 · similarity 0.88)
• One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention (ICLR 2024 · similarity 0.88)