FACTS: A Factored State-Space Framework for World Modelling

"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?"

The reviewer is right here. The current phrasing can be misleading. Invariance in the temporal order is not what we are trying to say here. In the third summary bullet point, we wanted to emphasize that FACTS achieves consistent factor representations over time, regardless of changes in the spatial order of input features over time. In the updated version, we have rephrased the text to avoid confusion.

"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?"

We thank the reviewer for this valuable question. Using FACTS, dynamic interactions among factors can be achieved by stacking multiple FACTS layers, and importantly, this can be done without disrupting parallel computation. For example, consider a two-FACTS-layer encoder: A FACTS operator takes two inputs—memory and input. Suppose the input signals have the shape $t \times m \times d$ (i.e., [seq_length x num_tokens x token_dim]). The first FACTS layer binds the inputs to their memory (e.g. $Z_{in}$ in parallel FACTS), which can either be the input itself or self-initialized, and outputs a set of $k$ factors, organized into a tensor $Z_{out}$ with the shape $t \times k \times d$. By “interaction of factors,” we refer to the modeling of information flow along the $k$ factors. In this setup, the second FACTS layer can take the first layer’s output, $Z_{out}$, as both its memory (which may need to be detached from the gradient graph during training) and input, performing the routing operation described in Eqn. 12, i.e., \text{rout}($Z_{out}$, $Z_{out}$). This effectively simulates self-attention operations and enables information flow across the $k$ factors/particles. Furthermore, the summation operation in Equation 19 will aggregate such interactive dynamics along the time-axis for state propagation. We will add this discussion to the paper.

"In equation 19, isn’t $X\_{t:s}$ $X_{s:t}$? Because $s$ is smaller than $t$. And, $\bar{\mathbf{A}}^{\times}(Z\_0, X\_{t:s})$ needs to be $\bar{\mathbf{A}}\_{s+1} \odot \cdots \odot \bar{\mathbf{A}}\_t$?"

We appreciate the reviewer for pointing out such detail. Both $\bar{\mathbf{A}}^{\times}\_{t:s}$ and $\bar{\mathbf{A}}\_{s:t}^{\times}$ are correct, as they lead to the same calculation due to commutativity of element-wise product: $\bar{ \mathbf{A} }^{\times}(Z\_0, X\_{t:s}) =\bar{\mathbf{A}}\_{t+1}\odot \bar{\mathbf{A}}\_{t}\odot\bar{\mathbf{A}}\_{t-1} \cdots \odot \bar{\mathbf{A}}\_{s+1}$. This difference is purely a matter of notational convention. We choose $\bar{**A**\_{t:s}}$ because it aligns more naturally with the subsequent terms, $\bar{**B**\_{s}}$ and $\bar{**U**\_{s}}$ , maintaining a descending order of subscripts from $t$ to $s$: $\bar{\mathbf{A}}^{\times}\_{t:s} \odot \bar{\mathbf{B}}\_{s} \odot \bar{\mathbf{U}}\_{s}$. This also ensures consistency with the conventions of other linearized SSMs (as shown in Equation 5). We have clarified this in the updated draft (line 280) for better readability.

"And, $\bar{\mathbf{A}}^{\times}(Z\_0, X\_{t:t})$ looks 1, it needs to be mentioned I think."

We appreciate the reviewer again for pointing out such detail. We have added relevant item in line 280 in the updated version.

"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?"

We thank the reviewer for raising this insightful question. The parallelization design outlined in our Equations 16–20 not only enhances computational efficiency but also provides flexibility for recurrent applications of the FACTS core operation. Specifically, while we denote $Z\_0$ as the “initial state,” it does not necessarily have to represent the actual initial state of a long sequence. Instead, a long sequence can be divided into segments based on a chosen strategy or algorithm, allowing FACTS to run in a parallelized manner within each segment. I.e. computations within individual segments are parallelised, while dependencies across the original long sequence remain sequential. In such cases, $Z\_0$ would represent the final memory state of the preceding segment. This design essentially draws inspiration from LSTM architectures and skip connections in ResNet. To provide further clarity, we have included an analysis of FACTS’ long-term forecasting performance, comparing its fully sequential mode to its fully parallel mode (adjusted by the segment window size) on the MTS Electricity dataset. These results are detailed in Figure 3. Each window size corresponds to a different update frequency of the “$Z\_0$”. For instance, windows_size $1$ corresponds to fully-recurrent FACTS while windows_size $96$ (the sequence length) corresponds to fully-parallel FACTS. As shown in Figure 3, the models achieved consistent performance over different segment window sizes.

"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)."

