# In-Context Learning with Representations: Contextual Generalization of Trained Transformers

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

## Key facts

- Cited by 38 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 39 references that are in the map.
- First posted on arXiv 19 Aug 2024.

## Builds on (most cited first, 30 of 34)

- [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)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016)
- [Larger language models do in-context learning differently](https://mlsanity.com/paper/2303.03846/) (2023)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [Many-Shot In-Context Learning](https://mlsanity.com/paper/2404.11018/) (2024)
- [Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers](https://mlsanity.com/paper/2212.10559/) (2022)
- [Inductive Biases and Variable Creation in Self-Attention Mechanisms](https://mlsanity.com/paper/2110.10090/) (2021)
- [The Learnability of In-Context Learning](https://mlsanity.com/paper/2303.07895/) (2023)
- [What and How does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization](https://mlsanity.com/paper/2305.19420/) (2023)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?](https://mlsanity.com/paper/2310.08391/) (2023)
- [A Theory of Emergent In-Context Learning as Implicit Structure Induction](https://mlsanity.com/paper/2303.07971/) (2023)
- [In-Context Learning through the Bayesian Prism](https://mlsanity.com/paper/2306.04891/) (2023)
- [How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations](https://mlsanity.com/paper/2310.10616/) (2023)
- [Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality](https://mlsanity.com/paper/2402.19442/) (2024)
- [Global Convergence of Deep Networks with One Wide Layer Followed by Pyramidal Topology](https://mlsanity.com/paper/2002.07867/) (2020)
- [Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape](https://mlsanity.com/paper/2402.01258/) (2024)

## Built on by (most cited first, 30 of 38)

- [Preconditioning Benefits of Spectral Orthogonalization in Muon](https://mlsanity.com/paper/2601.13474/) (2026)
- [On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures](https://mlsanity.com/paper/2410.11778/) (2024)
- [Short-length Adversarial Training Helps LLMs Defend Long-length Jailbreak Attacks: Theoretical and Empirical Evidence](https://mlsanity.com/paper/2502.04204/) (2025)
- [How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias](https://mlsanity.com/paper/2505.00926/) (2025)
- [Faster WIND: Accelerating Iterative Best-of-N Distillation for LLM Alignment](https://mlsanity.com/paper/2410.20727/) (2024)
- [On the Role of Transformer Feed-Forward Layers in Nonlinear In-Context Learning](https://mlsanity.com/paper/2501.18187/) (2025)
- [In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learning](https://mlsanity.com/paper/2409.12293/) (2024)
- [Out-of-Distribution Generalization of In-Context Learning: A Low-Dimensional Subspace Perspective](https://mlsanity.com/paper/2505.14808/) (2025)
- [Transformers Meet In-Context Learning: A Universal Approximation Theory](https://mlsanity.com/paper/2506.05200/) (2025)
- [On the Convergence of Gradient Descent on Learning Transformers with Residual Connections](https://mlsanity.com/paper/2506.05249/) (2025)
- [When and How Unlabeled Data Provably Improve In-Context Learning](https://mlsanity.com/paper/2506.15329/) (2025)
- [Decoupling Positional and Symbolic Attention Behavior in Transformers](https://mlsanity.com/paper/2511.11579/) (2025)
- [Can LLMs Reason Over Non-Text Modalities in a Training-Free Manner? A Case Study with In-Context Representation Learning](https://mlsanity.com/paper/2509.17552/) (2025)
- [In-Context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-Separation](https://mlsanity.com/paper/2502.08136/) (2025)
- [Provable In-Context Learning of Nonlinear Regression with Transformers](https://mlsanity.com/paper/2507.20443/) (2025)
- [Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descent](https://mlsanity.com/paper/2510.10425/) (2025)
- [Provable Benefit of Curriculum in Transformer Tree-Reasoning Post-Training](https://mlsanity.com/paper/2511.07372/) (2025)
- [On Theoretical Interpretations of Concept-Based In-Context Learning](https://mlsanity.com/paper/2509.20882/) (2025)
- [SpA2V: Harnessing Spatial Auditory Cues for Audio-driven Spatially-aware Video Generation](https://mlsanity.com/paper/2508.00782/) (2025)
- [Review of Hallucination Understanding in Large Language and Vision Models](https://mlsanity.com/paper/2510.00034/) (2025)
- [How Can Mamba Learn In Context with Outliers and Generalize Provably?](https://mlsanity.com/paper/2510.00399/) (2025)
- [Transformers Provably Learn Directed Acyclic Graphs via Kernel-Guided Mutual Information](https://mlsanity.com/paper/2510.25542/) (2025)
- [Learning to Adapt: In-Context Learning Beyond Stationarity](https://mlsanity.com/paper/2604.10946/) (2026)
- [Learning-to-Context Slope: Evaluating In-Context Learning Effectiveness Beyond Performance Illusions](https://mlsanity.com/paper/2506.23146/) (2025)
- [Multi-Layer Attention is the Amplifier of Demonstration Effectiveness](https://mlsanity.com/paper/2508.00385/) (2025)
- [In-Context Learning for Non-Stationary MIMO Equalization](https://mlsanity.com/paper/2510.08711/) (2025)
- [Understanding and Improving Continuous Adversarial Training for LLMs via In-context Learning Theory](https://mlsanity.com/paper/2604.12817/) (2026)
- [When Does Context Help? Error Dynamics of Contextual Information in Large Language Models](https://mlsanity.com/paper/2602.08294/) (2026)
- [Covariance-Aware Transformers for Quadratic Programming and Decision Making](https://mlsanity.com/paper/2602.14506/) (2026)
- [Demonstrations, CoT, and Prompting: A Theoretical Analysis of ICL](https://mlsanity.com/paper/2603.19611/) (2026)

## Abstract (arXiv)

In-context learning (ICL) refers to a remarkable capability of pretrained large language models, which can learn a new task given a few examples during inference. However, theoretical understanding of ICL is largely under-explored, particularly whether transformers can be trained to generalize to unseen examples in a prompt, which will require the model to acquire contextual knowledge of the prompt for generalization. This paper investigates the training dynamics of transformers by gradient descent through the lens of non-linear regression tasks. The contextual generalization here can be attained via learning the template function for each task in-context, where all template functions lie in a linear space with m basis functions. We analyze the training dynamics of one-layer multi-head transformers to in-contextly predict unlabeled inputs given partially labeled prompts, where the labels contain Gaussian noise and the number of examples in each prompt are not sufficient to determine the template. Under mild assumptions, we show that the training loss for a one-layer multi-head transformer converges linearly to a global minimum. Moreover, the transformer effectively learns to perform ridge regression over the basis functions. To our knowledge, this study is the first provable demonstration that transformers can learn contextual (i.e., template) information to generalize to both unseen examples and tasks when prompts contain only a small number of query-answer pairs.