Research map / Architectures and efficiency / Attention and transformers
What Can Transformer Learn with Varying Depth? Case Studies on Sequence Learning Tasks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Designs four sequence learning tasks (classification, in-context QA, template matching, in-context template matching) and uses theory plus synthetic experiments to show how many attention layers each requires in attention-only transformers.
Problem
Understanding the minimum number of attention layers an attention-only transformer needs to perform memorization, reasoning, generalization, and contextual generalization on sequence tasks.
Approach
Introduces a family of interconnected sequence tasks of increasing difficulty and proves that one attention layer suffices for memorization but not reasoning or generalization, two layers suffice for reasoning and generalization via copy-matching and parse-mapping, and three layers are constructed for contextual generalization (parsing-copying-matching).
Key results
- A single attention layer memorizes the sequence classification data (given enough heads) but cannot perfectly solve in-context QA or template matching, per the theorems and experiments.
- Two-layer single-head transformers learn in-context QA and template matching, with the experiments showing these tasks are learned with enough training steps.
- A three-layer transformer performs perfectly on in-context template matching while 1- and 2-layer models perform near random guessing.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use one attention layer for sequence classification | Attention-only transformer · Sequence classification | accuracy | Helped |
| use one attention layer instead of two | Attention-only transformer · In-context question answering | accuracy | Hurt |
| use two attention layers instead of one | Attention-only transformer · In-context question answering | accuracy | Helped |
| use one attention layer instead of two | Attention-only transformer · Template matching | accuracy | Hurt |
| use two attention layers instead of one | Attention-only transformer · Template matching | accuracy | Helped |
| use one or two layers instead of three | Attention-only transformer · In-context template matching | accuracy | Hurt |
| increase depth from two to three attention layers | Attention-only transformer · In-context template matching | accuracy | Helped |
| increase depth from three to four attention layers | Attention-only transformer · In-context template matching | learning speed | Helped |
| increase the number of attention heads | Attention-only transformer · Sequence classification | learning speed | Helped |
| increase the number of attention heads | Attention-only transformer · In-context question answering and template matching | performance | No clear effect |
Surprising findings
- Two-layer transformers struggle with in-context template matching in experiments, even failing to train well, though a rigorous proof is left open.
- Adding more attention heads does not significantly help reasoning and generalization tasks.
Limitations
- Analysis restricted to attention-only transformers with MLP removed and ReLU attention; softmax extension only for single head.
- Two-layer failure on contextual generalization is shown empirically, not proven.
- Tasks are synthetic and may not capture real-world language tasks.
Experimental setup
Datasets: synthetic sequence classification data synthetic in-context QA data synthetic template matching data synthetic in-context template matching data
Benchmarks: synthetic memorization synthetic in-context QA synthetic template matching synthetic in-context template matching
Metrics: accuracy loss
Baselines: single-layer transformer two-layer transformer four-layer transformer
Builds on: induction heads Tracr RASP
Closest papers by meaning
- Memorization Capacity of Multi-Head Attention in Transformers (ICLR 2024 · similarity 0.89)
- Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers (Student Abstract) (AAAI 2024 · similarity 0.89)
- Understanding the Expressive Power and Mechanisms of Transformer for Sequence Modeling (NeurIPS 2024 · similarity 0.88)
- An Evolved Universal Transformer Memory (ICLR 2025 · similarity 0.88)
- Born a Transformer -- Always a Transformer? On the Effect of Pretraining on Architectural Abilities (NeurIPS 2025 · similarity 0.87)
- Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.