We thank the reviewer for raising this interesting point and valuable suggestion. We have included analysis of the impact of the preset number of slots on FACTS for unsupervised object discovery -- during testing time. Specifically, we took a FACTS model that is trained under the "object dicscovery for video reconstruction" setting (using $k=11$ for training) 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 4 show that the video reconstruction quality can be improved by increasing $k$ 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 a ``sweet point'' that is both effective and computationally efficient.

Importantly, we note that this does not conflict with our MTS results, which examine the impact of different $k$ during training. From a practical perspective, the observation that increasing factors can worsen performance in MTS experiments is understandable, as more factors/parameters can lead to overfitting, especially with noisy real-world MTS data. From a causal perspective, "sweet point" in training may imply finding the true number of factors and potentially “true causal relations” among them. However, identifying causal structures or discovering physical laws solely from observable data is often intractable -- particularly when dealing with noisy real-world data (e.g., MTS data). This discrepancy explains why the best “number of factors” in our MTS experiments do not fully match our expectations. The practical robustness of FACTS against different $k$ in training, as evidenced by the standard errors in Table 6 (updated version), may ultimately be more valuable for a machine learning model.

"One integral component of the model ... similar kinds ... studied ... in various past works ... if the authors could present a comparison ... or at least highlight the differences. Specifically, [2] ... in more detail. "

Thank you for bringing to our attention these works, specifically SlotSSMs. First note that according to ICLR policy (https://iclr.cc/Conferences/2025/ReviewerGuide), papers are considered contemporaneous if published (i.e., at a peer-reviewed venue) within the last 4 months and the slotSSM does not meet this criteria. However, to satisfy the reviewer and to clarify the distinction between Slot State-Space Models (SlotSSM) and FACTS, we highlight several key methodological differences:

1. Intrinsic State-Space Design: While SlotSSMs build upon Mamba by introducing pre- and post-processing steps to make it suitable for modular data, FACTS presents an intrinsic general formulation of state-space models (alternative to Mamba directly) that incorporates modularity and permutation invariance. This alternative design inherently supports factorisation and dynamically-changing input relationships by construction, eliminating the need for additional processing layers.
2. Input-Only Selectivity vs Input-Memory Selectivity: Like Mamba, SlotSSMs base selectivity purely on the input data, assigning features to latent factors based on their relevance. In contrast, FACTS extends this concept by incorporating selectivity based on both input features and memory, achieved through an attention-augmented memory-input routing mechanism. This allows FACTS to better capture long-term dependencies and refine its memory representations over time.
3. Dynamic Routing vs Fixed Assignments: SlotSSMs rely on pre-SSM assignments between input features and latent states using inverted cross-attention. This rigidity can limit their flexibility in scenarios where input-to-latent relationships evolve over time. In contrast, FACTS formulation (Eqn. 10-11) allows to employ a dynamic memory-input routing mechanism, enabling it to assign input features to latent factors based on their relevance at each time step. This enhances FACTS' flexibility and adaptability, especially in scenarios with evolving input relationships.
4. Special Case Relationship: Ignoring the memory selectively, in an abstract level, "SlotEncoder + SlotSSM" can be interpreted as a special case of FACTS (plus an additional task-dependent adaptation). This corresponds to our fully parallel variant of FACTS (as described in equations 16-20). In contract, FACTS allows recurrent processing, with stable performance, as presented in Figure 3 in the updated manuscript.
5. Robustness in Diverse Scenarios: FACTS is tailored to a broader range of forecasting-based application, due to its flexible state-space memory design. SlotSSMs, on the other hand, focus primarily on adapting Mamba for modularity in slot-based inputs for object-centric learning, limiting their application scope, i.e., Facts is general-purpose SSM framework while SlotSSM requires application-specific formulation for each task, e.g., OC-SlotSSM in [2].

To illustrate this empirically, here we compare SoltSSM to FACTS on 3 different tasks. We use same experimental setup as described in our paper for both methods with hyper-parameter search:

• On MTS forecasting (Electricity, pred_len set to 720):

• Unsupervised object discovery with MOVi-A (Note that for FACTS training is not completed yet):

• Graph-data with METR-LA traffic forecasting for 1 hour (i.e.12-steps) ahead:

Note with FACTS clearly outperforms SlotSSM on 2 of the three tasks, with less than half the error rates on the graph-based METR-LA dataset. On MOVi-A, where SlotSSM slightly outperforms FACTS:

• FACTS, despite its general-purpose design, achieves nearly similar performance to the task-specific SlotSSM.
• As SlotSSM can be seen as a special case of FACTS, we are confident that, with more extensive hyperparameter search and longer training, FACTS can match or surpass SlotSSM.

To sum up, while there is some conceptual overlap between FACTS and SlotSSMs: while SlotSSM adapts Mamba for modular slot-based data with additional processing modules, FACTS provides a fundamentally different formulation of SSMs, intrinsically handling evolving input relationships with broader adaptability and strong performance multiple tasks. Furthermore, unlike SlotSSM's input-only selective routing, FACTS incorporates both input and memory through an attention-augmented routing mechanism. We have incorporated this discussion into the updated version of the paper.

"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."

There seems to be a misunderstanding. As a general and versatile framework, our extensive experiments show that FACTS is applicable for different tasks and different data modalities (time-series, videos, graphs). Specifically for the reviewer's interest in vision data, as discussed in Sec 4.2, we clarify that we conducted two groups of vision experiments:

1. For slot-dynamics prediction, as mentioned in the paper, for fair comparison, we used SlotFormer’s exact setting with their pre-trained encoder and decoder. The results of our slot-dynamics experiments with CLEVRER and OBJ3D datasets are presented in Table 2, clearly showing the superiority of FACTS.
2. For the unsupervised object discovery experiments, all of the used modules (CNN vision encoders, FACTS, and decoders) are trained end-to-end in a single run without any supervision. The results on unsupervised object discovery are presented in the Figure 7-8, this shows that FACTS can successfully identify objects for world modelling and indeed be trained end-to-end with other modules.

Nonetheless, for improved clarity, we've updated the discussion in section 4.2 and provided more details of our experimental setup in Appendix C.3.

"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."

Recent works on SSM, e.g. Mambas [1], highlight why input should be part of the selectivity design. In this work, we take this a step further and say that we need selectivity based on both input and memory. This can be efficiently done with an attention-augmented memory-input routing mechanism. This explains the main reasoning behind Eqn.14 and Eqn.15. Such formulation of the input-memory coupling in the selectivity process: i) allow us to design a permutable but structure persevering memory in an end-to-end manner. ii) provide more training feedback for the memory parameters due to the additional gradient generated through Eqn. 14-15. This leads to better generalization and stronger performance, as shown in all our results, particularly Table 1, by comparing FACTs to ‘input-only selectively’ approaches, e.g., S-mamba.

