FACTS: A Factored State-Space Framework for World Modelling
Authors: Li Nanbo
Firas Laakom
Yucheng XU
Wenyi Wang
Jürgen Schmidhuber
Venue: ICLR 2025 Poster (2025)
Keywords: sequence modelling
spatial-temporal modelling
world modelling
multivariate time-series forecasting
object-centric representation learning
unsupervised learning
self-supervised learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose the FACTS model, a novel factored state-space framework for spatial-temporal world modelling, which outperforms or matches state-of-the-art models in tasks like time series forecasting and object-centric world modelling.

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 FACTored State-space (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.

Published as a conference paper at ICLR 2025
FACTS: A FACTORED STATE-SPACE FRAMEWORK FOR WORLD MODELLING
Li Nanbo's ; Firas Laakom', Yucheng Xu², Wenyi Wang' .Jrgen Schmidhuber Center of Excellence for Generative AI, KAUST, Saudi Arabia School of Informatics, University of Edinburgh, United Kingdom
Thh Swiss Al Lab, IDSIA, USI & SUPSI, Switzerland
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 strue- tures, particularly in scenarios requiring long-term high-dimensional sequence modelling. To address these issues, we propose a novel recurrent framework, the FACTored State-space (FACTS) model, for The FACTS framework constructs graph-structtred memory with routing mechanism that learns permutable memory representations, ensuring invariance to input permutations while adapting through selective state-space propagation. Fur- thermore, 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-temporaa graph pre- diction, demonstrating that consistently outperforms or matches specialised state-of-the-ard models, despite its general-purpose world modelling design. Code available at: ? https: / /i thub. . com/NanboL.i hwmm/aabbbLi/FACTS
world modelling.
INTRODUCTION
World modelling (Schmidhuber 1990b; 2015 Ha & Schmidhuber 2018) aims to create an internal representation of the environment for an Al system, enabling it represent (Hafner et al. 2019 understand (Schrittwieser et al., 2020 Hafner et al. (222)), and predict (Ha & Schmidhuber
2023
2018 Micheli et al. (022) the dynamics of complex environments. This capability crucial for various domains, including autonomous systems, robotics, and financial forecasting, where accurate predictions depend on effectively capturing both spatial and temporal dependencies (Hafner et al.) 2019|Ha & Schmidhuber] 2018). Consequently, spatial-temporal learning (Liu et al.] [002]] Hochre- iter & Schmidhuber] 1997 Wu et al., 2023a Oreshkin et al. 2019) emerges as a key challenge in world modelling, as approaches must balance the complexities of modelling high-dimensional quential data while maintaining robust long- terg- predictive power.
Despite significant advancements, current ppatial-temporal learning frameworks, used in world mod- elling, based on Transformers (Schmidhuber 1992a Vaswani et al. 20177] Schlag et al., 202I] and
RNNs (Schmidhuber 2015 [Ha & Schmidhuber 2018 Hafner et 2019 (020)) backbones, face apatial-tempon data. Trans- former, though powerful (Chen al. 2022 Robine al. 2023 Micheli et al., (2022), are inefficient for long-term tasks due to their quadratic scaling and limited context windows (Zhang et al. 2022). On the other hand, RNNs provide a more structured approach sequential data. However, their efficacy is hindered by the vanishing gradients (Hochreiter 1991 Pascanu et al. 2013). The pri- mary challenges in necessity to preserve long-term dependencies (Hochreiter et al. 2001 Tallec & Ollivier 2018).
limitations in fully capturing the complexities of high- nigh-dimensional
learning arise from the high dimensionality of the data and the
Recently, there has been a growing interest in Structured State-Space Models (SSMs) for world modelling (Gu & Dao 2023 Hafner et al., 2023 Samsami et al. 2024) using latent state-space
' Coorrespnndence nanbo. @kkkautt... edu.s.

---

Published as a conference paper at ICLR 2025
Radoned M Manny
Inputs
Xt
X
Xit7
Figure 1: Overview of FACTored State-space (FACTS) Architecture. The FACTS framework constructs factored state-space memory, allowing for flexible representations (e.g. graphs and sets). Sequential inputs (e.g. X²) are processed through selective memory-inpun interaction mech- anism (denoted by the circular icon O), which determines how the inputs interact with and update factored memory. The different coloured pathways represent distinct latent factors, whose dynam- ics evolve over time based on these interactions. The design ensures that the memory update is commutation-invarian with respect the input features, enabling FACTS capture and track mean- ingful algorithmic regularities for accurate future predictions.
representations. These representations allow for the modelling of underlying dynamics, where la- tent states evolve over time according to the governing equations (Wang al. 2024b). However, while SSMs have demonstrated improved capacity for capturing temporal dynamics (Gu & Dao, Wang et al. 2024b Baron al. (023). they often lack efficient mechanisms to handle high-
2023
dimensional spatial data. To address this limitation, recent approaches often impose rigid structural constraints their state spaces, such diagonal (Gu & Dao 2023 Gupta et al.) 2022a b)) or block- diagonal structures (Dao & Gu (0044), to capture invariant components throughout the sequence. This assumption, that specific dimensions of the state space correspond consistent patterns over time, can be restrictive in world modelling scenarios where the relationship between state-space dimensions and imput features evolves dynamically.
For example, in dynamic system involving multiple agents (e.g. robots or sensors), where the positions of the agents change over time, to capture this dynamism, current SSMs require learn- ing distinct representations for essentially identical scenarios at each time step the agent loca- tions change. This redundancy leads an inefficient use of model capacity and data, ultimately limiting the model's ability to effectively capture the dynamics of the interactions between inputs and states. Therefore, there a need for a consistent dynamic mapping between inputs and latent states enhance apatial-tempern sion (Schmidhuber 1992b] 2003), which essential for robust long-term prediction power. Another limitation of current SSMs is their inability to capture redundancy in the input space itself. In many cases, each agent's s state (e.g., position, speed, direction) may contributes to the world's overall un- derstanding, but the identity or order of the agents do not matter, i.c., only their interactions are crucial for making accurate predictions. In such instances, swapping agents should not alter the predictions and the world understanding. However, current SSM-based methods, typically based on linear transformations, fail to account for this and can perceive identical scenarios as different based on imput order, hindering its ability to capture regularities and making them unsuitable for sequential modelling in various applications.
modelling capabilities and enable more efficient history compres-
To address these challenges, in this paper, we propose the FACTS model, novel recurrent frame- work for spatial-temporal world modelling. FACTS conceptualises the input a set of nodes and introduces permutable memory that can incorporate complex structures. Through selective memory-input routing, input features are dynamically assigned to distinct state-space factors, i.e., explanatory latent representations, that capture the underlying dynamics of the system. This formu- lation ensures input permutation invariance in the state-space memory, allowing the model to learn consistent factor representations over time, even when the spatial or temporal relationships between input features and factors change illustrated in Figure Additionally, by treating inputs set of
2

---

Published as a conference paper at ICLR 2025
nodes, FACTS: i) can incorporate wider range of input structures, e.g., images graphs or sets. ii) maintains consistent representations of inputs (e.g. agents), regardless of their order. This allows the model to capture regularities in both spatial/local modalities and temporal dependencies, enhancing its memory efficiency and long-term prediction capabilities through more efficient history compres- sion (Schmidhuber 1992b 2003). To validate our proposed world modelling approach, conduct an extensive empirical analysis across multiple tasks, such as multivariate time series forecasting and object-centric world modelling, demonstrating that FACTS consistently matches exceeds the performance of specialised state-of sfate-of-the-ar models. This confirms its robustness and versatility in addressing complex high- -ighddimessional sequential tasks.
In summary, our main contributions are as follows:
We introduce FACTored State-space (FACTS), a novel recurrent framework that incorporates permutable memory structure, enabling flexible and efficient modelling of complex spatial- temporal dependencies.
FACTS dynamically assigns input features to distinct latent state-space factors, ensuring effective history compression and enhancing long-term prediction power.
We formally and empirically show that FACTS achieves consistent factor representations over time, regardless of changes in the spatial order of input features over time, providing robustness in dynamically evolving environments.
We validate the robustness and predictive power of FACTS through extensive world modelling ex- periments, demonstrating its superior competitive performance across multivariate time-series forecasting, object-centric world modelling, and spatial-temporal graph prediction tasks.
PRELIMINARIES
Structured-state space Models (SSMs) have their roots in the classic Kalman filter (Kalman 1960).
where they process mn-dimensional input signal (+)) R into a d-dimensional latent state =(() e
Rd which is then projected onto an output signal y(t) IR.. The general form of an SSM is
expressed as follows:
=()= = A()=((
(1)
where =(t) =t) indicates the time derivative of the state. The matrices A(t) E Radd B(t) Riddo C(t) R2 axd and D(t) Rnxm present the state, input, output, and feed-forward matrices,
respectively. In systems without direct feedthrough, (tt becomes zero matrix. Furthermore.
since the original system operates in a continuous domain, discretisation is often used (Wang et al. 2024b
Smith et al. 2023). resulting in the general discrete time formulation of SSM:
2: A, + B Y: C,
(3)
(4)
with A.. B., and C, govern the dynamics driven by the input sequence Vgt, with different construc- tions (Wang et al. 2024b Gu & Dao 2023 Dao & Gu 2024) influencing the expressiveness and efficiency of the model. If we denote the state vector with h. we can see that equation equationneqquatiin equation
form is equivalent to the RNN dynamics. Hence, similarly to RNNs, the system, in equation is inherently sequential, which limits parallel processing.
Parallelisation and the selective mechanism As shown in Bielloch (1990); Smith et al (022) if B, constructed independently of ø;-1, the linear recurrence in equation can expanded as equa- tion 5
Σ A.y.a.. 3-0
where A A1-1...AA42AA A+t1 I; BDzD 2u with respect to some initialisation. This expansion not only allows for parallel computation of the linear terms, but also reveals the direct connection established between distant inputs/observations along the sequential dimension, e.g. xo
3

---

Published as a conference paper at ICLR 2025
and z, with O, thereby facilitating the capture of long-term dependencies. Furthermore, inte- grating a selective mechanism (Gu & Dao 2023 by constructing A, and B, as functions of each
input leads to the following formulation:
21 ( 2-1 (3
Σ A (zts)) B(x,)
(7)
y1 =(
(8)
This formulation enables content-aware compression of the historical information, addressing the modelling. Such selective mechanism is foundational to
issue of memory decay in the effectiveness of modern state-space models (Gu Dao 2023 Dao & Gu (024). Additionally,
as indicated in Eq. equation 6, SSMs can support parallel computation, since the term (xx))
remains independent of the preceding state 2t-1-
PROPOSED FRAMEWORK: FACTS
In this section, we introduce our proposed framework: FACTS (FACTored State-space) model, a novel class of recurrent neural networks designed with structured state-space memory. FACTS characterised by two key features: permutable state-space memory. which allows for flexible representation of system dynamics with more complex structures and invariant recurrence with respect to permutations of the input features, ensuring consistent modelling of underlying factors in the world across different time steps.
One key intuition behind the "permutable state-space memory" in FACTS is the principle of history compression (Schmidhuber] 1992b 2003), which emphasises the need to eliminate redundant infor- mation in sequence modelling while uncovering algorithmic regularities. This principle is essential for effective long sequence modelling with sigh-dimensional data, it improves generalisation by reducing the accumulation of unnecessary information. Existing SSMs often address this challenge by imposing fixed structural constraints their state spaces, such diagonal
block-diagona
structures, to capture invariant components that persist throughout the sequence (Gu & Dao 2023] Gupta al. 2222aa Dao Gui 2024). However, these fixed structural priors assume that specific dimensions of the state space correspond to consistent and specific factors over time. This as- sumption can be limiting in world-modellin scenarios where the relationship between state-space dimensions and input features evolves dynamically. For instance, in video sequence modelling, fac- tors may correspond moving objects, and the spatial location of these objects (i.e., pixel positions) changes from frame frame. In such cases, the model needs to adapt to these changes, but current SSMs formulations struggle to maintain consistent factor representations due to their rigid structural constraints, i.c., in equation 7 the matrices B(x;) must not only select relevant information for modelling sequence dynamics but also account for the changing relative orders between subspaces of 2t-1 and æt, which evolve time. This introduces additional complexity, leading noise and redundancy that can hinder effective history compression.
3.1 FACTS FORMULATION
The FACTS framework is formalised as a class of structured state-space model, which can capture the dynamic interactions between the latent factors and the input features. To facilitate this dynamic factorisation, i.c., the process of identifying and disentangling meaningful factors from the input data over time, at each time step f, we conceptualise the hidden state Z, as graph and hence, the state-space memory is represented as a set of nodes that correspond to the latent factors. The input features X, are also treated as another set of nodes. Formally, let
Z
{2] 2 2
2
X
m1
where Z. denote the set latent factors at time step t and m is the number of input features. By formulating both sets as nodes, FACTS is inherently invariant to permutations of both input features and factors. Then, to efficiently learn the optimal connections between these two sets, we propose graph-based routing mechanism that effectively match input features with the corresponding factors i.e., learn edges between nodes in Z, and X,, reflecting the strength of each correspondence between each input feature and each latent factor.
4

---

Published as a conference paper at ICLR 2025
Dynamic selective state-space updates In analogy with the standard SSM dynamics equation equation g the evolution of the latent factors of FACTS is governed by the following modified state-space dynamics:
Z: = A. 0 21-1 B, o U,
(10)
y/ = Dec C Zi)
(11)
Here, Ze represents the state-space memory time t, which stores the latent factors. The terms A., B,, and C, are selective state-space model parameters responsible for controlling the information flow between the previous memory Z₁-1 and the input features X,. The symbol denotes element- wise multiplication, while Dec is permutation permutationninvvriiat decoder applied to the latent factors.
Compared to the standard SSM dynamics, i.c., equation eequation 4 we note two key differences: (i) FACTS relies on element-wise multiplication, instead of matrix multiplication, to conserve the invariance properties. (ii) 3y in equation is replaced with U, (Zt-1, which is key element in FACTS that models the interactions between the memory Z and the input features X,,
The attention-based router Before diving into the details of the different parts of equation TO and equation [1 we first introduce the routing mechanism used in this work. To maintain the recurrent permutability of equation [O] the routing mechanism between memory and inputs must dynamically assign input features to consistent factors. This can be done using an attention-basel routing mechanism defined follows:
Ze-1 O Ce X. softmax (0(26-10=(()) (Xc) (x()))
(12)
Vd
where the operator C learns the relationships between the memory Z and input features X dynamically determining which features correspond to which latent factors. The functions ø, th and o represent the query, key, and value mappings, respectively, and are applied row-wise to the memory and input features.
Factorisation process The term U, = U((-,, in equation 10 is crucial for capturing the inter- actions between the memory and the input features. Note that in prior works (Gu & Dao 2023 Dao Gui 2024) U, typically constructed as function of the current input X, only. In this paper, we argue that, similar to the gating in RNNs vs LSTMs (Hochreiter & Schmidhuber] 1997), it is more effective use both X, and Z,-1 to conserve long term dependencies. This interaction plays a key role in factorisation, which refers to the process of binding the input features to specific memory items, effectively uncovering the underlying factors. In the FACTS framework, the memory at the previous time step Z₁-1 serves as the prior over the latent factors, and U, is the factor momentum that guides the evolution of these factors across time. This factor momentum is computed as:
U, = Z O do X
(13)
where du, shy, and you are its corresponding query, key, and value mappings.
Selectivity through memory-input routing The selective state-space model parameters A,, B,, and C2 are constructed through interactions between the memory and the input features, ensuring that both the memory and the inputs jointly decide which information should be retained updated. These parameters are computed as follows:
A Zr-1 C Cos.Vans X
A, = exp(((A.)))
(14)
B, = A, @ (Z--1 Cóg.an X,)
= C OoeNo.ec X
(15)
Here, A step size introduced for discretisation, and the functions Φs. V'a, and م (as well as their counterparts for B, and C₁) are responsible for mapping the memory and inputs to their respective selective parameters. The exponential function exp ensures that the selective parameters non-negative, while is a trainable scalar controlling the influence of A, By employing this selective mechanism, FACTS is capable of compressing long sequences in its state-space memory while maintaining the key properties of latent permutation equivariance and row-wise permutation invariance. Hence, FACTS can efficiently capture meaningful factors, e.g., objects in video frames or independent sources in a signal, even as their relationships with input features change over time.
Linearisation Although the framework presented so far in equation I0 has a permutable state-space memory and equivariant through the memory-input routing, which we formally show in Section 3.2
5

---

Published as a conference paper at ICLR 2025
the routing between Zz-1 and X, introduces dependencies of B X;) and ( on Ze-1. This results in a non-linear recurrence between Z,-1 and Z, in equation I0 which limits parallelisation and hinders training efficiency. To overcome this, substitute Z with Zo. ie. the initial memory or state, within the information routing processes. This leads to the final formulation of FACTS:
4. = (ACTSS(1-1,20)))
(16)
A(Zo, X1) @ Z1-1 o
(17)
A(Zo, @FACTS(2)22 Xt-1) )+ B(bs,
(18)
ΣA Xt:) R BZ X.) 0 (Zo,
(19)
FACTSS Z0; X1a)
(20)
where (Zo. A... o A, o Ar- A.t1 and A is filled with "1"; the initial state
Zo can be given a priori, learnable, or sampled from a prior distribution. This formulation in equa-
tion 17 linearise the recurrence in equation 10 breaking the non-linear dependency between Z and Z, That is, as shown equation 20 the inputs interact only with the initial memory, enabling fast computation of Z, using equation| 19 i.e., without recurrence. This significantly improves com- putational efficiency. Note that while Zo is denoted as the "initial state", it needs not represent the sequence's 's true start; long sequences can segmented (chunked) for parallel computation within segments while maintaining sequential dependencies across the original sequence.
3.2 THEORETICAL ANALYSIS OF FACTS
Here, we formally proof the permutation equivariance and invariance properties of FACTS. We first formally define the two fundamental properties, namely left permutation equivariant (L.P...) and right permutation invariant (R.P..) in Definitions D
and 2
respectively.
Definition 1. Let f: _22xx DXxgxknn > RR xns bivariate function with 71, 722, R3, nay N. f is permutation equivariant (L.P.E.) if for all S M1 E Rm1x=,, and M2 Igxxnaxna,
(i M2) = o [M,
where S% denotes the set of permutation matrices of size kook
Definition 2. Let R2ixx Ixtxagxma Iaaxsss be bivariate function with n1, n22 n2, n4, nz, N. f is right permutation invariant (R.P.I.) if for all σ1, 01,02,...,,0 Sng: Mj Ruixnng and M2, M2. M2,..... ,., M2 mRmss
[a (M,M, ) S(M,, [M2 M2 MG)).
These L.P.E. and R.P.I. properties, which formally describe the two fundamental aspects of FACTS: permutable memory and permutation invariant recurrence (w.rt. the features) - with memory Z and input X, serving as the left and right arguments of FACTS. They are thus essential not only for constructing the routing mechanism but also for the overall design of FACTS.
Using Definitions and 2, that by taking memory Z₁ and features X, as the left and right arguments in FACTS (equation 20 we can show the following result:
Theorem 1. FACTS as defined equation 20 is L.P.E. and R.P.I.
The proof of Theorem i is available in Appendix B Theorem I proves our main claim that FACTS: i) is invariant to input features permutation. ii) learns permutable state-space memory, Furthermore, is possible extend our results in Theorem to the more general case, where A, B, U are L.P.E. and R.P.I. functions of Z and X,- The main result is presented in Theorem
Theorem 2. if A. B,U are L.P.E. and R.P.I. functions of Z and X, any dynamics governed by equation T0 is L.P.E. and R.P.I.
The proof of Theorem is available in Appendix Theorem highlights the main condition on the variables A, B,B.U. U to ensure that the model is invariant to input features and has equivariant memory. This can spark future research to develop SSM models based on equation [0] that are efficient history compressors and suitable dynamic world modelling scenarios.

---

Published as a conference paper at ICLR 2025
4 EXPERIMENTS
We design experiments to evaluate the effectiveness of FACTS in world modelling. We frame world modelling as a prediction task, where the model must predict future events in complex environments based on observed history, and we assess its performance by prediction accuracy. Experiments are conducted on three environments: the multivariate time series (MTS) benchmark for forecasting, synthetic multi-object videos (Yi et al., 2020; Greff et al., 2022 Lin al.) 2020), and dynamic-
graph node prediction (Li et al. 2018). Additionally, Appendix presents an extensive ablation
study that highlights the robustness of FACTS.
4.1 LONG TERM MULTIVARIATE
FORECASTING
In many real-world applications, such as climate prediction, traffic flow management, or autonomous systems, predicting future states over long horizons is crucial for effective decision-making. World modelling, in these domains, often involves high-dimensional multivariate inputs- -uch as interact- ing agents, variables, or environmental
reiturrirrquirnggg the system to account for their complex dependencies and interactions across time. Long-term forecasting in this context challenging as demands accurate representation of temporal dynamics over extended periods. Additionally, input features may lack a predefined order, or this order could change dynamically. For example, a system may receive data from sensors (e.g.. temperature and pressure) without knowing which is which during testing. The world model must provide reliable long-term predictions even if the input order changes unexpectedly and generalise unseen configurations, without requiring exposure every possible permutation of input features during training.
We use the open-source Time Series Library (TSLib to evaluate long-term multivariate time-series forecasting (MSTF) tasks across 9 real-world datasets spanning multiple domains (c.g.. energy, weather, and finance). FACTS is compared against state-of-the-arr bascline models (Wang al. 2024c| Liu et al. 2024 Wu et al. 2023a Nie et al. 2023 Zeng et al. 2023 Zhang & Yan 2023 Zhou al. 2022 Wu et al., 20211. following TSLib's standardised settings: the input sequence length is fixed at 96, with prediction lengths of (96. 192, 336, 720). Performance is evaluated using
mean-squared error (MSE) and mean-absolute error (MAE). C.f. Appendix C.2 for more details.
4.1.1 FORECASTING WITH PREDEFINED ORDER (SCENARIO I)
We use the exact same setup to Wang et al. 2024c); Liu al. 2024 Wu et al. (023a)), with the exception of the pre- and post-ppocessseee modules (referred to as the "embedders" and "projectors" in TSLib). In our implementation, structure of FACTS (c.f. Appendix [C for more details). Note that in the standard setup of TSL.b, the arrangement of the input features in the test is not changed and is the identical to the arrangement to the one seen during the training. The average results over the different prediction windows of our proposed approach along with all competing methods are presented in Table I and the full results are available in Table Appendix E
replace these with set functions to accommodate the output
Results Table | highlights the strong performance of FACTS, which achieves competitive results in both metrics, compared to the competing state-of-the-an specialised MSTF models. As shown, FACTS achieves the best MAE and MSE on out of datasets and, terms of MAF, is always the top 2 in 8 out of 9 them. Even where it is not the top performer, FACTS remains highly competitive (2nd or 3rd place) as seen in Traffic and Solar-Energy underscoring its robustness and ability to adapt to different scenarios. FACTS' ability to efficiently capture long-term dependencies can be attributed to its structured state-space memory, which promotes the learning of modular patterns whose interactions explain the apatial-temport correlations of the multivariate observations.
Parallel vs Recurrent FACTS The parallelisation design outlined in our equation I7 20 not only en- hances computational efficiency but also provides flexibility for recurrent applications of the FACTS core operations. As mentioned, long sequences can be segmented for parallel computation within segments while preserving sequential dependencies across the entire sequence. To clarify further, we analysed FACTS' long-term forecasting performance by comparing its fully sequential mode to its fully parallel mode, adjusted by segment window size, the MTS Electricity dataset (c.f. Fig- ure Each window size represents a different update frequency for Zy: for example, a window
Time Series Library benchmark: https::/t/// htt///gitttutete wtt..thttb..u.umm.too.e weeemme-rerr---ee w.trriss-libbarr. l.ibbary.gi...
7

---

Published as a conference paper at ICLR 2025
ETTml
ETTm2
ETThi
ETTh2
Electricity
Exchange
Traflic
Weather
Solar Energy
Autoformer 2021
MSE
0.588
0.327
0.496
0.450
0.227
0.613
0.628
0.338
0.885
MAE
0.617
0.371
0.487
0.459
0,338
0.539
0.379
0.382
0.711
FEDformer 2022
MSE
0.440
0.437
0.214
0,519
0.309
0.291
MAE
0.452
0,349
0.460
0.449
0.327
0.429
0.376
0.360
0.381
TimesNet 2023a
MSE
0.400
0.291
0.458
0.414
0.192
0.416
0.620
0.259
0.301
MAE
0.406
0.333
0,450
0.427
0,295
0.443
0.336
0.287
0.319
PatchTST 2023
0.281
0.469
0.387
0.205
0.367
0.481
0.259
0.270
MAE
0,400
0.326
0.454
0.407
0.290
0.404
0.304
0.281
0.307
DLinear 2023
MSE
0,403
0.350
0.456
0.559
0.212
0.354
0.625
0.265
0.330
MAE
0.407
0.401
0.452
0,515
0,300
0.414
0,383
0.317
0.401
Crossformer 2023
MSE
0.513
0.757
0.529
0,942
0,244
0.940
0.550
0.259
0.641
0.522
0.334
0.707
0.304
0.315
0.639
iTransformer
MSE
0.407
0.288
0.454
0.383
0.178
0,360
0.428
0.258
0.233
MAE
0.410
0,332
0.447
0.407
0,270
0,403
0.282
0.278
0.262
S-Mamba 2024g
MSR
0.398
0.288
0.455
0.381
0.170
0.367
0.414
0.251
0.240
MAE
0.405
0.332
0.450
0.405
0.265
0.408
0.276
0.276
0.273
FACTS (Ours)
MSE
0.392
0.281
0.440
0.373
0.166
0.342
0.472
0.251
0.253
MAE
0.397
0,325
0.428
0.399
0.263
0.392
0,303
0.278
0.272
Table I: Average MSE and MAE of the different approaches on the multivariate time series forecast- ing tasks, For each metric and each dataset, the top performance and the second best are highlighted in red and blue, respectively.
EACT (419
S-Manbs
FAC Fours)
TACTS S4m1
S-Wande
Figure 2: Model robustness to input permutations on 4 MTSF datasets (left to right: Electricity, Traffic, ETTm1, SolarEnergy). Magenta bars represent original performance, salmon bars show performance using our/TSLib implementation, and yellow bars represent results under input permu- tation. Results are averaged over five random seeds, with error bars showing +2x standard deviation.
size of corresponds fully recurrent FACTS, while a window size 96 (the sequence length) cor- responds to fully parallel FACTS. As shown in Figure 3 the models consistently maintained strong performance across various segment window sizes.
4.1.2 FORECASTING WITH UNKNOWN ORDER (SCENARIO II)
To evaluate robustness under dynamic scenarios, we randomly permute the input features during the testing to simulate environments where the arrangement of agents or entities (c.g., robots, sensors) changes unpredictably. This mirrors real-world scenarios where input configurations vary, challeng- ing world models to adapt to unseen input orderings. We focus on top pretrained models from the first scenario (Table []. i.e., FACTS, iTransformer, and 5-Mamba. For the datasets, we use the chal- lenging ones from the first scenario and permute the feature embeddings five times during testing, reporting the average performance and two standard deviations.
Results The main results are presented in Figure While other models, iTransformer and S-Mamba, experience significant degradation in performance when the input feature orders are shuffled. FACTS consistently maintains its prediction performance across the different tasks. These findings corrob- orate the theoretical results of Section 3.2 For example, on the Traffic dataset, the MSE error of S-mamba, which is the model with the top performance in the standard setting (Table Dp. increases more than threefold, and the error for iTransformer doubles. In contrast, FACTS, leveraging its selective memory routing which consistently assigns input features the latent factors, preserves low error rates despite the permutation. This highlights FACTS' ability to handle dynamic and un- ordered environments. This adaptability further emphasizes the generalisation strength of FACTS, particularly in world modelling scenarios where input orders may be inconsistent or unknown.
4.2 OBJECT-CENTRIC WORLD MODELLING
Visual object-centric representation learning (OCRL) (Burgess et al. 2019 Greff et al., 2019 Nanbo
al. 2020) tackles the challenge of binding visual information to consistent factors, object
8

---

Published as a conference paper at ICLR 2025
LPIPS
Method
CLEVRER
OBJ3D
FG-ARIf
mloUt
PredRNN (017)
0.17
0.12
SAVi (2023)
0.64
0.43
VQFormer (2019)
0.18
0.11
OC-SIoLSSM (2024)
0.68
0.55
G-SWM (2020)
0.16
0.10
FACTS (Ours)
0.75
0.60
SlotFormer (2023b)
0.11
0.08
SAVi-dyn (2023)
0.19
0.12
Quantitative results for unsuper- vised object discovery: segmentation qual- ity on MOVI-A under the video reconstruc- tion setting, evaluated by FG-ARI and mIoU (higher indicate better).
Table 3:
FACTS (Ours)
0.09
0.07
Table 2: Quantitative results for slot dynamics prediction: visual quality of future frame pre- diction, measured by LPIPS (lower is better).
features dynamically permute with movement across pixels in videos (Kipf et al. 2023 ). This aligns with FACTS' objective of identifying regularities in dynamic environments for history compres- sion and future-event prediction, making OCRL an ideal evaluation benchmark. To evaluate how FACTS 1) leverages object information for future predictions and 2) aligns its discovered factors with objects, we conduct two OCRL experiments, slot dymamics prediction and unsupervised object discovery, set on widely-used OCRL, datasets (Yi et al. 2020 Lin et al., 2020 Greff et al., 2022).
Slot Dynamics Prediction The task involves having a world model capture object-centric dynamics in latent space: given the latent object representations of observed events ("burn-in"") the model predicts the future latent codes of the objects ("roll-out"). We conducted this experiment following the setup of our major baseline, SlotFormer (Wu et al., (023b). We evaluate the performance of the model by assessing 1) the visual quality of the predicted future frames and 2) the precision of future segmentation map rendered from the predicted latents. We quantify visual quality using the LPIPS metric, which provides stronger alignment with human perception than other commonly used metics such as PSNR and SSIM (Wu et al., [2023b] Sara et al.) (019),, and segmentation accuracy using the commonly used Mean Intersection over Union (mloU) and Adjusted Rand Index (ARI), with and/or without the foreground focus.
Results The results presented in Table and Table g highlight the strengths of the FACTS model in terms of both visual quality and segmentation accuracy for object dynamics prediction on the CLEVRER dataset. FACTS achieves the lowest LPIPS score of 0.09, indicating superior visual quality in the predicted frames. Additionally, it achieves competitive segmentation accuracy, achiev- ing leading FG-mloU of 48.1 48%. This highlights its effectiveness predicting object positions and interactions in future frames. We attribute these results to FACTS' selective history compres- contrast to SlotFormer, which predicts the next state by attending to all past
sion mechanism. inputs-resulting in inefficiency and noisy Prddictions-FACSS effectively compresses and retains only the most relevant information in memory, thereby filtering out noise and enabling more accurate dynamics modelling and future predictions.
Unsupervised Object Discovery In contrast to the slot dynamics task, where object slots or factors given as input, this experiment requires FACTS automatically discover relevant factors for future predictions in multi-object videos. This process enables us to understand the regularities that FACTS conceptualises as significant for forecasting future events. We utilised a CNN encoder to convert the input image into feature set, from which FACTS learns the object factors. These factors employed predict future object slots and are subsequently decoded back into video frames using a setup of SAVi and jointly train all the modules (including the CNN encoder and apatial-broadees decoder) end-to-end from scratch by minimising the reconstruction MSE and future prediction MSE. In addition to future prediction, we also conducted object discovery under the video reconstruction setting using the MOVi-A dataset to enable direct comparison with the state-of-the-an the- video-based OCRL models for videos, i.c., SAVi (Kipf et al. 2023) and OC-SlotSSM (Jiang et al., (024)).
decoder. Note that, in this experiment, we adopted the uullyeunseer
Results We visualise the discovered factors by independently rendering each factor's 's dynamics back into videos. In the "discovery for prediction" task, FACTS primarily identifies moving objects considered "usefull for future predictions- while treating static objects background. In contrast, for the "discovery for reconstruction" " task, FACTS identifies also static objects as explanatory fac- tors (see Figure like SAVi and OC-SlotSSM. We attribute this behaviour to the residual design

---

Published as a conference paper at ICLR 2025
Method
RMSEL
MAET
MAPET
HA (2018)
12.16
5,92
15.17%
0.30
LSTNet 2018
9.22
5.11
12.56%
STGCN 2018
7.92
3.87
10.05%
0.28
DCRNN 2018
7.87
3,85
10.01%
GWN (2019)
7.66
3.54
9.98%
0.26
MSE
ASTGCN (2019)
7.99
3.94
10.12%
Error
MAE
GMA (2020)
8.32
4.06
10.91%
0.24
MTGNN 2020b)
8.16
3.99
10.28%
0.22
AGCRN 20200
8.22
4.02
10.53%
DGCRN 20233
7.19
3.44
9.73%
STGM (2023)
7.10
3.23
9.39%
MegaCRN (2023)
7.23
3,38
9.72%
1
16 Window size
96
TESTAM (2024)
7.09
3.36
9.67%
FACTS (Ours)
6.97
3.11
9.08% %
Figure 3: Parallel vs Recurrent FACTS on long-term MTS (Electricity): MSE and MAE for different window/chunk sizes. In- put (observable) sequence length set to 96.
Table 4: RMSE, MAE, and MAPE of various ap- proaches for long-term graph node forecasting (1- hour, 12-step ahead) on the METR-LA dataset.
of the FACTS predictor, which is muted during reconstruction. Our quantitative results Table demonstrate that FACTS outperforms both OC-SlotSSM and SAVi in unsupervised object discovery (for reconstruction), highlighting its superior performance despite being more general framework. Additional visual results of our object-centric world modelling are available in the Appendix
4.3 LONG TERM PREDICTION WITH GRAPH DATA
To demonstrate FACTS' flexibility powerful general framework able handle different in- put types/modalities and efficiently solve diverse forecasting- based tasks, we also apply FACTS to dynamic-graph input data evaluated long-term node prediction task (12-step MAE) with the commonly used METR-LA dataset (Li al. 2018) and we compare against existing state-of-the-art approaches on this task. We refer Appendix for more experimental details.
Results As can be seen in T FACTS, leveraging its graph-structured memory, also outperforms all existing methods on this task, even those specialised for this task. For instance, TESTAM (Lee & Ko 2024) yields an MAPE of 9.67% whereas FACTS yields 9.08%. This further corroborates our main claim that FACTS is indeed versatile world model framework with consistently strong performance in several diverse forecasting tasks.
5 DISCUSSIONS & CONCLUSION
In this work, we introduced FACTS, novel recurrent framework designed for spatial-temporal world modelling. FACTS constructed permutable state-space memory, which offers the flex- ibility needed to capture complex dependencies across time and space. By employing selective memory-input routing, FACTS is able to dynamically assign input features to distinct latent fac- tors, enabling we formally showed that FACTS: i) is invariant to input feature permutation. ii) learns permutable state-space memory, maintaining consistent factor representations regardless of changes in the input order. Furthermore, through comprehensive empirical evaluations, FACTS demonstrated superior performance on variety of real-world datasets, consistently matching or outperforming specialised state-of-the-ar models in diverse tasks. Notably, FACTS maintained its predictive powers even in challenging settings where the order of input features was shuffled, highlighting its robustness and adaptability. These results underscore the model's 's potential for a wide range of applications, partic- ularly in world modelling scenarios where input configurations are variable uncertain. For future work, we plan to extend FACTS to larger-scale experiments, exploring its scalability and potential even more complex world modelling tasks.
efficient history compression and long-term prediction accuracy. Furthermore,
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
The research reported in this publication was supported by funding from King Abdullah University of Science and Technology (KAUST)- (AUUS) 5940.
Center of Excellence for Generative AI, under award number
REFERENCES
Lei Bai, Lina Yao, Can Li, Xianzhi Wang, and Can Wang. Adaptive graph convolutional recurrent network for traffic forecasting. Advances in neural information processing systems, 33:17804- 17815, 2020.
Ethan Baron, Itamar Zimerman, and Lior Wolf. 2-dimensional state space layer for spatial induc- tive bias. In The Twelfth International Conference Learning Representations, 2023.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xLSTM: Ex- tended long short-term memory. arXiv preprint arXiv:4405005517 2024.
Guy E Blelloch. Prefix sums and their applications. 1990.
Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geomet- ric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine, 34(4):18-44 2017.
Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. MONet: Unsupervised scene decomposition and representa- tion. arXiv preprint / 1390, 2019.
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.094: 094481, 2022.
Tianxiang Chen, Zhentao Tan, Tao Gong, Qi Chu, Yue Wu, Bin Liu, Jieping Ye, and Nenghai Yu. MiM-ISTD: Coamba-n--Mambee for efficient infrared small target detection. arXiv preprint arXiv:2403.02148, 2024.
Silvia Chiappa, Sébastien Racaniere, Daan Wierstra, and Shakir Mohamed, Recurrent environment simulators. arXiv preprint a7Xiv:1704.02254 2017.
Kyunghyun Cho. Learning phrase representations using RNN encoder-decoder for statistical ma- chine translation. arXiv preprint arXiv:/406. 1078, 2014.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060, 2024.
Mare Deisenroth and Carl E Rasmussen. PILCO: A model-based and data-efficient approach to policy search. In International Conference on machine learning, 2011.
Emily L Denton et al. Unsupervised learning of disentangled representations from video. Advances in neural information processing systems, 30, 2017.
Stefan Depeweg, José Miguel Hernández-Lobato, Finale Doshi- Velez, and Steffen Udluft. Learning and policy search in stochastic dynamical systems with Bayesian neural networks. arXiv preprint arXiv: (atf::605.01110 070500777777 2016.
Chengbin Du, Yanxi Li, and Chang Xu. Understanding robustness of visual state space models for image classification. arXiv preprint arXiv:2403. 10935, 2024a.
Yu Du, Xu Liu, and Yansong Chua. Spiking structured state space model for monaural speech enhancement. In ICASSP 2024-2024 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 766-770. IEEE, 2024b.
Zhengcong Fei, Mingyuan Fan, Changqian Yu, and Junshi Huang. Scalable diffusion models with state space backbone. arXiv preprint wrXiv:2402.06608 2024.

---

Published as a conference paper at ICLR 2025
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA). pp. 512-519. IEEE, 2016.
Yarin Gal, Rowan McAllister, and Carl Edward Rasmussen. Improving PILCO with Bayesian neural network dynamics models. In Data-efficient machine learning workshop, ICML, volume 4, pp- 25, 2016.
Zoubin Ghahramani. Factorial learning and the EM algorithm. Advances in neural information processing systems, 7, 1994.
Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Sergey Levine, Charles Blun- dell, Yoshua Bengio, and Michael Mozer. Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems. arXiv preprint arXiv:2006.162 2020.
Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Heess, Michael Mozer, and Yoshua Bengio. Neural production systems: Learning rule-govemed visual dynamics. CoRR, abs/2103.00 10.00.90...7... 2022.
Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. Is Mamba capable of in-context learning? arXiv preprint arXiv:2402.03170, 2024.
Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. Neural expectation maximization. In Advances in Neural Information Processing Systems, pp. 6691-6701, 2017.
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. In Proceedings of the 36th International Conference on Machine Learning. pp. 2223,2433, 2019.
Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arXiv preprint arXiv:2012.05200. 2020.
Klaus Greff, Francois Belletti, Lueas Beyer, Carl Doersch, Yilun Du, Daniel Duckworth, David Fleet, Dan Gnanapragasam, Florian Golemo, Charles Herrmann, Thomas Kipf, Abhijit Kundu, Dmitry Lagun, Issam Laradji, Hsueh-Ti (Derek) Liu, Henning Meyer, Yishu Miao, Derek Nowrouzezahrai, Cengiz Oztireli, Etienne Pot, Noha Radwan, Daniel Rebain, Sara Sabour, Mehdi S. M. Sajjddi, Matan Sela, Vincent Sitzmann, Austin Stone, Deqing Sun, Suhani Vora, Ziyu Wang. Tianhao Wu, Kwang Moo Yi, Fangcheng Zhong, and Andrea Tagliasacchi. Kubric: scalable dataset generator. 2022.
Albert Gu and Tri Dao. Mamba: Linear-tine sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 2021.
Shengnan Guo, Youfang Lin, Ning Feng, Chao Song, and Huaiyu Wan. Attention based spatial- temporal graph convolutional networks for traffic flow forecasting. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 922-929, 2019.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. Advances in Neural Information Processing Systems, 35:22982-22994, 2022a.
Ankit Gupta, Harsh Mehta, and Jonathan Berant. Simplifying and understanding state space models with diagonal linear RNNs. arXiv preprint arXiv22212.00768, 2022b.
David Ha and Jürgen Schmidhuber. World models. Preprint arXiv:1803.1011 (variant at NeurIPS 2018), 2018
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. International conference on machine learning, pp. 2555-2565. PMLR, 2019.
12

---

Published as a conference paper at ICLR 2025
Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world models. arXiv preprint arXiv:2010.02193, 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 0400044444 2023.
Wei He, Kai Han, Yehui Tang, Chengcheng Wang. Yujie Yang, Tianyu Guo, and Yunhe Wang. DenseMamba: State space models with dense hidden connection for efficient large language mod- els. arXiv preprint arXiv:2403.008// 2024.
Daniel Hein, Stefan Depeweg, Michel Tokic, Steffen Udluft, Alexander Hentschel, Thomas A Run- kler, and Volkmar Sterzing. A benchmark environment motivated by industrial control problems. In 2017 IEEE Symposium Series on Computational Intelligence (SSCI), pp. 1-8. IEEE, 2017.
Sepp Hochreiter. Untersuchungen dynamischen neuronalen Netzen. Diploma, Technische Uni- versität München, 91(1):31, 1991.
Sepp Hochreiter and Jürgen Schmidhuber. Bridging long time lags by weight guessing and "long short-term memory". Spatiotemporal models in biological and artificial systems, 37(65-72):11, 1996.
Sepp Hochreiter and Jürgen Schmidhuber. Long Short- Memory. Neural Computation, 9(8): 1735-1780, 1997.
Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, and Jürgen Schmidhuber. Gradient flow in recur- rent nets: the difficulty of learning long- term dependencies, 2001.
Harold Hotelling. Analysis of a complex statistical variables into principal components. Journal of educational psychology, 24(6))417, 1933.
Tao Huang, Xiaohuan Pei, Shan You, Fei Wang, Chen Qian, and Chang Xu. LocalMamba: Visual state space model with windowed selective scan. arXiv preprint XXiv:2403.00338 2024.
Md Mohaiminul Islam and Gedas Bertasius. Long movie clip classification with state-space video models. In European Conference on Computer Vision, pp. 87-104. Springer, 2022.
Jindong Jiang, Fei Deng, Gautam Singh, Minseung Lee, and Sungjin Ahn. Slot state space models. In Neural Information Processing Systems, 2024.
Renhe Jiang, Zhaonan Wang, Jiawei Yong, Puneet Jeph, Quanjun Chen, Yasumasa Kobayashi, Xuan Song, Shintaro Fukushima, and Toyotaro Suzumura. Spatio-temporal meta-graph learning for traffic forecasting. In Proceedings of the AAAI conference on artificial intelligence, volume 37, pp. 8078-8086, 2023.
Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. Journal of Basic Engineering, 82(1):35-45, 1960.
Diederik P Kingma. Auto-encoding variational bayes. arXiv preprini arXiv: 1rXiv:1312.61144 2013.
Thomas Kipf, Elise Van der Pol, and Max Welling. Contrastive learning of structured world models. In International Conference on Learning Representations, 2019.
Thomas Kipf, Gamaleldin Fathy Elsayed, Aravindh Mahendran, Austin Stone, Sara Sabour, Georg Heigold, Rico Jonschkowski, Alexey Dosovitskiy, and Klaus Greff. Conditional object-centric learning from video. In International Conference Learning Representations, 2023.
Igor Kononenko. Bayesian neural networks. Biological Cybernetics, 61(5)36113370, 1989.
Mourad Lablack and Yanming Shen. Spatio-temporal graph mixformer for traffic forecasting. Ex- pert systems with applications, 228:120288 20281, 2023.
Guokun Lai, Wei-Cheng Chang, Yiming Yang, and Hanxiao Liu. Modeling long-and short-term temporal patterns with deep neural networks. The 41st international ACM SIGIR conference research & development in information retrieval, pp. 104, 2018.
13

---

Published as a conference paper at ICLR 2025
Hyunwook Lee and Sungahn Ko. TESTAM: A time-enhanced apatio-temporrs attention model with mixture of experts. In The Tweifth International Conference on Learning Representations, 2024.
URL httpttts//peeeeevieieeeee net.
Fuxian Li, Jie Feng, Huan Yan, Guangyin Jin, Fan Yang. Funing Sun, Depeng Jin, and Yong Li. Dynamic graph convolutional recurrent network for traffic prediction: Benchmark and solution.
ACM Transactions Knowledge Discovery from Data, 17(1):1-21 2023.
Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. Diffusion convolutional recurrent neural net- work: Data-driven traffic forecasting. International Conference on Learning Representations, 2018.
Zhixuan Lin, Yi-Fu Wu, Skand Peri, Bofeng Fu, Jindong Jiang, and Sungjin Ahn. Improving gen- erative imagination object-centric world models. In International Conference on Machine Learning, pp. 6140-6149. PMLR, 2020.
Yong Liu, Tengge Hu, Haoran Zhang, Haixu Wu, Shiyu Wang, Lintao Ma, and Mingsheng Long. iTransformer: Inverted transformers are effective for time series forecasting. In International Conference on Learning Representations, 2024.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. In Advances in Neural Information Processing Systems, 2020.
Xianping Ma, Xiaokang Zhang, and Man-On Pun. RS 3 Mamba: Visual state space model for remote sensing image semantic segmentation. IEEE Geoscience and Remote Sensing Letters, 2024.
David JC MacKay et al. Introduction to Gaussian processes. NATO ASI series computer and systems sciences, 168: 133-166,, 1998.
Rowan McAllister and Carl Edward Rasmussen. Data-efficient reinforcement learning in continuous-state POMDPs. arXiv preprini arXiv:1602.02522 2016.
Harsh Mchta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language model- ing via gated state spaces. arXiv preprint arXiv:2206. /3947, 2022.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world mod- cls. arXiv preprint arXiv:2209.00588 2022.
Paul Munro. A dual back proppapppaation scheme for scalar reward learning. In Ninth Annual Con- ference of the Cognitive Science Society, pp. 165-176. Hillsdale, NJ. Cognitive Science Society Lawrence Erlbaum, 1987.
Li Nanbo, Cian Eastwood, and Robert Fisher. Learning object-centric representations of multi- object scenes from multiple views. In Advances in Neural Information Processing Systems, 2020.
Demrick Nguyen and Bernard Widrow. The truck backer-upper: An example of self-learning in neural networks. Advanced neural computers, pp- 11-19, Elsevier, 1990.
Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4ND: Modeling images and videos as multidimensional signals with state spaces. Advances neural information processing systems, 35:2846 2861, 2022.
Yuqi Nie, Nam H. Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam. A time series is worth 64 words: Long-term forecasting with transformers. In International Conference on Learning Representations, 2023.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard I. Lewis, and Satinder Singh. Action-conditional video prediction using deep networks in Atari games. Advances in neural information proceesssiioooooooniosooo processing systems, 28, 2015.
Boris N Oreshkin, Dmitri Carpov, Nicolas Chapados, and Yoshua Bengio. N-BEATS: Neural basis expansion analysis for interpretable time series forecasting. arXiv preprint 1905v:I905.10 10437, 2019.
14

---

Published as a conference paper at ICLR 2025
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. pp. III-1310, 2013.
Thomas Porter and Tom Duff. Compositing digital images. In Proceedings of the 11th anmual conference on Computer graphics and interactive techniques, pp. 253-259, 1984.
Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with VQ-VAB-2. In Advances in neural information processing systems, 2019.
Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-bass based world mod- els are happy with 100k interactions. arXiv preprint arXiv:2303.07 109, 2023.
Tony Robinson and Frank Fallside. Dynamic reinforcement driven error propagation networks with application game playing. In Proceedings of the Annual Meeting of the Cognitive Science Society, volume 11, 1989.
David Salinas, Valentin Flunkert, Jan Gasthaus, and Tim Januschowski. DeepAR: Probabilistic forecasting with autoregressive recurrent networks. International journal of forecasting, 36(3): 1181-1191, 2020.
Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar. Mastering memory tasks with world models. arXiv preprint 2xXi::2433004255: 2024.
Umme Sara, Morium Akter, and Mohammad Shorif Uddin. Image quality assessment through FSIM, SSIM, MSE and PSNR comparative study. Journal of Computer and Communications, 7(3):8-8,, 2019.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355 9366. PMLR, 2021.
Karl Schmeckpeper, Annie Xie, Oleh Rybkin, Stephen Tian, Kostas Danilidiis, Sergey Levine, and Chelsea Finn. Learning predictive models from observation and interaction. In European Con- ference on Computer Vision, pp. 708-725. Springer, 2020.
Jürgen Schmidhuber.
An on-line algorithm for dynamic reinforcement learning and planning in reactive environments. 1990 IJCNN international joint conference on neural networks, pp-
253-258 IEBE, 1990a.
Jürgen Schmidhuber. Making the world differentiable: On using fully recurrent self-supervised neural networks for dynamic reinforcement learning and planning in non-stationary environments. Institut für Informatik, Technische Universität München. Technical Report FKI-126, 90, 1990b.
Jürgen Schmidhuber. Reinforcement learning in Markovian and Markovvo environments. In D. S. Lippman, J. E. Moody, and D. S. Touretzky (eds.), Advances in Neural Information Pro- cessing Systems (NIPS 3), pp. 500-506. Morgan Kaufmann, 1991a.
Jürgen Schmidhuber. A possibility for implementing curiosity and boredom in model-building neu- ral controllers. pp. 222-227, 1991b.
Jürgen Schmidhuber. Learning to control fast weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-134) 1992a.
Jürgen Schmidhuber. Learning complex, extended sequences using the principle of history com- pression. Neural computation, 4(2):234 1992b.
Jürgen Schmidhuber. Learning factorial codes by predictability minimization. Neural computation, 4(6):863-879, 1992c.
Jürgen Schmidhuber. Exploring the predictable. In Advances evolutionary computing: theory and applications, pp. 579-612. Springer, 2003.
Jürgen Schmidhuber. On learning think: Algorithmic information theory for novel combina- tions reinforcement learning controllers and recurrent neural world models. arXiv preprint 10Xi::55/1002200 2015.
15

---

Published as a conference paper at ICLR 2025
Jürgen Schmidhuber, Daan Wierstra, Matteo Gagliolo, and Faustino Gomez. Training recurrent networks by EVOLINO. Neural computation, 1933:757-779, 2007.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering Atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-609) 2020.
Yuan Shi, Bin Xia, Xiaoyu Jin, Xing Wang, Tianyu Zhao, Xin Xia, Xuefeng Xiao, and Wen- ming Yang. VmambalR: Visual state space model for image restoration. arXiv preprint arXiv:2403.11422 2024.
David Silver, Hado Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac- Arnold, David Reichert, Neil Rabinowitz, Andre Barreto, et al. The Predictron: End-to-end learning and planning. In International Conference on Machine Learning, pp. 3191-3199. PMLR, 2017.
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. International Conference on Learning Representations, 2023.
Rupesh K Srivastava, Klaus Greff, and Jürgen Schmidhuber. Training very deep networks. Advances in neural information processing systems, 28, 2015.
Aleksandar Stanić, Yujin Tang, David Ha, and Jürgen Schmidhuher. Learning to generalize with object-centric agents in the open world survival game crafter. IEEE Transactions on Games, 2023.
Corentin Tallec and Yann Ollivier. Can recurrent neural networks warp time? arXiv preprint arXiv:1804. 11 188, 2018.
Michael E Tipping and Christopher M Bishop. Probabilistic principal component analysis. Journal of the Royal Statistical Society Series B: Statistical Methodology. 61(3):611-63 1-622, 1999.
Naftali Tishby, Fernando Percira, and William Bialck. The information bottleneck method. arXiv preprint physics/0004055, 2000.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances Neural Informa- fion Processing Systems, 2017.
Jue Wang, Wentao Zhu, Pichao Wang, Xiang Yu, Linda Liu, Mohamed Omar, and Raffay Hamid. Selective structured state-spaces for long-form video understanding. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 6387-6397, 2023.
Shiyu Wang, Haixu Wu, Xiaoming Shi, Tengge Hu, Huakun Luo, Lintao Ma, James Y Zhang, and JUN ZHOU. TimeMixer: Decomposable multiscale mixing for time series forecasting. In International Conference Learning Representations (ICLR), 2024a.
Xiao Wang, Shiao Wang. Yuhe Ding, Yuehang Li, Wentao Wu. Yao Rong, Weizhe Kong, Ju Huang, Shihao Li, Haoxiang Yang, et al. State space model for new-generation network alternative to 16, 2024b.
transformers: A survey. arXiv preprint
Yunbo Wang, Mingsheng Long, Jianmin Wang, Zhifeng Gao, and Philip S Yu. PredRNN: Recurrent neural networks for predictive learning using spatiotemporal LSTMs. In Advances neural information processing systems, 2017.
Zihan Wang, Fanheng Kong, Shi Feng, Ming Wang, Han Zhao, Daling Wang, and Yifei Zhang. Is Mamba effective for time series forecasting? arXiv preprint arXiv:2403.11144 2024c.
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual Interaction Networks: Learning a physics simulator from video. Advances in neural information processing systems, 30, 2017.
Nicholas Watters, Loic Matthey, Christopher P Burgess, and Alexander Lerehner. Spatial Broadcast Decoder: A simple architecture for learning disentangled representations in vaes. arXiv preprint arXiv: 1rXivv1901.07011 2019.
16

---

Published as a conference paper at ICLR 2025
Paul Werbos. Learning how the world works: Specifications for predictive networks in robots and brains. In Proceedings of IEEE International Conference on Systems, Man and Cybernetics, NY, 1987.
Paul J Werbos. Neural networks for control and system identification. In Proceedings of the 28th IEEE Conference on Decision and Control,, pp. 360-265, IEEE, 1989.
Christopher K I Williams and Michalis K Titsias. Greedy learning of multiple objects in images using robust statistics and factorial learning. Neural Computation, ((5):1039-1062. 2004.
Haixu Wu, Jiehui Xu, Jianmin Wang, and Mingsheng Long. Autoformer: Decomposition trans-
formers with auto-correlation for long-term series forecasting, In Advances Neural Information Processing Systems, 2021.
Haixu Wu, Tengge Hu, Yong Liu, Hang Zhou, Jianmin Wang, and Mingsheng Long. TimesNet: Temporal 2d-variation modeling for general time series analysis. In International Conference on Learning Representations, 2023a.
Ziyi Wu, Nikita Dvornik, Klaus Greff, Thomas Kipf. and Animesh Garg. SlotFormer: Unsupervised visual dynamics simulation with object-centric models. In International Conference on Learning Representations, 2023b.
Zonghan Wu, Shirui Pan, Guodong Long. Jing Jiang, and Chengqi Zhang. Graph wavenet for deep spatial-temporal graph modeling. In Proceedings of the Twenty-Fighth International Joint Conference on Artificial Intelligence, IJCAI-19, pp. 1907-1913. International Joint Conferences
on Artificial Intelligence Organization, 2019. doi: 1.249663/jcaa.2211/264. URL https:
//doi.org//i 20t.2224433a...213333333 2019//660./2/6
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning systems, 32(1):4-24, 2020a.
Zonghan Wu, Shirui Pan, Guodong Long, Jing Jiang, Xiaojun Chang, and Chengqi Zhang. Con- necting the dots: Multivariate time series forecasting with graph neural networks. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 753-763. 2020b.
Jing Nathan Yan, Jiatao Gu, and Alexander M Rush. Diffusion models without attention. In Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8239-8249, 2024.
Yuhuan Yang, Chaofan Ma, Jiangchao Yao, Zhun Zhong, Ya Zhang, and Yanfeng Wang. Remamber: Referring image segmentation with mamba twister. arXiv preprint arXiv:2403.178.. 2024a.
Zhichao Yang, Avijit Mitra, Sunjae Kwon, and Hong Yu. Clinicalamma:: A generative clinical language model on longitudinal clinical notes. arXiv preprint arXiv:2403.05795, 2024b.
Kexin Yi, Chuang Gan, Yunzhu Li, Pushmeet Kohli, Jiajun Wu, Antonio Torralba, and Joshua B Tenenbaum. CLEVRER: Collision events for video representation and reasoning. In International Conference on Learning Representations, 2020.
Bing Yu, Haoteng Yin, and Zhanxing Zhu. Spatio-temporal graph convolutional networks: A deep learning framework for traffic forecasting. In Proceedings of the Twenty-Seventh International Joint Conference Artificial Intelligence. IJCAI-18, pp. 3634-3640. International Joint Con- ferences on Artificial Intelligence Organization, 7 2018. doi: 10.24963/ijcai.21188 https://doi. wwt.or///.........../ 2492342964310 2a1....1.....55555
URL
Alan Yuille and Daniel Kersten. Vision as Bayesian inference: analysis by synthesis? Trends in cognitive sciences, 2006.
Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. Are transformers effective for time series forecasting? In Proceedings of the AAAI conference on artificial intelligence, volume 37, pp. 11121-11128, 2023.
17

---

Published as a conference paper at ICLR 2025
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. OPT: Open pre-trained transformer language models. arXiv preprint 2rX5v:200..01068, 2022.
Yunhao Zhang and Junchi Yan. Crossformer: Transformer utilizing cross-dimension dependency for multivariate time series forecasting. In International Conference on Learning Representations, 2023.
Chuampan Zheng, Xiaoliang Fan, Cheng Wang, and Jianzhong Qi. Gman: A graph multi-attention network for traffic prediction. In Proceedings of the AAAI conference artificial intelligence, volume 34, pp. 1234-1241, 2020.
Tian Zhou, Ziqing Ma, Qingsong Wen, Xue Wang. Liang Sun, and Rong Jin. FEDformer: Frequency enhanced decomposed transformer for long-tern series forecasting. In International Conference on Machine Learning, 2022.
18

---

Published as a conference paper at ICLR 2025
A RELATED WORKS
From a historical perspective, "world models" using models learn environmental dynamics and
leveraging them in policy training has an extensive literature, with early foundations laid in the 1980s
using feed-forward neural networks (Werbos 1987 1989 Munro 1987 Robinson & Fallside 1989
Nguyen& & Widrow| (190) and in the 1990s with RNNs (Schmidhuber 1990b a 1991b a). Notably,
PILCO (Deisenroth & Rasmussen 2011; McAllister & Rasmussen 2016) has emerged as a key
probabilistic model- based method, using Gaussian processes (GPs) MacKay et al. I 1998) to learn system dynamics from limited data and train controllers for tasks like pendulum su and uni- cycle balancing. While GPs perform well with small, low-dimensional datasets, their computational complexity limits scalability in high-dimensional scenarios. To address this, later works (Gal et al, 2016 Depeweg et al. 2016) have adopted Bayesian neural networks (Kononenko 1989). which have demonstrated success in control tasks with well-defined states (Hein et al. (017)). However, these methods remain limited when modelling digh-dimensional environments, such as sequences of raw pixel frames. In the context of reinforcement learning, using recurrent models to learn system dynamics from compressed latent spaces has significantly improved data efficiency (Schmeckpeperi jet al. 2020 Finn et al., 2016). While the development of internal models for reasoning about future states using RNNs dates back to the early 1990s, subsequent works, such as "Leerning to Think" (Schmidhuber 2015) and "World Models" (Ha & Schmidhuber] (0188), have extended this by introducing RNN-based frameworks that model environments and reason about future outcomes. These RNN-based models have been applied to future frame generation (Chiappa et al. 2017;; Oh
ct al. 2015 Denton et al., (017)) and reasoning about future outcomes (Silver al. [ | Watters et al.
2017). However, as RNNs suffer from the vanishing gradients problem (Hochreiter 1991
Pascanu et al. 2013), recently there has been a growing interest in using Transformers (Chen al.
2022
Robine et al.] 2023]] Micheli et al.) (022) and SSM-based appraoches (Gu & Dao
2023
Hafner
ctt al.
2023 Samsami et al. al.] (224)) for world modelling.
As world modelling is fundamentally intertwined with sequence modelling (Schmidhuber 1990b) often carries temporal implications that align with the principle of history compression (Schmid-
huber 1992b 2003). Temporal selectivity is essential in these models, with Recurrent Neural Net- works (RNNs), particulary those with gating mechanisms like LSTMs (Hochreiter & Schmidhuber 1997]. , GRUs (Chol 2014), and xLSTMs (Beck al. 2024), being well-suited for this task. However, learning from high-dimensional sequential data complicates the problem, posing a core challenge in spatial-temporal learning. This challenge is exacerbated by the quadratic computation scaling in transformers, despite their success. Approaches like dimensionality reduction (Hotelling, 1933
Tipping & Bishop 1999, Kingma, 2013) and predictability minimisation (Schmidhuber) 1992c Ghahramani] 1994) must adhere to the principle of history compression along the temporal axis, rather than compressing spatial information at each time step independently. From the perspective of information bottleneck principle (Tishby et al. (2000), the goal is to selectively extract the "bottle- neck" from high-dimensional sequences that is most useful for world modelling tasks like predicting future events.
Recently, the emergence of Mamba (Gu & Dao 2023) and other SSM-based frameworks (Gu et al, 202 Dao & Gul 2024 Wang et al.) b024b) has garnered widespread attention for their strong performance in efficient sequence modelling. Mamba structure is similar to LSTM (Hochreiter & Schmidhuber 1997). in the sense that it utilises forget gate, an input gate, and output gate. The key difference is that these gates depend only on the previous input (not on the hidden state representing the history of inputs so far). While this hinders their representation power (e.g., cannot solve the parity problem (Hochreiter & Schmidhuber] 1996] Schmidhuber et al. 2007 Srivastava ct al., 2015). this formulation enables parallel computation of selective history compression via sub-lineat sequential attention (Dao & Gu 2024), constructing dependencies between distant data points within the sequence. This sparks their successful applications across various tasks including language modelling (Mehta et al. [202] Grazzi et al., 2024 He et al., 2024), deep noise suppres- sion (Du et al., 2024b), and clinical note understanding (Yang et al.) 2024b). Additionally, many SSM-based vision models have been proposed for tasks such as classification (Du et al.] 2024a]
Shi et al.] 2024| Baron et al., 2023) Huang et al., 20224] Smith et al., 2023| Nguyen et al. 2022).
detection CCen et al.; 2024), segmentation (Yang et al. 2024at Ma et al. (2024), generation ( (Yan) (et al.) 2024; et al. 2024), and video understanding (Islam & Bertasius] 2022] Wang et al.) 2023).
Despite their success, existing SSMs often lack efficient mechanisms for handling high-dimensional
19

