# How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalization

arXiv: https://arxiv.org/abs/2505.15009 · HTML: https://mlsanity.com/paper/2505.15009/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 33 references that are in the map.
- First posted on arXiv 21 May 2025.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Linear attention is (maybe) all you need (to understand transformer optimization)](https://mlsanity.com/paper/2310.01082/) (2023)
- [Transformers as Support Vector Machines](https://mlsanity.com/paper/2308.16898/) (2023)
- [Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context](https://mlsanity.com/paper/2312.06528/) (2023)
- [Max-Margin Token Selection in Attention Mechanism](https://mlsanity.com/paper/2306.13596/) (2023)
- [JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention](https://mlsanity.com/paper/2310.00535/) (2023)
- [Understanding Factual Recall in Transformers via Associative Memories](https://mlsanity.com/paper/2412.06538/) (2024)
- [On the Optimization and Generalization of Multi-head Attention](https://mlsanity.com/paper/2310.12680/) (2023)
- [Mechanics of Next Token Prediction with Self-Attention](https://mlsanity.com/paper/2403.08081/) (2024)
- [Implicit Bias and Fast Convergence Rates for Self-attention](https://mlsanity.com/paper/2402.05738/) (2024)
- [Superiority of Multi-Head Attention in In-Context Linear Regression](https://mlsanity.com/paper/2401.17426/) (2024)
- [Implicit Optimization Bias of Next-Token Prediction in Linear Models](https://mlsanity.com/paper/2402.18551/) (2024)
- [Non-asymptotic Convergence of Training Transformers for Next-token Prediction](https://mlsanity.com/paper/2409.17335/) (2024)
- [Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers](https://mlsanity.com/paper/2406.03068/) (2024)
- [How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias](https://mlsanity.com/paper/2505.00926/) (2025)
- [Astral Space: Convex Analysis at Infinity](https://mlsanity.com/paper/2205.03260/) (2022)

## Built on by (most cited first, 1 of 1)

- [Emergence of Superposition: Unveiling the Training Dynamics of Chain of Continuous Thought](https://mlsanity.com/paper/2509.23365/) (2025)

## Abstract (arXiv)

We study the approximation capabilities, convergence speeds and on-convergence behaviors of transformers trained on in-context recall tasks -- which requires to recognize the positional association between a pair of tokens from in-context examples. Existing theoretical results only focus on the in-context reasoning behavior of transformers after being trained for the one gradient descent step. It remains unclear what is the on-convergence behavior of transformers being trained by gradient descent and how fast the convergence rate is. In addition, the generalization of transformers in one-step in-context reasoning has not been formally investigated. This work addresses these gaps. We first show that a class of transformers with either linear, ReLU or softmax attentions, is provably Bayes-optimal for an in-context recall task. When being trained with gradient descent, we show via a finite-sample analysis that the expected loss converges at linear rate to the Bayes risks. Moreover, we show that the trained transformers exhibit out-of-distribution (OOD) generalization, i.e., generalizing to samples outside of the population distribution. Our theoretical findings are further supported by extensive empirical validations, showing that without proper parameterization, models with larger expressive power surprisingly fail to generalize OOD after being trained by gradient descent.