Research map / Architectures and efficiency / Attention and transformers
Transformers Struggle to Learn to Search
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.
Trains small GPT-2-style transformers on DAG graph-connectivity search with a carefully balanced lookahead distribution, reverse-engineers the learned path-merging algorithm, and shows learning becomes harder with larger graphs and does not improve with model scale or chain-of-thought-style intermediate steps.
Problem
Whether transformers can learn to perform search (graph connectivity / next-step path finding in DAGs) robustly, and whether failures stem from data, parameter count, or architectural limits.
Approach
Uses a balanced DAG generator that makes lookahead uniformly distributed and removes shortcut heuristics, enabling small transformers to learn search; introduces a computation-graph reconstruction method via attention perturbation and activation patching to reveal a parallel exponential path-merging algorithm.
Key results
- Model trained on the balanced distribution performs near-perfectly across lookaheads and generalizes to unobserved lookaheads, while the naive-distribution model fails at larger lookaheads.
- As maximum input graph size increases (with 8 layers, hidden size 16), the probability that models reach accuracy >= 0.995 becomes vanishingly small, and increasing model size does not fix this.
- Depth-first search and selection-inference with chain-of-thought-style intermediate steps are learnable with few layers but still fail to scale to larger graphs; scaling model size does not help.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train on the balanced distribution instead of the naïve distribution | Transformer graph search · Naïve, star, and balanced graph distributions | accuracy | Helped |
| Train on the naïve distribution | Transformer graph search · Graphs with larger lookaheads | accuracy | Hurt |
| Train on the star distribution rather than the balanced distribution | Transformer graph search · Balanced and naïve graph distributions | accuracy | Mixed |
| Train on the balanced distribution with lookahead ≤12 | Transformer graph search · Lookaheads 13, 14, and larger | accuracy | Mixed |
| Increase maximum input graph size | 8-layer transformer, hidden size 16 · Balanced training; naïve-distribution test | test loss | Hurt |
| Increase maximum input graph size | 8-layer transformer, hidden size 16 · Balanced training | convergence rate | Hurt |
| Increase model dimension | Transformer graph search on 31-vertex graphs · Balanced training; naïve-distribution test | training examples to global minimum | No clear effect |
| Use learned token embeddings and rotary positional embeddings (ablation) | Decoder-only transformer graph search · Larger input graphs | test loss | No clear effect |
| Allow depth-first search intermediate steps | Transformer DFS · Larger graphs | test loss | Hurt |
| Increase model scale | Transformer DFS · Fixed maximum graph size of 35 vertices | training FLOPs to learn the task | Hurt |
| Increase model scale | Transformer selection-inference · Fixed maximum graph size of 45 vertices | test loss | No clear effect |
| Add random padding between vertices in the DFS trace (ablation) | Transformer DFS · Larger backtrack distances | accuracy | Helped |
Surprising findings
- Increasing model scale does not alleviate the difficulty of learning search on larger graphs, even with chain-of-thought-style intermediate steps.
- Random padding augmentation enables DFS backtracking to larger distances but does not help learning on larger graphs.
Limitations
- Mechanistic analysis requires many forward passes and is only applied to small models.
- Model does not learn maximal path-merge operations and may use other heuristics on some examples.
- Experiments use small models (hidden size 16, 6-8 layers) and synthetic graphs, so generalization to LLM scale is untested.
Experimental setup
Datasets: synthetic DAG search examples (naive, star, balanced distributions) synthetic DFS traces synthetic selection-inference examples natural-language implicational propositional logic proof search (synthetic)
Benchmarks: held-out synthetic naive and balanced graph search test sets
Metrics: accuracy test cross-entropy loss fraction of seeds converging (training accuracy >= 0.995) proportion of explained examples for path-merging algorithm proportion of maximal path-merge operations
Baselines: naive Erdős–Rényi DAG distribution star distribution selection-inference prompting depth-first search
Builds on: GPT-2 Sophia activation patching RoPE selection-inference prompting
Closest papers by meaning
- Understanding Transformer Reasoning Capabilities via Graph Algorithms (NeurIPS 2024 · similarity 0.89)
- Transformers need glasses! Information over-squashing in language tasks (NeurIPS 2024 · similarity 0.88)
- Do Efficient Transformers Really Save Computation? (ICML 2024 · similarity 0.87)
- Towards an Understanding of Stepwise Inference in Transformers: A Synthetic Graph Navigation Model (ICML 2024 · similarity 0.87)
- A Little Depth Goes a Long Way: The Expressive Power of Log-Depth Transformers (NeurIPS 2025 · similarity 0.87)
- Understanding Transformers via N-Gram Statistics (NeurIPS 2024 · 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.