---

Published as a conference paper at ICLR 2025
spatial data, relying primarily linear and rigid structural biases (Gu et al. 2021 Dao & Gu 2024).
particularly when dealing with permutable spatial structures.
Multivariate time-series forecasting (MTSF) and object-centric representation learning (OCRL) both involve working with noisy, high-dimensional sequential data, making them critical benchmarks for evaluating world models with significant real-world impact. We chose assess our models on these tasks because they exemplify the primary challenge of spatial-temporal learning. Existing MTSF methods either struggle effectively model long-term temporal dependencies (Salinas al.)
2020 Zhang & Yan 2023; Nie et al. 2023) or fail to effectively leverage the cross-variate regular- ities in high-dimensional imputs (Wu al. 2021 Zhou al.) [022] Wu et al.) 2023a Zeng et al.) 2023 Wang et al. (2024)), resulting in inaccuracies in future- state forecasting. Two notable models in the field, iTrranffrmmer (Liu et al.) 2024) and S-Mamba (Wang et al., 2024c), also have limita- tions. i'Transfomeer suffers from the quadratic scaling of transformers, making it difficult to capture long-term dependencies, while S-Mamba struggles with handling the spatial structures of the data. Object-centric representations, or "slots", are designed to capture "objects". i.e. solving the binding problem (Greff et al. 2020). Our goal is more general: we aim to capture modularities, or "factors", that remain invariant across sequences, framing the discovery of spatial regularities in history com- pression as another instance of the binding problem. Although a philosophical discussion on whether these "factors" should align with common-sense "objects" is beyond the scope of this paper. OCRL is closely related and serves as a good demonstration of our approach. OCRL, originated from the framework (Yuille & Kersten (006). initially applied iet al. 20199; Greff et al. 2019; Locatello et (2020). later extended to videos (Nanbo et al. 2020 Kipl al. 2023). and developed into object-centric world models (Lin et al.) 2020 Kipl et al.
images (Burgess
2019 Wu et al.) 20233 Stanic et al., 2023). Recent OCRL works heavily rely on the Slot Atten-
tion (SA) mechanism (Locatello et al. (2020) for object discovery, which closely related to our routing modules. We view the SA, which also satisfies the LPE and RPI properties,
suitable but computationally expensive alternative to equation 12 addition to MTSF and OCRL, because FACTS leverages its graph-structured memory for apatial-tempore prediction, it is also closely lated to geometric modelling and learning (Bronstein et al. [017]] Wu et al.] (020a)). In this paper, demonstrate FACTS' potential for handling non-Huclidean data through long-term graph node prediction experiments.
It is worth mentioning Goyal et al. (2020 2022). which propose latent state factorisation and equivariance in iitention-augmentte LSTM/GRU--ase frameworks. However, unlike these works, FACTS employs structured state-space memory, enabling dynamic input-to-factor assignments with explicit permutation invariance and efficient training. Recently, SlotSSMs (Jiang et al. 2024) intro- duced factorisation into SSMs by stacking slot attentions on top of Mamba blocks, adding an extra layer to handle modular patterns. In contrast, FACTS offers a fundamentally different state-space formulation, seamlessly integrating factorisation, geometric modelling, and selectivity within a sin- gle block through dynamic memory-input routing. While SlotSSMs can be viewed as special case of FACTS under specific conditions (akin to the fully parallel FACTS), FACTS' versatile architec- ture ensures robust performance across diverse tasks and evolving environments without task- or undality-specife modifications.
B PROOFS
Theorem (Restatement of Theorem 1. FACTS as defined in equation is L.P.E. and R.P.I.
ox E S... Zo. Z, E kRd X2,,,,..., X, E Ruxd By equa- tion 17 equation [8 equation| 19. and equation 20 it sufficient show
Proof. Let az E Se, t xx
oz
EACTS X²).
for all k N, k [1,t].
0z FACTS X²)
=0z (A Xx) 0 Z B X Xt))
=oz =z A(Zo, 002000022222-21022000 0z)
20

