Research map / Optimization / Implicit bias and generalization
Transformers Implement Functional Gradient Descent to Learn Non-Linear 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 theoretically and empirically that non-linear-attention Transformers can implement functional gradient descent in an RKHS induced by the attention non-linearity, reaching Bayes-optimal prediction for matching kernel Gaussian-process labels.
Problem
How non-linear Transformers (softmax, ReLU attention) learn non-linear functions in context, i.e., which learning algorithm they implement and whether training finds it.
Approach
A simple parameter construction makes a non-linear attention module with activation equal to a kernel K implement functional gradient descent in the RKHS of K; when labels come from a K-Gaussian process, the iterates converge to the Bayes-optimal predictor. Multi-head attention with different activations composes kernels, and stationary-point analysis shows training converges to this construction.
Key results
- When the attention activation matches the generating kernel, the Transformer achieves the lowest in-context loss, close to the Bayes loss, for linear, ReLU and exponential kernel labels.
- A 2-head Transformer with different activations matches the best single-head model on linear and exponential data and outperforms both single-head models on a composite kernel.
- Trained parameters converge toward the predicted stationary point (B^T C approaching identity) across kernel and activation combinations.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Match the attention activation to the label-generating kernel | Single-head Transformer · Linear, ReLU, and exponential-kernel GP data | test in-context loss | Helped |
| Use softmax attention instead of exponential-kernel attention | Single-head Transformer · Exponential-kernel GP data, 3 layers, n=6, 8, 10, 12, or 14 | test in-context loss | Helped |
| Use exponential-kernel attention instead of softmax attention | Single-head Transformer · Exponential-kernel GP data, n=6, at least 5 layers | test in-context loss | Helped |
| Increase the number of layers with exponential-kernel attention | Single-head Transformer · Exponential-kernel GP data, n=14 | test in-context loss | No clear effect |
| Use two heads with linear and exponential-kernel activations rather than either single-head model | Multi-head Transformer · Composite linear-plus-exponential-kernel GP data | test in-context loss | Helped |
| Use two heads rather than one linear-activation head | Multi-head Transformer · Linear-kernel GP data | test in-context loss | No clear effect |
| Use two heads rather than one exponential-activation head | Multi-head Transformer · Exponential-kernel GP data | test in-context loss | No clear effect |
| Use ReLU attention on ReLU-kernel GP data | Single-head Transformer · ReLU-kernel GP data | test in-context loss | Hurt |
Surprising findings
- Softmax attention outperforms exponential-kernel attention on exponential-kernel data when the number of layers is small, though the gap closes with more layers.
- ReLU attention keeps a significant gap to the Bayes estimator that does not shrink with depth.
Limitations
- Proposition on Bayes optimality requires the activation to match a PSD kernel; ReLU and softmax are not kernels and are only tested empirically.
- Interpretation of the unconstrained-stationary-point algorithm (functional GD++) is left as future work.
- Some parameter matrices plateau around 0.2 distance from identity, unclear whether optimization or a different stationary point.
Experimental setup
Datasets: synthetic K-Gaussian process regression data (linear, ReLU, exponential kernels)
Benchmarks: synthetic K-Gaussian process in-context regression
Metrics: log test ICL loss normalized Frobenius distance to identity of B^T C
Baselines: Bayes estimator linear attention ReLU attention softmax attention exponential-kernel attention
Builds on: linear transformer in-context learning construction (gradient descent) preconditioned GD / GD++ analyses Gaussian process label model
Closest papers by meaning
- In-context Learning on Function Classes Unveiled for Transformers (ICML 2024 · similarity 0.92)
- Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context (NeurIPS 2024 · similarity 0.89)
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.89)
- Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR 2024 · similarity 0.88)
- Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions (NeurIPS 2025 · similarity 0.88)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.88)
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.