# Learning Local Causal World Models with State Space Models and Attention

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in World models for language agents, one of 11 directions in World models.
- Builds on 16 of its 19 references that are in the map.
- First posted on arXiv 4 May 2025.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)

## Builds on (most cited first, 16 of 16)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [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)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [SlotFormer: Unsupervised Visual Dynamics Simulation with Object-Centric Models](https://mlsanity.com/paper/2210.05861/) (2022)
- [CITRIS: Causal Identifiability from Temporal Intervened Sequences](https://mlsanity.com/paper/2202.03169/) (2022)
- [Counterfactual Data Augmentation using Locally Factored Dynamics](https://mlsanity.com/paper/2007.02863/) (2020)
- [Robust agents learn causal world models](https://mlsanity.com/paper/2402.10877/) (2024)
- [A Tutorial on Learning With Bayesian Networks](https://mlsanity.com/paper/2002.00269/) (2020)
- [Facing Off World Model Backbones: RNNs, Transformers, and S4](https://mlsanity.com/paper/2307.02064/) (2023)
- [Explainable Reinforcement Learning via a Causal World Model](https://mlsanity.com/paper/2305.02749/) (2023)
- [Slot State Space Models](https://mlsanity.com/paper/2406.12272/) (2024)
- [Causal Intervention Improves Implicit Sentiment Analysis](https://mlsanity.com/paper/2208.09329/) (2022)
- [SPARTAN: A Sparse Transformer World Model Attending to What Matters](https://mlsanity.com/paper/2411.06890/) (2024)

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

- [A Comprehensive Survey on World Models for Embodied AI](https://mlsanity.com/paper/2510.16732/) (2025)
- [Why We Need World Models for AGI: Where LLMs Fail and How World Models May Outperform](https://mlsanity.com/paper/2605.23972/) (2026)

## Abstract (arXiv)

World modelling, i.e. building a representation of the rules that govern the world so as to predict its evolution, is an essential ability for any agent interacting with the physical world. Despite their impressive performance, many solutions fail to learn a causal representation of the environment they are trying to model, which would be necessary to gain a deep enough understanding of the world to perform complex tasks. With this work, we aim to broaden the research in the intersection of causality theory and neural world modelling by assessing the potential for causal discovery of the State Space Model (SSM) architecture, which has been shown to have several advantages over the widespread Transformer. We show empirically that, compared to an equivalent Transformer, a SSM can model the dynamics of a simple environment and learn a causal model at the same time with equivalent or better performance, thus paving the way for further experiments that lean into the strength of SSMs and further enhance them with causal awareness.