Research map / Architectures and efficiency / Attention and transformers
Nonlinear transformers can perform inference-time feature learning
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 single-layer softmax transformer pretrained by one gradient step plus ridge regression can perform inference-time feature learning for Gaussian single-index models, achieving in-context sample complexity r^{Θ(ge)} that beats kernel and CSQ bounds.
Problem
Whether pretrained transformers can perform in-context learning of single-index functions with per-prompt varying feature vectors, with end-to-end optimization and statistical guarantees, and surpass kernel and CSQ sample complexity.
Approach
Softmax attention's exponential transformation of labels, exp(y/ρ), reduces the effective information exponent of the link function to its generative exponent, so the attention extracts the per-prompt feature direction β from in-context examples, enabling inference-time sample complexity independent of the information exponent.
Key results
- Pretrained GPT-2 in-context sample complexity estimated at roughly d^1.1 for degree-3 single-index task, versus kernel lower bound d^3 and CSQ bound d^1.5.
- Theorem 1: single-layer softmax transformer pretrained with Tpt = Õ(d^{Θ(ie)}) achieves od(1) ICL error with test prompt length Õ(r^{Θ(ge)}), ge ≤ 2.
- Figure 6.1 shows kernel ridge regression error deteriorates rapidly with d while pretrained transformer scales more favorably.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a pretrained Softmax transformer with nonlinear attention instead of fixed-basis kernel methods | Gaussian single-index learning · Gaussian single-index models | inference-time sample complexity: rΘ(deg(σ∗)) -> rΘ(ge(σ∗)) | Helped |
| Use nonlinear Softmax attention instead of a CSQ learner | Gaussian single-index learning · Gaussian single-index models | inference-time sample complexity: dΘ(ie(σ∗)) -> rΘ(ge(σ∗)) | Helped |
| Pretrain GPT-2 for in-context feature learning rather than fit a Gaussian RBF kernel from scratch | Gaussian single-index task with σ∗(·) = He3(·) · Synthetic Gaussian single-index tasks | generalization error versus context length: ddeg(σ∗) = d3 | Helped |
| Pretrain GPT-2 to learn the task’s feature direction in-context | Gaussian single-index task with σ∗(·) = He3(·) · Synthetic Gaussian single-index tasks | inference-time sample-complexity scaling: n ≃ d1.1; kernel lower bound of d3; CSQ lower bound of n ≳ d1.5 | Helped |
Surprising findings
- Inference-time sample complexity is independent of both the degree and the information exponent of the link function, beating CSQ lower bounds despite pretraining with CSQ-like single gradient step.
Limitations
- Analysis restricted to single-layer transformer with one gradient step on attention and fixed link function across tasks.
- Pretraining sample complexity still scales with information exponent d^{Θ(ie)} rather than generative exponent.
- Validation limited to small synthetic experiments with a 6-layer GPT-2 and degree-3 link.
Experimental setup
Datasets: synthetic Gaussian single-index tasks
Benchmarks: synthetic Gaussian single-index task (He3 link)
Metrics: generalization error (squared error) minimal test prompt length to reach error threshold
Baselines: kernel ridge regression Oko et al. 2024b ICL construction
Builds on: Oko et al. 2024b in-context single-index learning Damian et al. 2024 generative exponent Lee et al. 2024 data reuse / generative exponent Garg et al. 2022 ICL of functions
Closest papers by meaning
- In-context Convergence of Transformers (ICML 2024 · similarity 0.90)
- Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context (NeurIPS 2024 · similarity 0.89)
- In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness (NeurIPS 2024 · similarity 0.89)
- Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR 2024 · similarity 0.89)
- Can Transformers Learn Full Bayesian Inference in Context? (ICML 2025 · similarity 0.89)
- Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape (ICML 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.