Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning Local Causal World Models with State Space Models and Attention

Francesco Petri, Luigi Asprino, Aldo Gangemi

TL;DR

This paper addresses the challenge of learning causal world models by integrating causal discovery with a state-space modeling approach. It introduces Sparse Slot State Space Model (S2-SSM), which combines object-centric slot representations with sparsity-regularized attention to infer local causal graphs among objects and environment slots. Empirical results on Interventional Pong show that S2-SSM can predict future frames and recover meaningful causal graphs, matching or surpassing Transformer-based baselines in several settings, with sparsity regularization proving essential for robust causal discovery. The work suggests that state-space models can deliver memory-efficient, causally aware world models suitable for robust planning in changing environments.

Abstract

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.

Learning Local Causal World Models with State Space Models and Attention

TL;DR

This paper addresses the challenge of learning causal world models by integrating causal discovery with a state-space modeling approach. It introduces Sparse Slot State Space Model (S2-SSM), which combines object-centric slot representations with sparsity-regularized attention to infer local causal graphs among objects and environment slots. Empirical results on Interventional Pong show that S2-SSM can predict future frames and recover meaningful causal graphs, matching or surpassing Transformer-based baselines in several settings, with sparsity regularization proving essential for robust causal discovery. The work suggests that state-space models can deliver memory-efficient, causally aware world models suitable for robust planning in changing environments.

Abstract

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.
Paper Structure (21 sections, 6 equations, 3 figures, 2 tables)

This paper contains 21 sections, 6 equations, 3 figures, 2 tables.

Table of Contents

  1. INTRODUCTION
  2. Related work
  3. Causal Discovery
  4. World Modelling
  5. METHOD
  6. Obtaining Object Representations
  7. Environment Representation
  8. World Model
  9. Learning a Causal Graph
  10. Training Objective
  11. Reconstruction loss
  12. Sparsity regularization
  13. EXPERIMENTS
  14. Experimental Setup
  15. Environment description
  16. ...and 6 more sections

Figures (3)

  • Figure 1: Scheme of the S2-SSM architecture.
  • Figure 2: Qualitative examples showing the reconstructed image and the causal graph in two different environments. The nodes of the graph, from the top clockwise, are: environment slot, score, right pong, ball, left pong, border.
  • Figure 3: Error metrics separated by environment. The MEAN column is the average of all the others.