Research map / Architectures and efficiency / Attention and transformers
One-Layer Transformer Provably Learns One-Nearest Neighbor 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 that a single frozen-value softmax attention layer trained with gradient descent on a 1-nearest-neighbor ICL task converges to zero loss despite nonconvexity, and that the learned model behaves like 1-NN under distribution shift.
Problem
Whether transformers can provably learn a nonparametric algorithm (one-nearest neighbor) in context under gradient-based training, given the nonconvex loss landscape of softmax attention.
Approach
Reduces the gradient dynamics of a one-layer softmax attention model (frozen value matrix, merged query-key matrix, specific initialization) to a two-parameter dynamical system, proves loss convergence to zero, and shows robustness of the learned predictor under prompt/query distribution shift.
Key results
- With initialization sigma above a stated threshold and N >= O(sqrt(d) log d), gradient descent drives the training loss L(W) to zero (theorem convergence result).
- Under distribution shift, the test MSE to the 1-NN predictor decays faster than training loss in experiments with SGD and random Gaussian initialization, for N in {16,32,64} and d in {8,16}.
- Theory bounds the expected squared error to the 1-NN predictor by O(R^2 N^2 K^{-poly(N,d) delta} + R^2 P(A_delta^c)) under distribution shift.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train the one-layer softmax-attention transformer by gradient descent with the specified initialization and sufficiently large σ | One-layer softmax-attention transformer · 1-NN in-context learning | training loss: converges to 0 | Helped |
| Train the transformer for K gradient-descent iterations under the stated margin and bounded-label conditions | One-layer softmax-attention transformer · Distribution-shifted 1-NN prediction | squared prediction error: O\big(\inf_{\delta}\big\{R^2N^2K^{-\poly(N,d)\delta} + R^2\PP^{\text{test}}(A_\delta^c)\big\}\big) | Helped |
| Train for enough gradient steps when test queries are separated from the decision boundary | One-layer softmax-attention transformer · Distribution-shifted classification with boundary separation | classification error after rounding: 0 | Helped |
| Use SGD and random Gaussian initialization instead of the theorem's specified initialization and full-batch gradient descent | Single-layer transformer · 1-NN in-context learning | training loss | Helped |
| Increase context length N or input dimension d | Single-layer transformer trained on 1-NN · 1-NN in-context learning | convergence speed | Hurt |
| Evaluate on a test distribution whose queries are separated from the decision boundary | Trained single-layer transformer · Well-separated distribution-shifted test set | MSE to 1-NN predictor | Helped |
Surprising findings
- Experiments indicate the model still learns 1-NN with SGD and random initialization, beyond the theoretical initialization assumption.
Limitations
- Analysis restricted to one-layer attention with frozen value matrix and a specific zero/sigma initialization.
- Data distribution assumes uniform sphere inputs and zero-mean uncorrelated labels; convergence result requires large sigma.
Experimental setup
Datasets: synthetic 1-NN prompts sampled uniformly on the sphere
Benchmarks: synthetic 1-NN prompts (training distribution) synthetic well-separated test distribution
Metrics: MSE loss MSE to 1-NN predictor
Builds on: linear self-attention in-context learning theory gradient flow analysis of transformers
Closest papers by meaning
- One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought (NeurIPS 2026 · similarity 0.90)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.89)
- Non-asymptotic Convergence of Training Transformers for Next-token Prediction (NeurIPS 2024 · similarity 0.88)
- Transformer Learns Optimal Variable Selection in Group-Sparse Classification (ICLR 2025 · similarity 0.88)
- Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? (ICLR 2024 · similarity 0.88)
- In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness (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.