"What is the precise meaning of "permutable state-space memory" in the context of this work?"

Thank you for your question. In this work, "state-space memory" is the $Z_t$ and "permutable state-space memory" means that instead of assuming a fixed structure or order for the latent state representations $Z_t$, FACTS allows these representations to be updated flexibly, where the input features can be permuted without compromising the consistency of the outputs (as shown in Figure 2).

This enables the model to leverage the concept of history compression and learn consistent representations over time. The permutable memory structure is key to efficiently capture dynamic dependencies and unleashes the full potential of the model as it does need to capture order-related characteristics of the input and of the memory. This simplifies the learning task allowing FACTS to maintain strong performance and adaptability in real-world modeling scenarios.

"Object centric video modelling results ... it will be interesting to report results also on OBJ3D (another benchmark ... It will also be helpful to report downstream results like Predictive VQA on CLEVRER, Physion (similar experiments as in Slotformer paper). "

We appreciate the reviewer’s suggestion regarding the inclusion of additional benchmarks and downstream tasks. In the current empirical validation of FACTS, we wanted to emphasize its predictive power to demonstrate its flexibility as a powerful general framework able to handle different input types/modalities and efficiently solve diverse forecasting-based tasks. To this end, we opted to evaluate it across diverse tasks (without emphasis on a particular one): multivariate time-series forecasting, object-centric video modeling, etc.. However, to satisfy the reviewer, we now also include results on another forecasting task, namely object dynamics modeling with OBJ3D dataset. As can be seen through Table 2, the results are consistent with the CLEVRER dataset, FACTS outperforms existing state-of-the-art competing methods designed for these tasks, despite its general purpose design. For the visual results, we now include two more demonstrative examples for object dynamics modeling: (Figures 5 for CLEVRER) and (Figures 6-7 for OBJ3D) that visually highlight FACTS world modeling abilities.

Furthermore, to further showcase FACTS’s versatility, in the revised manuscript, we now include results with a different input modality, particularly dynamic-graph input data evaluated on node prediction task (long-term prediction (12-step MAE)) with METR-LA dataset. As can be seen in Table 3. FACTS, leveraging our graph-structured memory, also outperforms all existing methods on this task. This further corroborated our main claim that FACTS is indeed a versatile worldmodel framework with consistent strong performance in several diverse forecasting tasks.

"The SCOFF [1] and Neural Production Systems work [2] is very relevant... 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."

We thank the reviewer for bringing to our attention the works of SCOFF [1] and Neural Production Systems [2]. These works indeed share similarities with FACTS in their focus on factoring latent states and ensuring equivariance, particularly in attention-augmented LSTM/GRU-based architectures. While our approach differs fundamentally in its state-space modelling foundation, we acknowledge the conceptual overlap and have now included a discussion of these works in the Related Work section in Appendix.