Research map / Architectures and efficiency / Attention and transformers
In-context Learning on Function Classes Unveiled for Transformers
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.
Theoretical paper showing by construction that transformers with general activation functions can implement approximate gradient descent on n-layer neural networks in-context, enabling learning of linear, indicator, and smooth function classes and algorithm selection.
Problem
Explaining why transformer-based sequence models can learn different function classes in-context, via neural networks as intermediaries.
Approach
Constructs transformers (with general decay-condition activation functions) that implement approximate gradient descent steps on n-layer neural network parameters in context, with resource bounds on heads, hidden dimension, and layers; uses neural network approximation to cover indicator, linear, and smooth function classes, and a two-task indicator construction for algorithm selection.
Key results
- Transformers with O(ε^-2) heads and hidden dimension can implement one approximate gradient descent step on a 2-layer network (Theorem 1); n-layer networks need a_n transformer layers with a_n = O(n^2) per Remark 3.5.
- Learning indicator functions of L2 balls via 3-layer networks gives width O(d^2) versus exponential width for 2-layer networks, so 2-layer bridges make transformer size explode exponentially in 1/δ.
- A transformer pre-trained on linear regression and classification can perform algorithm selection with an attention layer of 6 heads for the indicator (Theorem 6).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a 3-layer rather than a 2-layer neural network to approximate the indicator function (ablation) | Transformer in-context learning of indicator functions of L2 balls · indicator functions of L2 balls | number of attention heads and hidden dimension: O(δ^−1/2) wide neural network; O(1/(δϵ^2)) versus O(exp(δ^−1/4)ϵ^−2) | Helped |
| Use a deep network instead of a shallow network as the approximator (ablation) | Transformer learning smooth functions · smooth functions | number of attention heads and hidden dimension: O(δ^−d/n) versus O(δ^−1/(L−1)) | Helped |
| Control neural-network width rather than depth while keeping nK^2 constant (ablation) | Transformer approximate gradient descent on neural networks · n-layer neural networks | computational resource: nK^2 | Helped |
| Pre-train on both linear regression and classification tasks and select the task-specific algorithm from the prompt | Transformer · linear regression and classification | prediction error: O(ϵ + δ) | Helped |
Surprising findings
- Learning 2-layer neural networks as bridges may require transformer size exponential in the prediction error, so deeper and narrower networks are more resource-efficient.
Limitations
- No experiments; results are purely theoretical upper bounds and constructions.
- Assumes Barron-norm and decay conditions on activations and loss; relies on existence of weight constructions rather than trained models.
Experimental setup
Builds on: Bai et al. 2024 (transformers as algorithm selectors) Siegel & Xu 2020 (neural network approximation rates) Von Oswald et al. 2023 (transformers perform gradient descent in context) Safran & Shamir 2017 (depth-width trade-offs)
Closest papers by meaning
- Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context (NeurIPS 2024 · similarity 0.93)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.92)
- Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR 2024 · similarity 0.92)
- Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context (ICML 2024 · similarity 0.92)
- How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations (ICLR 2024 · similarity 0.89)
- Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context (ICLR 2025 · similarity 0.89)
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.