Learning Local Causal World Models with State Space Models and Attention

2025-05-04

AI Summary

Existing world models—particularly Transformer-based architectures—exhibit fundamental limitations in learning causal representations, especially in capturing local causal structures inherent in physical dynamics. Method: This paper proposes a local causal world modeling framework grounded in State Space Models (SSMs), the first to integrate SSMs into neural world modeling. Leveraging their implicit temporal inductive bias and low-rank dynamical modeling capacity, the approach explicitly discovers local causal dependencies in physical environments. The architecture synergistically combines SSMs, lightweight attention, and a differentiable causal discovery module to enable interpretable and generalizable dynamics prediction. Results: Experiments on standard physics simulation benchmarks demonstrate that the model matches or exceeds comparably sized Transformers in dynamic prediction accuracy, while substantially improving causal structure identification accuracy. These results validate the intrinsic suitability and promise of SSMs for causal world modeling.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Relational Probabilistic Models
• Knowledge Representation and Reasoning: Action, Change, and Causality

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Learning causal world models for environment prediction
• Comparing State Space Models and Transformers for causality
• Enhancing causal discovery in neural world modelling

Innovation

• Uses State Space Models for causal discovery
• Combines SSM with attention mechanisms
• Learns local causal world models effectively