mlsanityData to 14 Oct 2026

Learning Local Causal World Models with State Space Models and Attention

arXiv 2505.02074 ↗First posted 4 May 2025cs.LG stat.MLcited by 2

On the map: World models → World models for language agents

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.

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.

Builds on

16 of 19 references.

Closest ancestors · 2

Same direction, newest first

  1. 2024Robust agents learn causal world models16 Feb
  2. 2023Explainable Reinforcement Learning via a Causal World Model4 May

Other references · 14

  1. 2017Attention Is All You Needcited by 60,361
  2. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,718
  3. 2020Mastering Atari with Discrete World Modelscited by 1,007
  4. 2020Object-Centric Learning with Slot Attentioncited by 876
  5. 2022DayDreamer: World Models for Physical Robot Learningcited by 400
  6. 2022Transformers are Sample-Efficient World Modelscited by 328
  7. 2022SlotFormer: Unsupervised Visual Dynamics Simulation with Object-Centric Modelscited by 147
  8. 2022CITRIS: Causal Identifiability from Temporal Intervened Sequencescited by 133
  9. 2020Counterfactual Data Augmentation using Locally Factored Dynamicscited by 93
  10. 2020A Tutorial on Learning With Bayesian Networkscited by 57
  11. 2023Facing Off World Model Backbones: RNNs, Transformers, and S4cited by 31
  12. 2024Slot State Space Modelscited by 17
2 more
  1. 2022Causal Intervention Improves Implicit Sentiment Analysiscited by 16
  2. 2024SPARTAN: A Sparse Transformer World Model Attending to What Matterscited by 10

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 1; 2026: 1.

2 papers.

Similar, not cited

Close in wording, no citation either way

  1. 2026A Unifying Perspective on Causal World Models: From Observations to Representations to Structure
  2. 2022Intrinsically Motivated Learning of Causal World Models
  3. 2025Causal Process Models: Reframing Dynamic Causal Graph Discovery as a Reinforcement Learning Problem
  4. 2026HCLSM: Hierarchical Causal Latent State Machines for Object-Centric World Modeling
  5. 2020Causal World Models by Unsupervised Deconfounding of Physical Dynamics
  6. 2026Abductive World Modeling via Causal Representation Learning
  7. 2025SmallWorlds: Assessing Dynamics Understanding of World Models in Isolated Environments
  8. 2024Learning World Models With Hierarchical Temporal Abstractions: A Probabilistic Perspective