Research map / Learning theory and algorithms / Sample complexity
Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context
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 a one-layer transformer with a ReLU MLP embedding, pretrained by layer-wise gradient descent, learns Gaussian single-index functions in context with prompt length depending only on the target-subspace dimension r rather than the ambient dimension d.
Problem
In-context learning of nonlinear function classes (single-index models) by transformers: whether pretraining lets the model be statistically more efficient than estimators that only see the test prompt, which need n ≳ d samples.
Approach
Uses a linear-attention transformer with a nonlinear MLP embedding; one gradient step on the MLP weights aligns features with the r-dimensional subspace containing the index vectors, and the attention layer then performs regression on a polynomial basis of size r^{Θ(P)} in context.
Key results
- GPT-2 pretrained on synthetic single-index tasks shows in-context risk that overlaps across ambient dimensions d=16,32,64 for fixed r=8, indicating d-independent in-context sample complexity.
- Pretrained transformer outperforms kernel ridge regression and a two-layer ReLU network trained with Adam for r=8, d=16 and 32 in the experiments.
- Theorem gives ICL risk bound with test prompt length N* ~ r^{Θ(P)} and pretraining T ~ d^{Θ(Q)}.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Pretrain GPT-2 on Gaussian single-index tasks | GPT-2 · Gaussian single-index tasks, r=8, d=16 and 32 | in-context prediction risk | Helped |
| Increase ambient dimension from d=16 to d=32 or 64 while holding r=8 | Pretrained GPT-2 · Gaussian single-index tasks, r=8 | in-context prediction risk | No clear effect |
| Increase target subspace dimension r | Pretrained GPT-2 · Gaussian single-index tasks | in-context sample complexity | Hurt |
| Increase ambient dimension d | Kernel ridge regression and two-layer neural network · Gaussian single-index tasks, r=8, d=16 and 32 | in-context prediction risk | Hurt |
| Use a pretrained transformer rather than learning directly from the test prompt | GPT-2 versus KRR and two-layer NN · Gaussian single-index tasks, r=8, d=16 and 32 | in-context sample complexity | Helped |
Surprising findings
- The in-context sample complexity scales with the target subspace dimension r rather than the ambient dimension d, so kernel and neural-network baselines degrade with d while the transformer does not.
Limitations
- In-context complexity r^{Θ(P)} matches kernel/polynomial regression in r dimensions and may not be optimal.
- Analysis relies on linear attention with removed softmax and a specific parameterization.
- Pretraining complexity depends on the information exponent and uses a one-step gradient analysis; conjectured to be improvable.
Experimental setup
Datasets: synthetic Gaussian single-index regression tasks
Benchmarks: synthetic Gaussian single-index regression
Metrics: in-context prediction risk
Baselines: kernel ridge regression two-layer neural network trained with Adam
Builds on: linear transformer in-context learning analysis single-index model feature learning (information exponent) GPT-2
Closest papers by meaning
- In-context Learning on Function Classes Unveiled for Transformers (ICML 2024 · similarity 0.93)
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.91)
- In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization (NeurIPS 2024 · similarity 0.91)
- How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations (ICLR 2024 · similarity 0.90)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.90)
- Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.