mlsanityData to 14 Oct 2026

One-Layer Transformer Provably Learns One-Nearest Neighbor In Context

arXiv 2411.10830 ↗First posted 16 Nov 2024cs.LG cs.AI math.OCcited by 16

Not on the map. Why?

Abstract (arXiv)

Transformers have achieved great success in recent years. Interestingly, transformers have shown particularly strong in-context learning capability -- even without fine-tuning, they are still able to solve unseen tasks well purely based on task-specific prompts. In this paper, we study the capability of one-layer transformers in learning one of the most classical nonparametric estimators, the one-nearest neighbor prediction rule. Under a theoretical framework where the prompt contains a sequence of labeled training data and unlabeled test data, we show that, although the loss function is nonconvex when trained with gradient descent, a single softmax attention layer can successfully learn to behave like a one-nearest neighbor classifier. Our result gives a concrete example of how transformers can be trained to implement nonparametric machine learning algorithms, and sheds light on the role of softmax attention in transformer models.

Key facts. Cited by 16 papers in the map (data through 14 Oct 2026). Builds on 33 of its 43 references that are in the map. First posted on arXiv 16 Nov 2024.

Builds on

33 of 43 references.

References · 33

  1. 2017Attention Is All You Needcited by 60,361
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
  4. 2023GPT-4 Technical Reportcited by 18,396
  5. 2021Masked Autoencoders Are Scalable Vision Learnerscited by 7,403
  6. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,546
  7. 2022Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?cited by 1,434
  8. 2022MaskGIT: Masked Generative Image Transformercited by 1,099
  9. 2021An Explanation of In-context Learning as Implicit Bayesian Inferencecited by 938
  10. 2022Transformers learn in-context by gradient descentcited by 732
  11. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 715
  12. 2017Spectrally-normalized margin bounds for neural networkscited by 686
21 more
  1. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 684
  2. 2021Offline Reinforcement Learning as One Big Sequence Modeling Problemcited by 644
  3. 2016Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Conditioncited by 539
  4. 2018A Mean Field View of the Landscape of Two-Layers Neural Networkscited by 466
  5. 2023Larger language models do in-context learning differentlycited by 394
  6. 2023Trained Transformers Learn Linear Models In-Contextcited by 337
  7. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 298
  8. 2023Transformers learn to implement preconditioned gradient descent for in-context learningcited by 279
  9. 2019Stabilizing Transformers for Reinforcement Learningcited by 270
  10. 2023Transformers as Algorithms: Generalization and Stability in In-context Learningcited by 222
  11. 2023In-Context Convergence of Transformerscited by 115
  12. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 102
  13. 2022Vision Transformers provably learn spatial structurecited by 93
  14. 2023How Do Transformers Learn Topic Structure: Towards a Mechanistic Understandingcited by 86
  15. 2022MAE-AST: Masked Autoencoding Audio Spectrogram Transformercited by 82
  16. 2024Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimalitycited by 74
  17. 2020A Mean-field Analysis of Deep ResNet and Beyond: Towards Provable Optimization Via Overparameterization From Depthcited by 56
  18. 2014Rates of Convergence for Nearest Neighbor Classificationcited by 46
  19. 2020Modeling from Features: a Mean-field Framework for Over-parameterized Deep Neural Networkscited by 44
  20. 2024How Well Can Transformers Emulate In-context Newton's Method?cited by 25
  21. 2023A Mechanism for Sample-Efficient In-Context Learning for Sparse Retrieval Taskscited by 17

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 8; 2026: 8.

16 papers.

Sort:

By area:

  1. 2025Lost in Cultural Translation: Do LLMs Struggle with Math Across Cultural Contexts?cited by 9
  2. 2025Decoupling Positional and Symbolic Attention Behavior in Transformerscited by 6
  3. 2025Transformers versus the EM Algorithm in Multi-class Clusteringcited by 5
  4. 2025In-Context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-Separationcited by 4
  5. 2025Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descentcited by 4
  6. 2025Attention-based clusteringcited by 3
  7. 2025Bilevel Data Curation for LLM Fine-tuning: Offline Selection and Online Self-Refining Generationcited by 2
  8. 2025Towards Understanding Transformers in Learning Random Walkscited by 2
  9. 2026Attention-Based Sampler for Diffusion Language Modelscited by 1
  10. 2026Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regressioncited by 1
  11. 2026A Theoretical Analysis of Why Masked Diffusion Models Mitigate the Reversal Cursecited by 0
  12. 2026Attention-based PCAcited by 0
4 more
  1. 2026One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Contextcited by 0
  2. 2026Theory on Attention Dynamics for Out-of-Distribution In-Context Learningcited by 0
  3. 2026Retrieval-Based In-Context Learning: A Domain Adaptation Frameworkcited by 0
  4. 2026Adaptive Mean Estimation by In-Context Learning: A Gradient-Flow Analysiscited by 0