# On the Role of Depth and Looping for In-Context Learning with Task Diversity

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 36 references that are in the map.
- First posted on arXiv 29 Oct 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)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (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)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Confident Adaptive Language Modeling](https://mlsanity.com/paper/2207.07061/) (2022)
- [Data Distributional Properties Drive Emergent In-Context Learning in Transformers](https://mlsanity.com/paper/2205.05055/) (2022)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers](https://mlsanity.com/paper/2212.10559/) (2022)
- [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)
- [Provable Meta-Learning of Linear Representations](https://mlsanity.com/paper/2002.11684/) (2020)
- [Looped Transformers are Better at Learning Learning Algorithms](https://mlsanity.com/paper/2311.12424/) (2023)
- [Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression](https://mlsanity.com/paper/2306.15063/) (2023)
- [Deep Learning Through the Lens of Example Difficulty](https://mlsanity.com/paper/2106.09647/) (2021)
- [On the Turing Completeness of Modern Neural Network Architectures](https://mlsanity.com/paper/1901.03429/) (2019)
- [General-Purpose In-Context Learning by Meta-Learning Transformers](https://mlsanity.com/paper/2212.04458/) (2022)
- [Uncovering mesa-optimization algorithms in Transformers](https://mlsanity.com/paper/2309.05858/) (2023)
- [Linear attention is (maybe) all you need (to understand transformer optimization)](https://mlsanity.com/paper/2310.01082/) (2023)
- [Meta-Learning without Memorization](https://mlsanity.com/paper/1912.03820/) (2019)
- [Physics of Language Models: Part 1, Learning Hierarchical Language Structures](https://mlsanity.com/paper/2305.13673/) (2023)
- [Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?](https://mlsanity.com/paper/2410.08292/) (2024)
- [Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression](https://mlsanity.com/paper/2310.17086/) (2023)

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

- [Reasoning with Latent Thoughts: On the Power of Looped Transformers](https://mlsanity.com/paper/2502.17416/) (2025)
- [Scaling Latent Reasoning via Looped Language Models](https://mlsanity.com/paper/2510.25741/) (2025)
- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [Recurrent-Depth VLA: Implicit Test-Time Compute Scaling of Vision-Language-Action Models via Latent Iterative Reasoning](https://mlsanity.com/paper/2602.07845/) (2026)
- [Theory of Scaling Laws for In-Context Regression: Depth, Width, Context and Time](https://mlsanity.com/paper/2510.01098/) (2025)
- [Solve the Loop: Attractor Models for Language and Reasoning](https://mlsanity.com/paper/2605.12466/) (2026)
- [Prioritize the Process, Not Just the Outcome: Rewarding Latent Thought Trajectories Improves Reasoning in Looped Language Models](https://mlsanity.com/paper/2602.10520/) (2026)
- [Loop the Loopies!](https://mlsanity.com/paper/2607.16051/) (2026)
- [LT2: Linear-Time Looped Transformers](https://mlsanity.com/paper/2605.20670/) (2026)
- [Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMs](https://mlsanity.com/paper/2602.00576/) (2026)
- [Beyond What Seems Necessary: Hidden Gains from Scaling Training-Time Reasoning Length under Outcome Supervision](https://mlsanity.com/paper/2602.00927/) (2026)

## Abstract (arXiv)

The intriguing in-context learning (ICL) abilities of deep Transformer models have lately garnered significant attention. By studying in-context linear regression on unimodal Gaussian data, recent empirical and theoretical works have argued that ICL emerges from Transformers' abilities to simulate learning algorithms like gradient descent. However, these works fail to capture the remarkable ability of Transformers to learn multiple tasks in context. To this end, we study in-context learning for linear regression with diverse tasks, characterized by data covariance matrices with condition numbers ranging from [1, kappa], and highlight the importance of depth in this setting. More specifically, (a) we show theoretical lower bounds of log(kappa) (or √kappa) linear attention layers in the unrestricted (or restricted) attention setting and, (b) we show that multilayer Transformers can indeed solve such tasks with a number of layers that matches the lower bounds. However, we show that this expressivity of multilayer Transformer comes at the price of robustness. In particular, multilayer Transformers are not robust to even distributional shifts as small as O(e-L) in Wasserstein distance, where L is the depth of the network. We then demonstrate that Looped Transformers -- a special class of multilayer Transformers with weight-sharing -- not only exhibit similar expressive power but are also provably robust under mild assumptions. Besides out-of-distribution generalization, we also show that Looped Transformers are the only models that exhibit a monotonic behavior of loss with respect to depth.