Learning Local Causal World Models with State Space Models and Attention

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Petri, Francesco, Asprino, Luigi, Gangemi, Aldo
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866913819713863680
author Petri, Francesco
Asprino, Luigi
Gangemi, Aldo
author_facet Petri, Francesco
Asprino, Luigi
Gangemi, Aldo
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_02074
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Learning Local Causal World Models with State Space Models and Attention
Petri, Francesco
Asprino, Luigi
Gangemi, Aldo
Machine Learning
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.
title Learning Local Causal World Models with State Space Models and Attention
topic Machine Learning
url https://arxiv.org/abs/2505.02074