---

Published as a conference paper at ICLR 2025
Algorithm FACTS Module: Pseudo Implementation
Input: X1a E [RRxxmxd
D l-sequential axis, m-spatial axis
2: Output: Zr E &Rxxoxd
3: Initialise: Zo [Rx
E.g., init using learnables, or Gaussian samples
4: Param: A a
5:
6: 11= Rearrange(Linear (X1:), m. - ml)
Input projection
7: (X1a, Aia, B Iie, Cia) - Conv2D 1z;; ksz= 1)) sp))ssspi(aaxs==))
D Conv along (-axis
8: (U A C1+) + Rout( Zo. (X 1 B1r, (1)))
D. Rootinn for the elements
9: (Ara B1+)= Discretisation( (SoiolluussSS((( (A 14), Ant, B1:))
10: Zu StateSpaceerooogatiinn/
L1 Art. B1a))
D Zu E RRxxxdd c.f. equation 17 19
11: Zia C1. 1 Zıat
D Zr Rexked for selective output, e.g. Dec Z1+)
12:
Return: Zrc= Zra
Zizi the output; Zi:i the state representation
Since A, B, U are I.P.P.. and R.P.I.,
oz 72;
=ozA(Zo, 00)24 0z)
X2)
=(ooz20, Xx) B Zo. ox Xx) @ UU
= -FACTS 02 24 Xz.
Theorem (Restatement of Theorem any dynamics governed by equation 10 is L.P.E. and R.P.I.
if A.B.U are L.P.E. and R.P.I. functions of Z and X,,
Proof. Let 7--1 Rekxd X2, R.uxd .z E Sk, and ox Sm be matrices. Assume A, B,U are I.P.E. and R.P.I. functions of Z--1 and X, By expanding equation 10
(21)
(22)
=oz (4 (, 00 11
(23)
=02Zi Zt.
(24)
C IMPLEMENTATION DETAILS
C.I FACTS MODULE
We provide the pseudo implementation of FACTS in Algorithm 1 and refer to ? Github page for more details. All results reported for FACTS this paper were generated using a single NVIDIA A A100 GPU (80 GB).
C.2 LONG-TERM MULTIVARIATE TIME-SERIES FORECASTING
Datasets We use the open-source Time Series Library (TSLib), widely-used benchmark for train- ing and evaluating time-series models. TSLib provides standardized settings and a leaderboard of top-performing models, ensuring fair and consistent comparisons. Our focus is on long-term multivariate time-series forecasting (MSTF) tasks. We use collection of diverse real-world datasets, presented in Table 5 These datasets span various domains, including energy produc- tion/consumption, finance exchange rate, traffic monitoring, and weather forecasting, offering diverse range of variates and time granularities. The temporal resolutions of the datasets range from
We omit the reshaping operation of the Conv2D output tensors ([,tt,m) [t. m. *]) for simplicity.
21

---

Published as a conference paper at ICLR 2025
minutes to days, making them ideal for evaluating the performance of forecasting models over both short-term and long-term horizons.
Table 5: The datasets for MTSF evaluation.
Datasets
Variates
Timesteps
Granularity
ETTml & ETTm2
7
17,420
15min
ETTh1 & ETTh2
7
69,680
1hour
Electricity
321
26,304
lhour
Traffic
862
17,544
1hour
Exchange
8
7,588
1day
Weather
52,696
10min
Solar-Energy
137
52,560
10min
Baselines We compare FACTS against 8 baseline models, including state-of-the-art MSTF ap-
proaches that top the TSLib leaderboard (Wang al. 2024c| Liu et al. 2024 Wu et al. 2023a
Nie et al., 2023 Zeng et al. 2023 Zhang & Yan [023 Zhou et al. 2022
Wu et al. 2021). Below
brief summary of 5 most notable baseline approaches:
S-Mamha (Wang et al., 2024c): This baseline adapts Mamba models for MTS data by utilising a bidirectional scan on variates, achieving superior results compared the previous leading method, iTransformer.
Transformer (Liu et al. 2024): An inverted transformer architecture that captures univari- ate history and cross-vas variate dependencies within look-back window, though limited by the quadratic scaling of transformers. i'Trasformer has been leading the long-term forecasting task
TimesNer (Wu et al. 2023a): Specialises in modelling multi-periodicity and interactions among periodic signals in MTS data.
CrossFormer (Zhang & Yan 2023): The emphasis is modelling cross-dimension (spatial)
interactions within MTS data.
PutchTST (Nie et al. 2023): Uses patching techniques to segment sub-time sequences and model channel-wise transitions, improving temporal modelling.
FACTS for MISF Due to the noisy nature of raw input data, a single time step (represented as multivariate vector) often carries limited meaningful information. A common approach handle this is to introduce feature encoders to pre-process the data-as as those temporal and positional em- bedders used the baselines. We employ a set embedder to map the input multivariate sequences of size i x m into / x d, augmenting the tensor with an additional dimension that allows each time step to be represented as a set of m features, each of d-dimensional size. This resulting ten- d, serves the direct input to the FACTS model. For prediction, we adhere to the
sor, t x standard practice in the Time Series Library (TSLib) benchmark, which treats time-series models as encoders designed for single-step predictions, rather than auto-regressive forecasting. We show the MTSF process of FACTS in Algorithm 2 Note that our decoder, namely "factor-graph graph decoder",
Algorithm FACTS for Multivariate Time Series Forecasting
1: Input: 3ia e IR
2: Output: 1 e fxm
3:
4: X1 (21)))
XI plxmxd pee-ppeeoonne pre-normalisation embedding Ze TR txted FACTS() encoder encapsulating Algo.
5: Z FACTS( (
6: Z4+1#4 Predictor Z
Z I fxxxd where f is the prediction length
7: Pt+1dtf FactorGraahDeccoer((( Zeo l.t5)
D Bet /
8: Fultt= Ft+1zef ResidualPredd (x))
Predictive residual RJ feed-through D Invert pre-normalisation
9: Pt+1i-f + PostProcessing(2e... (Za
11: Return: Ie+1:t+1
22

---

Published as a conference paper at ICLR 2025
is erucially designed to be invariant the permutation of factors and processes the latent factors in parallel. Each factor independently makes predictions without relying on others. Specifically, the noftmax-weighten
decoder aggregates the individual predictions of each factor by applying ensuring that the final prediction effectively combines the contributions of all factors while main- taining permutation invariance. We details the aforementioned modules, i.e. the three embedders and the factor-graph decoder in the following list:
sum,
Discrete-Fouries Transform (DFT) decomposer: this embedder applies the Fast Fourier Transform to each variate and decompose it into multiple spectral components. The top-k low-frequency signals are treated as the "trend component", while the high frequency signals represent the "seasonal component", By concatenating the "trend" and "seasonal" components, we embed each univariate signal at each time point into a 2D (or + 1 dimensional, if the top-k frequencies are not combined) vector representation, resulting in set output that is compatible with FACTS.
Conv2d embedder: this embedder applies a convolution along the temporal axis of the input data to aggregate information from nearby time points (i.e., "local context", with the look- back window controlling the kernel size. We implement this using PyTorch's 's standard Conv2d module, setting the kernel size to [lookback, 1] and padding with zeros at the initial time steps. Although such embedders learn the input-feature mapping directly from data, making them flexible in capturing different relationships in different datasets, the choice of look-back window size is often intuitive.
Multi-scale Conv2d (MS-Conv2d) Embedder: this embedder retains the learning flexibility of a Conv2d embedder while extending the single look-back window to multiple scales. By combining different scales, it captures features at varying granularities, making it the most ro- bust among the three embedders and consistently delivering strong results (as shown in Table b
Factor Graph Decoder (FGD): this decoder takes in the a set of predicted latents Zj [2))=14 (c.f. Algorithm 2 for its definition) and first project each 2 E Rd (can run paral- lel) to ā' IR" (the logits) and á Rn (the prediction of the i-th factor). Then the k logits, which correspond to the k factor predictions, will be processed to l categorical probabilities by soft-max function: at softmax (ā/1;=z) The output prediction is the weighted
sum of the factor predictions w.r.t. their corresponding probabilities as: Fy Σx a² a'y. similar to the processes of spatial mixing and alpha blending in the vision and graphics com- munities (Porter& & Duff 1984 [Williams & Titsias 2004 Greff et al.) (217). In vision tasks, the Spatial Broadcast Decoder (SBD, Watters et al. 2019)) which shares similar properties with FGD but more computationally expensive, is more commonly used option.
C.3 OBJECT-CENTRIC WORLD MODELLING
Benchmark We conducted both the slot dynamics prediction on the CLEVRER (Yi al. (020) and OBJ3D (Lin et al., 2020) datasets, and the unsupervised object discovery experiments on both the CLEVRER (Yi et 2020) and MOVi-A (Greff et al., 2022) datasets. These datasets all consist of synthetic vision data and are used as standard benchmark for OCRL research.
Slot-dynamics prediction For the slot dynamics prediction task, we follow the setup in (Wu al., 2023b), filtering out video clips with new objects entering the scene during the rollout period to ensure consistent evaluation setting. The input to FACTS is the latent object representations ex- tracted using pre-trained object-slot encoder (S (SAVi, Kipf et al. (223)) from video frames. That is, we consider the output of the pre-trained SAVi our data in such task same (Wu al. 2023b). We extract the input latent object representations using pre-trained object-slot encoder (SAVi Kipf jet al. (2023) from video frames, and train an auto-regressive roll-outer with a single FACTS layer to predict the latent representations for the next 10 frames, based on the latent codes from 6 ob- served frames. During testing, to ensure a fair comparison with SlotFormer, we burn-in the first frames and roll out (predict) 48 frames. The predieted object representations
visualised using
a pre-trained, frozen SAVi decoder (a spatial broadcast decoder, Wu et al. (2023b)) to render video frames.
Unsupervised object discovery For unsupervised object discovery, we follow the fully- unconditional setting of SAVi, using unbiased slot initialisation and relying solely on RGB video
23

---

Published as a conference paper at ICLR 2025
Leed.Len
MS-Conv2d Emb.
Conv2d Emb.
DFT Emb.
Arg.Std.Err.
96
0.143
0.144
0.147
0.145+0.002
192
0.158
0.159
0.161
0.159=0.001
336
0.171
0.170
0.171
0.171
720
0.198
0.228
0.236
0.219+0.015
Table 6: Ablation: MSTF Performance of FACTS vs. different embedders (MSE)).
k \d
128
512
Arg.
1
0.349
0.330
0.328
0.336=0.007
3
0.349
0.337
0.326
0.337=0.007
5
0.349
0.331
0.349
0.343=0.006
7
0.349
0.331
0,330
0.337=0.006
9
0.349
0.332
0.352
0.344
Avg.:.Std.Err
0.349+0.000
0.332+0.001
0.336=0.006
Table 7: Ablation: MSTF Performance of FACTS (#factors k, #dimensions d) (MSE)).
frames input, with no additional information. The primary modification is replacing SAVI's recurrent slot attention modules with FACTS. Importantly, all of the used modules (CNN vision encoders, FACTS, and decoders) end-to-end in single run without any supervision. We evaluated object discovery under two settings, i.e. video reconstruction and future-frame prediction. For both setting, the input number of frames are set 6, while for prediction FACTS asked predict, in addition to the observed 6 frames, future 10 frames. The prediction setting is conducted with loss function and predictive residual design to emphasise future frame prediction accuracy. The predicted noise will be added to the decoded prediction, i.e., in the pixel space of the future framess similar design also used our MSTF experiments, c.f. the ResidualPred in Algorithm 2 (line 8). Such predictive residual designs allow FACTS to avoid handling pixel-level noise, enabling it to focus on environment dynamics more effectively.
our OCRL experiments, we define such a design mapping between input t-frame video to the pixel-wise noise of future video frames, denoted as PixelResPredichhe experiments, found this quite effective in capturing video dynamics as it sets free the slots (latent states) from capturing objects that are "irrelevant" for modelling their dynamics-i.c. the static ob- jects. The static objects will be merged into the "background". In object reconstruction experiments, object reconstruction experiments, this residual predictor is turned off to ensure a fair comparison.
( e It x f'x x x w. As mentioned, in the future prediction
C.4 DYNAMIC GRAPH NODE PREDICTION
Benchmark We conducted experiments on the METR-LA traffic dataset, which a benchmark data set for traffic prediction, capturing 207 county highway traffic sensor speed observations in Los Angeles metropolitan area. contains traffic data collected from sensors placed on road segments, represented a dynamic graph with nodes corresponding to sensors and edges capturing traffic correlations.
Setup We follow the experimental setup of STGM and incorporate FACTS with masking mech- anism. Specifically, masking is applied to the routing processes described in equation 12 which underpins the construction of selective state-space model (SSM) parameters, A and B. FACTS captures this mask information and integrates it into state-space dynamics modelling, ensuring that only relevant historical data informs the predictions.
D ABLATION STUDY
FACTS requires the of set embedder and the predefined selection of the number of factors in the state-space memory prior to training. Our ablation study aims to investigate the impact of different set embedders and the choice of the number of predefined factors model performance.
Impact of different set encoders We conduct our experiments on the Electricity dataset, testing four prediction lengths: 96, 192, 336, and 720. Three different set encoders evaluated, each
24

---

Published as a conference paper at ICLR 2025
employing different priors: Discrete-Ford Fourier Transform (DFT) decomposer, trainable Conv2d embedder, and a multi-scale Conv2d embedder (inspired by Wu et al. (223a)). Table presents the comparison results, with further details of these embedders provided in Appendix C2 The multi- scale periodic embedder consistently outperforms the others across all prediction lengths, while the DFT-based embedder shows declining performance the prediction length increases. The standard error further indicates that longer forecasting horizons amplify the impact of encoder choice, making it a more critical factor in model aceuracy. This highlights the importance of using unbiased, learnalle set encoder to improve generalisation.
Impact of different number of factors Previous work, such as Mamba (Gu Dao 2023 and xLSTM (Beck et al.) 2024), shows that state-space memory size significantly impacts performance. In FACTS, memory size determined by the number of factors and the dimension of each factor (d). To provide comprehensive analysis, we examine the impact of the preset number of factors (k) on FACTS' performance using the MOVi-A videos and the ETTml dataset. Specifically, we show the impact of k during testing time in the MOVi-A video reconstruction experiments and the robustness of FACTS training time in ETTml MTS forecasting experiments (with 96-prediction length setting).
For the MOVi-A experiments, we took a FACTS model that is trained under the "object dicscovery for video reconstruction" setting and evaluated its video reconstruction performance (measured by the image visual quality measure, LPIPS) against different preset k at testing time on the MOVi- A data. Our results in Figure show that the video reconstruction quality can be improved by increasing & up to certain level (> 11). Knowing that the maximum number of objects in these videos (the true causal factors) is 11 (10 objects plus 1 background), suggesting "sweet point" that is both effective and computationally efficient. However, in practice, identifying such a point in advance can be challenging, Figure 4 indicates that larger is preferred.
As the choice of k could largely affect FACTS's performance, we want to investigate how robust FACTS is against different choices of k. We examined this on the MTS forecasting task using the ETTml data, which consists of 7 variates. To isolate the effect of the number of factors, we train the model across various settings, gridded by different numbers of factors and factor dimensions. As shown in Table 7 FACTS achieves consistent performance across different number of factors (during training) and also the factor dimensions, demonstrating FACTS' robustness to these hyper- parameters.
LPIPS
0.30
Method
ARI t
FG-ARI 1
FG-mIoU t
0.25
LPIPS
G-SWM 2020)
57.14
49,61
24.44
SlotFormer 2023b)
63.45
63.00
29.81
SAVi-dyn (023)
8.64
64.32
18.25
FACTS (Ours)
58.25
62.34
48.11
2
4
8
11
12
16
32
number of slots
Table 8: Quantitative results for slot dynamics prediction: segmentation quality of predicted future frames, measured by ARI, FG-ARI, and FG-mIoU (all reported in %, higher indicate better).
Figure 4: Video reconstruction quality number of slots. With video reconstruction quality mea- sure by LPIPSI.
E ADDITIONAL RESULTS
25

---

Published as a conference paper at ICLR 2025
Pred
Object-0
Object-1
Object-2
1=6
1=10-
12:12
1:
Figure 5: Qualitative results of slot dynamics prediction CLEVRER. The first two columns show the ground truth (GT) and model predictions (Pred) for future frames. The subsequent columns represent independently rendered dynamics of individual objects (Object-O, Object-1, and Object- 2) identified by the model. We show object-centric dynamics in the remaining columns: two columns for each object: the left displays the predicted object dynamies, and the right shows the corresponding object masks.
26

---

Published as a conference paper at ICLR 2025
GT
Pred
Object-0
Object-1
Object-2
18.10 10
1=18-
26
Figure 6: Qualitative results of slot dynamics prediction OBJ3D. The first columns show the ground truth (GT) and model predictions (Pred) for future frames. The subsequent columns represent independently rendered dynamics of individual objects (Object-O, Object-1, and Object- 2) identified by the model. We show object-centric dynamics in the remaining columns: two columns for each object: the left displays the predicted object dynamies, and the right shows the corresponding object masks.
1
Fs 10
ES T6
222
26
30
Video
Video
Figure 7: Qualitative results of slot dynamies prediction on OBJ3D. The top row shows ground truth (GT) video frames, with burn-in frames used for initialisation. The middle presents the predicted future frames (Pred Video) generated by the model. The bottom row illustrates the object segmentation masks predicted by the model.
27

---

Published as a conference paper at ICLR 2025
GT
Recon
Object-0
Object-1
Object-2
[=0
1=3
126
1-0
1=8
1m3
1=0
1:6
Figure 8: Qualitative results of unsupervised object discovery. The first column shows selected ground truth future frames, while the second column presents our predicted future object-cetric states rendered as frames. We show object-centric dynamics in the remaining columns: two columns for each object: the left displays the predicted object dynamics, and the right shows the corresponding attention masks. FACTS effectively discover and captures interpretable "factors". i.e. objects, for modelling video dynamics.
28

---

Published as a conference paper at ICLR 2025
Clanki
WITT
6.00
6.30
3.01
1.00
1.88
01000
6450
340
6.01
2300
L.00
140
6200
0111
2211
1780
1.30
6.30
8.38
0700
6.00
640
6.30
1.00
0111
1100
0110
1100
1100
1.00
3.11
0108
6.00
3.07
0100
6.01
0.40
3300
6.10
1.00
0.00
1.00
0.0
1.00
6.00
0.00
1.60
0108
6.60
3..0
6.00
300
6311
6.00
3.06
6-102
6.00
2301
1.90
0200
0.70
17M
0.08
631
6314
139
0112
0.20
3.0
170
617
k.
1720
6.11
1344
6.00
177
0101
6.70
686
140
0.88
0.06
1.4
601
6.77
027
6417
3.38
0.00
130
0179
6.10
1300
6.71
36.4
6.10
666
13m
0150
642
1401
018
F.06
5.00
8.95
6.71
0221
6.12
010
170
1220
5.30
1.22
1.00
1.00
1.00
6.2%
130
679
610
630
0227
010
030
130
L10
638
618
616
130
179
0401
630
6411
1181
Table 9: Full results for the MTS long-term forecasting task (in MSE, and MAE,). We compare extensive competitive models under different prediction lengths. The input sequence length to 96 for all the datasets above. "Avg" represents the average across all four prediction lengths. For each metric and each dataset, the top performance and the second best are highlighted in and blue, respectively.
29


Reviews:
[{"review_id": "8ahuWB2OhD", "reviewer": "Reviewer_YG1h", "summary": "The paper introduces a framework that constructs a graph-structured memory for predicting the dynamics of complex systems with the selective state space architecture. The proposed architecture incorporates a graph-structured memory that dynamically routes input features to distinct latent factors and hence achieving permutation invariance. The proposed framework leverages parallel computation and selective state-space propagation (via attention) allowing it to handle sequences with complex dynamics efficiently. \n\nThrough experiments on multivariate time series forecasting and object-centric world modeling, the paper shows that FACTS performs competitively or better than state-of-the-art baselines. The interesting result is its robustness to changes in input order, which simulates real-world scenarios where the configuration of inputs may vary unexpectedly.", "questions": "- The SCOFF [1] and Neural Production Systems work [2] is very relevant. These works also proposed the idea of factorizing the latent state and ensuring the latent state is equivariant. These works also use the previous hidden state as a prior for modelling object consistency between consecutive time-steps. The reviewer understands that the proposed work attempts to model the factorization in state space models, whereas the related work mentioned by the reviewer is within attention augmented LSTMs/GRUs.\n\nPlease refer to weakness section.\n\n[1] https://arxiv.org/abs/2006.16225\n[2] https://arxiv.org/abs/2103.01937", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed architecture introduces a permutable memory structure, allowing flexible handling of unordered or dynamically changing inputs. The paper achieves improved performance over baselines by compressing history efficiently, and hence capturing long-term dependencies. \n- The paper is easy to read and comprehend.\n- The results shown on long term forecasting are interesting, and helps the reviewer to understand the implications of the proposed work better (especially forecasting with pre-defined and unknown order).", "weaknesses": "- Object centric video modelling results are a bit weak. it will be interesting to report results also on OBJ3D (another benchmark used in Slotformer paper). It will also be helpful to report downstream results like Predictive VQA on CLEVRER, Physion (similar experiments as in Slotformer paper).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jBxiUbEWkN", "reviewer": "Reviewer_HXLz", "summary": "The paper introduces the FACTored State-space (FACTS) framework to address input feature variance by employing a novel memory-input routing mechanism that ensures robust invariance to input permutations. This approach was empirically validated across tasks like multivariate time series forecasting and object-centric world modeling, demonstrating superior or competitive performance in benchmarks, showcasing its effectiveness and robustness in handling diverse input scenarios.", "questions": "1. What’s the rationale behind designing selectivity through memory-input routing in the manner presented? Can you provide a more detailed justification for the reasonableness of equations 14 and 15? The current explanation appears overly simplistic and could benefit from a deeper analytical discussion.\n\n2. What is the precise meaning of \"permutable state-space memory\" in the context of this work?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper addresses the critical challenge of input feature variance, an interesting issue in spatial-temporal learning, by introducing a novel method that utilizes a memory-input routing mechanism. This approach effectively manages the dynamic relationships between input features, ensuring robust modeling even when input orders change.\n\n2. The proposed FACTS model is both simple and highly effective due to its memory-input routing mechanism, which dynamically assigns input features to latent state-space factors. This straightforward approach simplifies handling high-dimensional data and enhances the model's ability to capture temporal and spatial dependencies, making it practical for real-world applications.\n\n3. The authors have conducted extensive experiments across various tasks, including multivariate time series forecasting and object-centric world modeling. These experiments provide strong empirical support for their claims, demonstrating that FACTS consistently outperforms or matches state-of-the-art models. The thorough evaluation across diverse datasets not only highlights the method's versatility but also its robustness in capturing complex temporal dynamics, further validating its effectiveness in real-world scenarios.", "weaknesses": "For the slot dynamics prediction experiment, the method proposed in the paper relies on a pre-trained encoder and is not end-to-end, which may limit its applicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gIawIoROZS", "reviewer": "Reviewer_Cexo", "summary": "This paper introduces FACTS, a modular or factored SSM model for world modelling applications. FACTS employs an SSM-based recurrent architecture with a permutable memory. This design uses an attention mechanism to dynamically assign input features to latent factors, enabling permutation invariance and efficient long-term prediction.  The experiments show that FACTS demonstrates robust performance on time-series forecasting and object-centric tasks, often outperforming state-of-the-art models.", "questions": "- I would like to ask the authors to clarify the difference between slot dynamics prediction and unsupervised object discovery experiments. It seems that slot dynamics prediction adopts the setup from slotformer which first trains a slot attention model and then uses a transformer to predict future slots conditioned on past slots. From my understanding, the proposed FACTS model is used instead of the transformer in slotformer. Considering this setup for the slot dynamics prediction task, I wonder what is different in the unsupervised object discovery ? Is it that both the slot attention module and FACTS model are trained jointly in this case?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and easy to follow. \n- Modular recurrent architectures are well-studied for world modelling in recent literature, this paper introduces modularity into SSMs while also maintaining their parallel processing capabilities thus the approach seems promising in terms of effeciency.", "weaknesses": "- One integral component of the model is the attention mechanism which assigns input nodes to latent factors. I believe that similar kinds of attention mechanism for the tasks similar to the ones studied in this paper have  already been explored before in various past works [1, 2, 3]. I wonder if the authors could present a comparison of their method to these approaches or atleast highlight the differences. Specifically, [2] proposes to also incorporate modularity and factorization into SSMs, it would be nice to highlight the differences with it and discuss in more detail. \n- The paper says that it is not feasible to compare to other approaches for unsupervised object discovery because of the unique way in which their model works. I would like to point to a few other modular world modelling methods which follow a similar setup as the unsupervised object discovery experiments in the paper which are described in [4]. This work uses a number of metrics which specifically also track moving objects. I wonder if there is any reason for not using the above metrics? Also it would be useful to have even a qualitative comparison to other methdos for unsupervised object discovery, currently there is not baseline used for comparison and only two examples are shown for the proposed approach. \n\n[1] Neural Production Systems https://arxiv.org/abs/2103.01937\n\n[2] Slot state space models https://arxiv.org/abs/2406.12272\n\n[3] Recurrent Independent Mechanisms https://arxiv.org/abs/1909.10893\n\n[4] Benchmarking Unsupervised Object Representations for Video Sequences\n https://arxiv.org/abs/2006.07034", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B6wcjtsmw9", "reviewer": "Reviewer_mfMm", "summary": "In this paper, the authors investigated the factorized State Space Model based World Modeling. Different from the previous works, they applied the attention mechanism to extract the factors from the inputs. With this, they can implement the input-order invariant factorized world model. Additionally, they also showed their modeling could be learned in parallel, so it doesn’t reduce the efficiency from other SSMs. They empirically analyzed their modeling on long-term multivariate time-series forecasting (MSTF) tasks and object-centric world modeling problems, and their model showed on par with the state-of-the-art models or better performance than them.", "questions": "- What is the temporal order change? Is it literally order changing between inputs? Why is it important and why should the world model be invariant for that?\n- In your experimental results, it looks like the interactions between factors are considered in your modeling, but it is not explicitly. How are the interactions considered in your modeling?\n-  In equation 19, isn’t $X_{t:s}$ $X_{s:t}$? Because $s$ is smaller than $t$. And, $\\bar{A}^\\times (Z_0, X_{t:s})$ needs to be $\\bar{A}_{s+1} \\odot \\dots \\odot \\bar{A}_t$?\n- And, $\\bar{A}^\\times (Z_0, X_{t:t})$ looks 1, it needs to be mentioned I think.\n- In equations 16-20, the meaning of them is to replace the recurrent attentions to the attentions between the initial state ($Z_0$) and each input $X_t$, as I understand it. I am okay with the equations, but I am not sure intuitively, how it works (how the recurrent attention can be replaced to the dot products of the attentions between initial state and inputs). Can you explain it more how it works and can you compare the parallel and sequential training for simple tasks?\n- In table 4, more factors make worse performance even the task is multivariate time-series forecasting. I expected there is a sweet spot. Can you explain it? And can you do this ablation studies for unsupervised object discovery problem? The results from the problem could be clearer (e.g., the MSE will be smallest when the number of factors is same with the number of objects or +1 from that for background).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- It is a concrete paper with strong evidences in the theoretical and empirical aspects.\n    - They discussed their modeling clearly, especially, showed the theoretical backgrounds and plenty of empirical results.\n- It is based on a good motivation, which is to design the factorized world model with the contextually represented factors.\n- The experiments are well designed, which is helpful to understand the strengths of their modeling. For instance, in section 4.1.2, they clearly showed the strength of the input invariant factorized world model and in section 4.2, they showed the FACTS performances on the dynamic predictions with the factorized inputs and the factorization performance (Unsupervised object discovery).", "weaknesses": "- The paper highlights the FACTS can model the consistent factor representation regardless of changes in the spatial or temporal order of input features (e.g., contribution summary in Introduction section). The spatial order change is reasonable and the authors showed it as an experiment in section 4.1.2, but for the temporal order change, it is unclear what it is and why it is important at least for me.\n- The interactions between factors are not clearly designed in their modeling, FACTS. The FACTS is to update the hidden with the factors from the inputs through the attention mechanism, but there is not explicitly designed the module or mechanism to learn the interactions between factors.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)