# What Can Transformer Learn with Varying Depth? Case Studies on Sequence Learning Tasks

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

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 41 references that are in the map.
- First posted on arXiv 2 Apr 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Galactica: A Large Language Model for Science](https://mlsanity.com/paper/2211.09085/) (2022)
- [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)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [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)
- [Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective](https://mlsanity.com/paper/2305.15408/) (2023)
- [Transformers Learn Shortcuts to Automata](https://mlsanity.com/paper/2210.10749/) (2022)
- [Chain of Thought Empowers Transformers to Solve Inherently Serial Problems](https://mlsanity.com/paper/2402.12875/) (2024)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [What Algorithms can Transformers Learn? A Study in Length Generalization](https://mlsanity.com/paper/2310.16028/) (2023)
- [Thinking Like Transformers](https://mlsanity.com/paper/2106.06981/) (2021)
- [Tracr: Compiled Transformers as a Laboratory for Interpretability](https://mlsanity.com/paper/2301.05062/) (2023)
- [How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations](https://mlsanity.com/paper/2310.10616/) (2023)
- [Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators?](https://mlsanity.com/paper/2307.14023/) (2023)

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

- [Reasoning with Latent Thoughts: On the Power of Looped Transformers](https://mlsanity.com/paper/2502.17416/) (2025)
- [LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation](https://mlsanity.com/paper/2602.11451/) (2026)
- [In-Context Learning with Representations: Contextual Generalization of Trained Transformers](https://mlsanity.com/paper/2408.10147/) (2024)
- [Out-of-distribution generalization via composition: a lens through induction heads in Transformers](https://mlsanity.com/paper/2408.09503/) (2024)
- [Encode, Think, Decode: Scaling test-time reasoning with recursive latent thoughts](https://mlsanity.com/paper/2510.07358/) (2025)
- [Mixture of Parrots: Experts improve memorization more than reasoning](https://mlsanity.com/paper/2410.19034/) (2024)
- [Towards Understanding Fine-Tuning Mechanisms of LLMs via Circuit Analysis](https://mlsanity.com/paper/2502.11812/) (2025)
- [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)
- [How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression](https://mlsanity.com/paper/2408.04532/) (2024)
- [On the Optimal Memorization Capacity of Transformers](https://mlsanity.com/paper/2409.17677/) (2024)
- [How Transformers Get Rich: Approximation and Dynamics Analysis](https://mlsanity.com/paper/2410.11474/) (2024)
- [Transformers Can Overcome the Curse of Dimensionality: A Theoretical Study from an Approximation Perspective](https://mlsanity.com/paper/2504.13558/) (2025)
- [On the Robustness of Transformers against Context Hijacking for Linear Classification](https://mlsanity.com/paper/2502.15609/) (2025)
- [Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regression](https://mlsanity.com/paper/2508.07571/) (2025)
- [Sparse Growing Transformer: Training-Time Sparse Depth Allocation via Progressive Attention Looping](https://mlsanity.com/paper/2603.23998/) (2026)
- [Dynamic Reasoning Chains through Depth-Specialized Mixture-of-Experts in Transformer Architectures](https://mlsanity.com/paper/2509.20577/) (2025)
- [Limit Analysis for Symbolic Multi-step Reasoning Tasks with Information Propagation Rules Based on Transformers](https://mlsanity.com/paper/2509.23178/) (2025)
- [Performance-Efficiency Tradeoffs in Transformers: An Approximation Theory Perspective](https://mlsanity.com/paper/2510.03784/) (2025)
- [Distributed Multi-Layer Editing for Rule-Level Knowledge in Large Language Models](https://mlsanity.com/paper/2604.08284/) (2026)

## Abstract (arXiv)

We study the capabilities of the transformer architecture with varying depth. Specifically, we designed a novel set of sequence learning tasks to systematically evaluate and comprehend how the depth of transformer affects its ability to perform memorization, reasoning, generalization, and contextual generalization. We show a transformer with only one attention layer can excel in memorization but falls short in other tasks. Then, we show that exhibiting reasoning and generalization ability requires the transformer to have at least two attention layers, while context generalization ability may necessitate three attention layers. Additionally, we identify a class of simple operations that a single attention layer can execute, and show that the complex tasks can be approached as the combinations of these simple operations and thus can be resolved by stacking multiple attention layers. This sheds light on studying more practical and complex tasks beyond our design. Numerical experiments corroborate our theoretical findings.