# SPARTAN: A Sparse Transformer World Model Attending to What Matters

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 45 of its 56 references that are in the map.
- First posted on arXiv 11 Nov 2024.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Scalability in Perception for Autonomous Driving: Waymo Open Dataset](https://mlsanity.com/paper/1912.04838/) (2019)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Object-Centric Learning with Slot Attention](https://mlsanity.com/paper/2006.15055/) (2020)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [Sim-to-Real Transfer of Robotic Control with Dynamics Randomization](https://mlsanity.com/paper/1710.06537/) (2017)
- [GAIA-1: A Generative World Model for Autonomous Driving](https://mlsanity.com/paper/2309.17080/) (2023)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Motion Transformer with Global Intention Localization and Local Movement Refinement](https://mlsanity.com/paper/2209.13508/) (2022)
- [Causality for Machine Learning](https://mlsanity.com/paper/1911.10500/) (2019)
- [Differentiable Causal Discovery from Interventional Data](https://mlsanity.com/paper/2007.01754/) (2020)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [Disentanglement via Mechanism Sparsity Regularization: A New Principle for Nonlinear ICA](https://mlsanity.com/paper/2107.10098/) (2021)
- [A Meta-Transfer Objective for Learning to Disentangle Causal Mechanisms](https://mlsanity.com/paper/1901.10912/) (2019)
- [Adaptively Sparse Transformers](https://mlsanity.com/paper/1909.00015/) (2019)
- [SlotFormer: Unsupervised Visual Dynamics Simulation with Object-Centric Models](https://mlsanity.com/paper/2210.05861/) (2022)
- [Transformer-based World Models Are Happy With 100k Interactions](https://mlsanity.com/paper/2303.07109/) (2023)
- [PHYRE: A New Benchmark for Physical Reasoning](https://mlsanity.com/paper/1908.05656/) (2019)
- [CITRIS: Causal Identifiability from Temporal Intervened Sequences](https://mlsanity.com/paper/2202.03169/) (2022)
- [Entity Abstraction in Visual Model-Based Reinforcement Learning](https://mlsanity.com/paper/1910.12827/) (2019)
- [GENESIS-V2: Inferring Unordered Object Representations without Iterative Refinement](https://mlsanity.com/paper/2104.09958/) (2021)
- [SCALOR: Generative World Models with Scalable Object Representations](https://mlsanity.com/paper/1910.02384/) (2019)
- [The Consciousness Prior](https://mlsanity.com/paper/1709.08568/) (2017)
- [Learning Independent Causal Mechanisms](https://mlsanity.com/paper/1712.00961/) (2017)
- [Relational Neural Expectation Maximization: Unsupervised Discovery of Objects and their Interactions](https://mlsanity.com/paper/1802.10353/) (2018)
- [Amortized Causal Discovery: Learning to Infer Causal Graphs from Time-Series Data](https://mlsanity.com/paper/2006.10833/) (2020)

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

- [Causal-JEPA: Learning World Models through Object-Level Latent Masking](https://mlsanity.com/paper/2602.11389/) (2026)
- [Intrinsically Interpretable Attention via Sparse Post-Training](https://mlsanity.com/paper/2512.05865/) (2025)
- [From Generative Engines to Actionable Simulators: The Imperative of Physical Grounding in World Models](https://mlsanity.com/paper/2601.15533/) (2026)
- [From Observation to Insight: Mechanistic World Models and the Quest for Autonomous Discovery](https://mlsanity.com/paper/2607.12474/) (2026)
- [Better Decisions through the Right Causal World Model](https://mlsanity.com/paper/2504.07257/) (2025)
- [AXIOM: Learning to Play Games in Minutes with Expanding Object-Centric Models](https://mlsanity.com/paper/2505.24784/) (2025)
- [Learning Local Causal World Models with State Space Models and Attention](https://mlsanity.com/paper/2505.02074/) (2025)
- [ECG-WM: A Physiology-Informed ECG World Model for Clinical Intervention Simulation](https://mlsanity.com/paper/2605.17580/) (2026)
- [Dynamic Sparsity: Challenging Common Sparsity Assumptions for Learning World Models in Robotic Reinforcement Learning Benchmarks](https://mlsanity.com/paper/2511.08086/) (2025)
- [World-as-Graph: Relational World Modeling Through Latent Space Graphs](https://mlsanity.com/paper/2609.38927/) (2026)

## Abstract (arXiv)

Capturing the interactions between entities in a structured way plays a central role in world models that flexibly adapt to changes in the environment. Recent works motivate the benefits of models that explicitly represent the structure of interactions and formulate the problem as discovering local causal structures. In this work, we demonstrate that reliably capturing these relationships in complex settings remains challenging. To remedy this shortcoming, we postulate that sparsity is a critical ingredient for the discovery of such local structures. To this end, we present the SPARse TrANsformer World model (SPARTAN), a Transformer-based world model that learns context-dependent interaction structures between entities in a scene. By applying sparsity regularisation on the attention patterns between object-factored tokens, SPARTAN learns sparse, context-dependent interaction graphs that accurately predict future object states. We further extend our model to adapt to sparse interventions with unknown targets in the dynamics of the environment. This results in a highly interpretable world model that can efficiently adapt to changes. Empirically, we evaluate SPARTAN against the current state-of-the-art in object-centric world models in observation-based environments and demonstrate that our model can learn local causal graphs that accurately reflect the underlying interactions between objects, achieving significantly improved few-shot adaptation to dynamics changes, as well as robustness against distractors.