FACTS: A Factored State-Space Framework For World Modelling

TL;DR

The paper introduces FACTS, a novel recurrent framework for spatial-temporal world modelling that efficiently captures complex dependencies, supports parallel computation, and outperforms existing models across various tasks.

Contribution

FACTS is a new graph-structured memory framework with permutation invariance and parallel processing, advancing state-space models for world modelling.

Findings

Outperforms or matches state-of-the-art models in diverse tasks

Supports parallel computation for high-dimensional sequences

Provides permutation-invariant and adaptable memory representations

Abstract

World modelling is essential for understanding and predicting the dynamics of complex systems by learning both spatial and temporal dependencies. However, current frameworks, such as Transformers and selective state-space models like Mambas, exhibit limitations in efficiently encoding spatial and temporal structures, particularly in scenarios requiring long-term high-dimensional sequence modelling. To address these issues, we propose a novel recurrent framework, the \textbf{FACT}ored \textbf{S}tate-space (\textbf{FACTS}) model, for spatial-temporal world modelling. The FACTS framework constructs a graph-structured memory with a routing mechanism that learns permutable memory representations, ensuring invariance to input permutations while adapting through selective state-space propagation. Furthermore, FACTS supports parallel computation of high-dimensional sequences. We empirically evaluate FACTS across diverse tasks, including multivariate time series forecasting, object-centric world modelling, and spatial-temporal graph prediction, demonstrating that it consistently outperforms or matches specialised state-of-the-art models, despite its general-purpose world modelling design.