# How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations

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

## Key facts

- Cited by 75 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 42 references that are in the map.
- First posted on arXiv 16 Oct 2023.

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

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [Calibrate Before Use: Improving Few-Shot Performance of Language Models](https://mlsanity.com/paper/2102.09690/) (2021)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity](https://mlsanity.com/paper/2104.08786/) (2021)
- [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)
- [Learning To Retrieve Prompts for In-Context Learning](https://mlsanity.com/paper/2112.08633/) (2021)
- [Convex Optimization: Algorithms and Complexity](https://mlsanity.com/paper/1405.4980/) (2014)
- [MetaICL: Learning to Learn In Context](https://mlsanity.com/paper/2110.15943/) (2021)
- [Causal Abstractions of Neural Networks](https://mlsanity.com/paper/2106.02997/) (2021)
- [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)
- [Data Distributional Properties Drive Emergent In-Context Learning in Transformers](https://mlsanity.com/paper/2205.05055/) (2022)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Noisy Channel Language Model Prompting for Few-Shot Text Classification](https://mlsanity.com/paper/2108.04106/) (2021)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression](https://mlsanity.com/paper/2306.15063/) (2023)
- [Impact of Pretraining Term Frequencies on Few-Shot Reasoning](https://mlsanity.com/paper/2202.07206/) (2022)

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

- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [Can large language models explore in-context?](https://mlsanity.com/paper/2403.15371/) (2024)
- [From Words to Numbers: Your Large Language Model Is Secretly A Capable Regressor When Given In-Context Examples](https://mlsanity.com/paper/2404.07544/) (2024)
- [Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs](https://mlsanity.com/paper/2410.13835/) (2024)
- [Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape](https://mlsanity.com/paper/2402.01258/) (2024)
- [Why Larger Language Models Do In-context Learning Differently?](https://mlsanity.com/paper/2405.19592/) (2024)
- [In-Context Learning with Representations: Contextual Generalization of Trained Transformers](https://mlsanity.com/paper/2408.10147/) (2024)
- [Transformers are Minimax Optimal Nonparametric In-Context Learners](https://mlsanity.com/paper/2408.12186/) (2024)
- [Bayesian Teaching Enables Probabilistic Reasoning in Large Language Models](https://mlsanity.com/paper/2503.17523/) (2025)
- [How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?](https://mlsanity.com/paper/2402.15607/) (2024)
- [The Mystery of In-Context Learning: A Comprehensive Survey on Interpretation and Analysis](https://mlsanity.com/paper/2311.00237/) (2023)
- [In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization](https://mlsanity.com/paper/2402.14951/) (2024)
- [Iteration Head: A Mechanistic Study of Chain-of-Thought](https://mlsanity.com/paper/2406.02128/) (2024)
- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [A phase transition between positional and semantic learning in a solvable model of dot-product attention](https://mlsanity.com/paper/2402.03902/) (2024)
- [Linear Transformers are Versatile In-Context Learners](https://mlsanity.com/paper/2402.14180/) (2024)
- [A Percolation Model of Emergence: Analyzing Transformers Trained on a Formal Language](https://mlsanity.com/paper/2408.12578/) (2024)
- [How Well Can Transformers Emulate In-context Newton's Method?](https://mlsanity.com/paper/2403.03183/) (2024)
- [In-context Time Series Predictor](https://mlsanity.com/paper/2405.14982/) (2024)
- [In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness](https://mlsanity.com/paper/2402.11639/) (2024)
- [Bypassing the Exponential Dependency: Looped Transformers Efficiently Learn In-context by Multi-step Gradient Descent](https://mlsanity.com/paper/2410.11268/) (2024)
- [How Do LLMs Perform Two-Hop Reasoning in Context?](https://mlsanity.com/paper/2502.13913/) (2025)
- [What Can Transformer Learn with Varying Depth? Case Studies on Sequence Learning Tasks](https://mlsanity.com/paper/2404.01601/) (2024)
- [Understanding the Language Model to Solve the Symbolic Multi-Step Reasoning Problem from the Perspective of Buffer Mechanism](https://mlsanity.com/paper/2405.15302/) (2024)
- [Intrinsic Self-correction for Enhanced Morality: An Analysis of Internal Mechanisms and the Superficial Hypothesis](https://mlsanity.com/paper/2407.15286/) (2024)
- [On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures](https://mlsanity.com/paper/2410.11778/) (2024)
- [A Survey to Recent Progress Towards Understanding In-Context Learning](https://mlsanity.com/paper/2402.02212/) (2024)
- [Reinforced In-Context Black-Box Optimization](https://mlsanity.com/paper/2402.17423/) (2024)
- [Transformers Learn Low Sensitivity Functions: Investigations and Implications](https://mlsanity.com/paper/2403.06925/) (2024)
- [How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression](https://mlsanity.com/paper/2408.04532/) (2024)

## Abstract (arXiv)

While large language models based on the transformer architecture have demonstrated remarkable in-context learning (ICL) capabilities, understandings of such capabilities are still in an early stage, where existing theory and mechanistic understanding focus mostly on simple scenarios such as learning simple function classes. This paper takes initial steps on understanding ICL in more complex scenarios, by studying learning with representations. Concretely, we construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through a possibly complex but fixed representation function, composed with a linear function that differs in each instance. By construction, the optimal ICL algorithm first transforms the inputs by the representation function, and then performs linear ICL on top of the transformed dataset. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size. Empirically, we find trained transformers consistently achieve near-optimal ICL performance in this setting, and exhibit the desired dissection where lower layers transforms the dataset and upper layers perform linear ICL. Through extensive probing and a new pasting experiment, we further reveal several mechanisms within the trained transformers, such as concrete copying behaviors on both the inputs and the representations, linear ICL capability of the upper layers alone, and a post-ICL representation selection mechanism in a harder mixture setting. These observed mechanisms align well with our theory and may shed light on how transformers perform ICL in more realistic scenarios.