Research map / Graph learning / Knowledge and temporal graphs
Towards an Understanding of Stepwise Inference in Transformers: A Synthetic Graph Navigation Model
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.
Studies stepwise inference (scratchpad/CoT-like) in small autoregressive Transformers trained on a synthetic graph path-existence/navigation task, reproducing gaps, a diversity-accuracy tradeoff, simplicity bias, and in-context controllability effects.
Problem
Stepwise inference protocols such as scratchpads and chain-of-thought improve LLM performance on multi-step problems, but their underlying mechanisms are poorly understood.
Approach
Casts multi-step reasoning as graph navigation on synthetic Bernoulli and hierarchical DAGs, training small GPT-style Transformers from scratch with and without intermediate path tokens, and uses this abstraction to reproduce scale-observed phenomena and give mechanistic hypotheses (approximate distance-based greedy navigation).
Key results
- Stepwise inference yields a clear accuracy gain over direct inference on Bernoulli and hierarchical DAGs, larger on hierarchical graphs, and the gap grows as training path lengths shrink.
- Increasing sampling temperature from 0.0 to 3.0 increases the number of unique valid paths but lowers accuracy, revealing a diversity-accuracy tradeoff.
- A simplified attention-based distance-greedy algorithm matches the full model's accuracy and yields over 75% identical paths to the trained model's outputs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use stepwise inference instead of direct inference | Transformer graph navigation · Bernoulli DAG and hierarchical DAG | path-classification accuracy | Helped |
| Limit training paths to shorter lengths | Stepwise versus direct inference · Hierarchical DAG; test paths with length at least the training threshold | stepwise inference gap in classification accuracy | Helped |
| Increase sampling temperature | Autoregressive Transformer · Fixed start and goal node pair | path diversity | Helped |
| Increase sampling temperature | Autoregressive Transformer · Fixed start and goal node pair | path accuracy | Hurt |
| Use stepwise inference with randomly corrupted training tokens | Transformer graph navigation · Single-graph scenario with 5%, 10%, and 20% token corruption | stepwise inference gap | Helped |
| Increase graph edge density | Stepwise versus direct inference · Hierarchical DAG, edge density from 0.08 to 0.12 | stepwise inference gap | Hurt |
| Generate paths at higher sampling temperatures | Trained Transformer · Bernoulli DAG | generated path length relative to ground-truth path length | Helped |
| Use in-context exemplars specifying a chain of motifs | Transformer motif navigation · Unseen orders of motifs, up to the training chain length | successful navigation | Helped |
| Increase the number of chained motifs beyond the maximum seen in training | Transformer motif navigation · Longer motif chains | successful steering | Hurt |
| Place conflicting motif-chain exemplars in the prompt | Transformer motif navigation · Two chains with the same initial and terminal motifs | selection of the first chain | Helped |
| Use the simplified algorithm extracted from the attention-only model | Full trained Transformer · Held-out node pairs | classification accuracy | No clear effect |
| Vary embedding dimensionality in the one-layer attention-only model (ablation) | One-layer Transformer without MLP or post-LayerNorm · Bernoulli DAG with 200 nodes and $p=0.05$ | generalization to held-out node pairs | Helped |
Surprising findings
- Models generated paths are biased toward shorter-than-ground-truth paths (simplicity bias), which is traced to the goal-distance logit structure.
- Under conflicting in-context exemplars, the model shows a primacy bias, preferring the path defined by the first chain.
Limitations
- Results are on a small synthetic task with a 2-layer Transformer trained from scratch, not validated at LLM scale.
- Assumes perfect skills so failures are attributed only to planning.
- Hypotheses about LLMs (compositional generalization limits, primacy bias at scale) are not tested in large models.
Experimental setup
Datasets: synthetic Bernoulli DAGs synthetic hierarchical DAGs synthetic motif-chain exemplar sequences
Benchmarks: held-out node pairs in synthetic graphs
Metrics: classification accuracy (path / no path) path validity unique valid path diversity path length misstep rate planning failure rate
Baselines: direct inference (no intermediate steps)
Builds on: chain-of-thought scratchpad nanoGPT
Closest papers by meaning
- Emergence of Superposition: Unveiling the Training Dynamics of Chain of Continuous Thought (ICLR 2026 · similarity 0.87)
- Discrete Diffusion Models Exploit Asymmetry to Solve Lookahead Planning Tasks (NeurIPS 2026 · similarity 0.87)
- Reasoning by Superposition: A Theoretical Perspective on Chain of Continuous Thought (NeurIPS 2025 · similarity 0.87)
- From Sequence to Structure: Uncovering Substructure Reasoning in Transformers (NeurIPS 2025 · similarity 0.87)
- Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable Tasks (ICML 2024 · similarity 0.87)
- Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regression (ICML 2026 · 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.