Research map / Probabilistic and scientific modeling / Molecules and proteins
SPARTAN: A Sparse Transformer World Model Attending to What Matters
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.
Proposes SPARTAN, a transformer world model with sparsity-regularised hard attention (sampled Bernoulli adjacency over object tokens) that learns state-dependent local causal graphs and adapts to unknown interventions via intervention tokens.
Problem
Transformer world models over object-centric tokens lack structured, state-dependent interaction structure, so they learn spurious edges, are sensitive to distractors, and adapt poorly to dynamics changes.
Approach
Apply sparsity regularisation and hard, Gumbel-sampled attention masks to object-factored transformer tokens, with path-counting across layers so the learned adjacency is a local causal graph; add learnable intervention tokens for unknown-target interventions and adapt by optimising a token at test time.
Key results
- On Interventional Pong, SPARTAN reaches SHD 1.51 vs 6.37 for the Transformer and 5.42 for Global Graph, with prediction error 8.60 vs 8.83.
- On CREATE, SPARTAN reaches SHD 1.17 vs 6.29 for the Transformer, with prediction error 246.6 vs 265.5.
- Under removal of non-causal objects, SPARTAN's error change is 24.5% on Interventional Pong versus 1140.2% for the Transformer, and 28.6% on CREATE versus 138.2%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SPARTAN instead of Transformer | World model · Interventional Pong | SHD: 6.37 -> 1.51 | Helped |
| Use SPARTAN instead of Transformer | World model · Interventional Pong | Pred. Err.: 8.83 -> 8.60 | Helped |
| Use SPARTAN instead of Transformer | World model · CREATE | SHD: 6.29 -> 1.17 | Helped |
| Use SPARTAN instead of Transformer | World model · CREATE | Pred. Err.: 265.5 -> 246.6 | Helped |
| Use SPARTAN instead of Transformer | World model · Traffic | SHD: 10.6 -> 6.84 | Helped |
| Use SPARTAN instead of Transformer | World model · Traffic | Pred. Err.: 0.51 -> 0.52 | Hurt |
| Use SPARTAN instead of Transformer | World model · Interventional Pong | Percentage change of error after removing non-causal entities: 1140.2 $\pm$ 15 -> 24.5 $\pm$ 4.4 | Helped |
| Use SPARTAN instead of Transformer | World model · CREATE | Percentage change of error after removing non-causal entities: 138.2 $\pm$ 5.6 -> \textbf{28.6} $\pm$ 1.7 | Helped |
| Use SPARTAN instead of Transformer | World model · Traffic | Percentage change of error after removing non-causal entities: 32.8 $\pm$ 0.4 -> \textbf{13.9} $\pm$ 0.2 | Helped |
| Use local K-nearest-neighbour attention instead of SPARTAN | MTR · Traffic | Percentage change of error after removing non-causal entities: \textbf{13.9} $\pm$ 0.2 -> 25.5 $\pm$ 0.3 | Hurt |
| Use SPARTAN instead of Sparse ACD | Graph learner · Interventional Pong | SHD: 2.67 -> 1.50 | Helped |
| Use SPARTAN instead of Sparse ACD | Graph learner · CREATE | SHD: 5.45 -> 1.17 | Helped |
Surprising findings
- On Traffic, the Transformer baseline has marginally lower prediction error (0.51) than SPARTAN (0.52), so sparsity does not improve accuracy there.
- Soft attention with thresholding (CoDA-style) is insufficient for local causal discovery on image-based observations, according to the authors.
Limitations
- No theoretical identifiability guarantees when scene composition varies across samples.
- Adaptation is limited to the learned intervention space, so entirely new behaviours may not be captured.
- Relies on pre-disentangled object representations; experiments use ground-truth object masks.
Experimental setup
Datasets: Interventional Pong CREATE Waymo Open Dataset
Benchmarks: Interventional Pong CREATE Waymo Open Dataset (Traffic)
Metrics: rollout prediction error structural hamming distance percentage change in error under non-causal removal few-shot adaptation error
Baselines: Transformer (SlotFormer-style) Global Graph (VCD/AdaRL-style) CoDA MTR local attention Amortised Causal Discovery
Builds on: SlotFormer MTR CoDA VCD AdaRL Gumbel softmax
Closest papers by meaning
- Trimming the Long-Tail of Visual World Modeling Evaluation (NeurIPS 2026 · similarity 0.87)
- Simple, Good, Fast: Self-Supervised World Models Free of Baggage (ICLR 2025 · similarity 0.86)
- Learning to Evolve Scenes: Reasoning about Human Activities with Scene Graphs (NeurIPS 2026 · similarity 0.85)
- StateSpaceDiffuser: Bringing Long Context to Diffusion World Models (NeurIPS 2025 · similarity 0.84)
- Sim-to-Real Causal Transfer: A Metric Learning Approach to Causally-Aware Interaction Representations (CVPR 2025 · similarity 0.84)
- Learning Interactive World Model for Object-Centric Reinforcement Learning (NeurIPS 2025 · similarity 0.84)
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.