FACTS: A Factored State-Space Framework for World Modelling

## Abstract

Abstract The abstract paragraph should be indented 1/2 inch (3 picas) on both left and
right-hand margins. Use 10 point type, with a vertical spacing of 11 points.
The word Abstract must be centered, in small caps, and in point size 12. Two
line spaces precede the abstract. The abstract must be limited to one
paragraph.

## 1 Submission of conference papers to ICLR 2024

ICLR requires electronic submissions, processed by
[https://openreview.net/](https://openreview.net/). See ICLR’s website for more instructions.

If your paper is ultimately accepted, the statement \iclrfinalcopy should be inserted to adjust the
format to the camera ready requirements.

The format for the submissions is a variant of the NeurIPS format.
Please read carefully the instructions below, and follow them
faithfully.

### 1.1 Style

Papers to be submitted to ICLR 2024 must be prepared according to the
instructions presented here.

Authors are required to use the ICLR LaTeX style files obtainable at the
ICLR website. Please make sure you use the current files and
not previous versions. Tweaking the style files may be grounds for rejection.

### 1.2 Retrieval of style files

The style files for ICLR and other conference information are available online at:

[http://www.iclr.cc/](http://www.iclr.cc/)

The file iclr2024_conference.pdf contains these
instructions and illustrates the
various formatting requirements your ICLR paper must satisfy.
Submissions must be made using LaTeX and the style files
iclr2024_conference.sty and iclr2024_conference.bst (to be used with LaTeX2e). The file
iclr2024_conference.tex may be used as a “shell” for writing your paper. All you
have to do is replace the author, title, abstract, and text of the paper with
your own.

The formatting instructions contained in these style files are summarized in
sections 2, 3, and 4 below.

## 2 General formatting instructions

The text must be confined within a rectangle 5.5 inches (33 picas) wide and
9 inches (54 picas) long. The left margin is 1.5 inch (9 picas).
Use 10 point type with a vertical spacing of 11 points. Times New Roman is the
preferred typeface throughout. Paragraphs are separated by 1/2 line space,
with no indentation.

Paper title is 17 point, in small caps and left-aligned.
All pages should start at 1 inch (6 picas) from the top of the page.

Authors’ names are
set in boldface, and each name is placed above its corresponding
address. The lead author’s name is to be listed first, and
the co-authors’ names are set to follow. Authors sharing the
same address can be on the same line.

Please pay special attention to the instructions in section 4
regarding figures, tables, acknowledgments, and references.

There will be a strict upper limit of 9 pages for the main text of the initial submission, with unlimited additional pages for citations.

## 3 Headings: first level

First level headings are in small caps,
flush left and in point size 12. One line space before the first level
heading and 1/2 line space after the first level heading.

### 3.1 Headings: second level

Second level headings are in small caps,
flush left and in point size 10. One line space before the second level
heading and 1/2 line space after the second level heading.

#### 3.1.1 Headings: third level

Third level headings are in small caps,
flush left and in point size 10. One line space before the third level
heading and 1/2 line space after the third level heading.

## 4 Citations, figures, tables, references

These instructions apply to everyone, regardless of the formatter being used.

### 4.1 Citations within the text

Citations within the text should be based on the natbib package
and include the authors’ last names and year (with the “et al.” construct
for more than two authors). When the authors or the publication are
included in the sentence, the citation should not be in parenthesis using \citet{} (as
in “See \citetHinton06 for more information.”). Otherwise, the citation
should be in parenthesis using \citep{} (as in “Deep learning shows promise to make progress
towards AI \citepBengio+chapter2007.”).

The corresponding references are to be listed in alphabetical order of
authors, in the References section. As to the format of the
references themselves, any style is acceptable as long as it is used
consistently.

### 4.2 Footnotes

Indicate footnotes with a number(^1^11Sample of the first footnote) in the
text. Place the footnotes at the bottom of the page on which they appear.
Precede the footnote with a horizontal rule of 2 inches
(12 picas).(^2^22Sample of the second footnote)

### 4.3 Figures

All artwork must be neat, clean, and legible. Lines should be dark
enough for purposes of reproduction; art work should not be
hand-drawn. The figure number and caption always appear after the
figure. Place one line space before the figure caption, and one line
space after the figure. The figure caption is lower case (except for
first word and proper nouns); figures are numbered consecutively.

Make sure the figure caption does not get separated from the figure.
Leave sufficient space to avoid splitting the figure and figure caption.

You may use color figures.
However, it is best for the
figure captions and the paper body to make sense if the paper is printed
either in black/white or in color.

Figure: Figure 1: Sample figure caption.

### 4.4 Tables

All tables must be centered, neat, clean and legible. Do not use hand-drawn
tables. The table number and title always appear before the table. See
Table 1.

Place one line space before the table title, one line space after the table
title, and one line space after the table. The table title must be lower case
(except for first word and proper nouns); tables are numbered consecutively.

**Table 1: Sample table title**
| PART | DESCRIPTION |
| --- | --- |
| Dendrite | Input terminal |
| Axon | Output terminal |
| Soma | Cell body (contains cell nucleus) |

## 5 Default Notation

In an attempt to encourage standardized notation, we have included the
notation file from the textbook, Deep Learning
[goodfellow2016deep] available at
[https://github.com/goodfeli/dlbook_notation/](https://github.com/goodfeli/dlbook_notation/). Use of this style
is not required and can be disabled by commenting out
math_commands.tex.

Numbers and Arrays

| $\displaystyle a$ | A scalar (integer or real) |
| --- | --- |
| $\displaystyle{\bm{a}}$ | A vector |
| $\displaystyle{\bm{A}}$ | A matrix |
| $\displaystyle{\bm{\mathsfit{A}}}$ | A tensor |
| $\displaystyle{\bm{I}}_{n}$ | Identity matrix with $n$ rows and $n$ columns |
| $\displaystyle{\bm{I}}$ | Identity matrix with dimensionality implied by context |
| $\displaystyle{\bm{e}}^{(i)}$ | Standard basis vector $[0,\dots,0,1,0,\dots,0]$ with a 1 at position $i$ |
| $\displaystyle\text{diag}({\bm{a}})$ | A square, diagonal matrix with diagonal entries given by ${\bm{a}}$ |
| a | A scalar random variable |
| $\displaystyle{\mathbf{a}}$ | A vector-valued random variable |
| $\displaystyle{\mathbf{A}}$ | A matrix-valued random variable |

Sets and Graphs

| $\displaystyle{\mathbb{A}}$ | A set |
| --- | --- |
| $\displaystyle\mathbb{R}$ | The set of real numbers |
| $\displaystyle\{0,1\}$ | The set containing 0 and 1 |
| $\displaystyle\{0,1,\dots,n\}$ | The set of all integers between 0 0 and $n$ |
| $\displaystyle[a,b]$ | The real interval including $a$ and $b$ |
| $\displaystyle(a,b]$ | The real interval excluding $a$ but including $b$ |
| $\displaystyle{\mathbb{A}}\backslash{\mathbb{B}}$ | Set subtraction, i.e., the set containing the elements of ${\mathbb{A}}$ that are not in ${\mathbb{B}}$ |
| $\displaystyle{\mathcal{G}}$ | A graph |
| $\displaystyle Pa_{\mathcal{G}}({\textnormal{x}}_{i})$ | The parents of ${\textnormal{x}}_{i}$ in ${\mathcal{G}}$ |

Indexing

| $\displaystyle{a}_{i}$ | Element $i$ of vector ${\bm{a}}$, with indexing starting at 1 |
| --- | --- |
| $\displaystyle{a}_{-i}$ | All elements of vector ${\bm{a}}$ except for element $i$ |
| $\displaystyle{A}_{i,j}$ | Element $i,j$ of matrix ${\bm{A}}$ |
| $\displaystyle{\bm{A}}_{i,:}$ | Row $i$ of matrix ${\bm{A}}$ |
| $\displaystyle{\bm{A}}_{:,i}$ | Column $i$ of matrix ${\bm{A}}$ |
| $\displaystyle{\mathsfit{A}}_{i,j,k}$ | Element $(i,j,k)$ of a 3-D tensor ${\bm{\mathsfit{A}}}$ |
| $\displaystyle{\bm{\mathsfit{A}}}_{:,:,i}$ | 2-D slice of a 3-D tensor |
| $\displaystyle{\textnormal{a}}_{i}$ | Element $i$ of the random vector ${\mathbf{a}}$ |

Calculus

| $\displaystyle\frac{dy}{dx}$ | Derivative of $y$ with respect to $x$ |
| --- | --- |
| $\displaystyle\frac{\partial y}{\partial x}$ | Partial derivative of $y$ with respect to $x$ |
| $\displaystyle\nabla_{\bm{x}}y$ | Gradient of $y$ with respect to ${\bm{x}}$ |
| $\displaystyle\nabla_{\bm{X}}y$ | Matrix derivatives of $y$ with respect to ${\bm{X}}$ |
| $\displaystyle\nabla_{\bm{\mathsfit{X}}}y$ | Tensor containing derivatives of $y$ with respect to ${\bm{\mathsfit{X}}}$ |
| $\displaystyle\frac{\partial f}{\partial{\bm{x}}}$ | Jacobian matrix ${\bm{J}}\in\mathbb{R}^{m\times n}$ of $f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{m}$ |
| $\displaystyle\nabla_{\bm{x}}^{2}f({\bm{x}})\text{ or }{\bm{H}}(f)({\bm{x}})$ | The Hessian matrix of $f$ at input point ${\bm{x}}$ |
| $\displaystyle\int f({\bm{x}})d{\bm{x}}$ | Definite integral over the entire domain of ${\bm{x}}$ |
| $\displaystyle\int_{\mathbb{S}}f({\bm{x}})d{\bm{x}}$ | Definite integral with respect to ${\bm{x}}$ over the set ${\mathbb{S}}$ |

Probability and Information Theory

| $\displaystyle P({\textnormal{a}})$ | A probability distribution over a discrete variable |
| --- | --- |
| $\displaystyle p({\textnormal{a}})$ | A probability distribution over a continuous variable, or over<br>a variable whose type has not been specified |
| $\displaystyle{\textnormal{a}}\sim P$ | Random variable a has distribution $P$ |
| $\displaystyle\mathbb{E}_{{\textnormal{x}}\sim P}[f(x)]\text{ or }\mathbb{E}f(x)$ | Expectation of $f(x)$ with respect to $P({\textnormal{x}})$ |
| $\displaystyle\mathrm{Var}(f(x))$ | Variance of $f(x)$ under $P({\textnormal{x}})$ |
| $\displaystyle\mathrm{Cov}(f(x),g(x))$ | Covariance of $f(x)$ and $g(x)$ under $P({\textnormal{x}})$ |
| $\displaystyle H({\textnormal{x}})$ | Shannon entropy of the random variable x |
| $\displaystyle D_{\mathrm{KL}}(P\|Q)$ | Kullback-Leibler divergence of P and Q |
| $\displaystyle\mathcal{N}({\bm{x}};{\bm{\mu}},{\bm{\Sigma}})$ | Gaussian distribution over ${\bm{x}}$ with mean ${\bm{\mu}}$ and covariance ${\bm{\Sigma}}$ |

Functions

| $\displaystyle f:{\mathbb{A}}\rightarrow{\mathbb{B}}$ | The function $f$ with domain ${\mathbb{A}}$ and range ${\mathbb{B}}$ |
| --- | --- |
| $\displaystyle f\circ g$ | Composition of the functions $f$ and $g$ |
| $\displaystyle f({\bm{x}};{\bm{\theta}})$ | A function of ${\bm{x}}$ parametrized by ${\bm{\theta}}$.<br>(Sometimes we write $f({\bm{x}})$ and omit the argument ${\bm{\theta}}$ to lighten notation) |
| $\displaystyle\log x$ | Natural logarithm of $x$ |
| $\displaystyle\sigma(x)$ | Logistic sigmoid, $\displaystyle\frac{1}{1+\exp(-x)}$ |
| $\displaystyle\zeta(x)$ | Softplus, $\log(1+\exp(x))$ |
| $\displaystyle||{\bm{x}}||_{p}$ | $L^{p}$ norm of ${\bm{x}}$ |
| $\displaystyle||{\bm{x}}||$ | $L^{2}$ norm of ${\bm{x}}$ |
| $\displaystyle x^{+}$ | Positive part of $x$, i.e., $\max(0,x)$ |
| $\displaystyle\bm{1}_{\mathrm{condition}}$ | is 1 if the condition is true, 0 otherwise |

## 6 Final instructions

Do not change any aspects of the formatting parameters in the style files.
In particular, do not modify the width or length of the rectangle the text
should fit into, and do not change font sizes (except perhaps in the
References section; see below). Please note that pages should be
numbered.

## 7 Preparing PostScript or PDF files

Please prepare PostScript or PDF files with paper size “US Letter”, and
not, for example, “A4”. The -t
letter option on dvips will produce US Letter files.

Consider directly generating PDF files using pdflatex
(especially if you are a MiKTeX user).
PDF figures must be substituted for EPS figures, however.

Otherwise, please generate your PostScript and PDF files with the following commands:

### 7.1 Margins in LaTeX

Most of the margin problems come from figures positioned by hand using
\special or other commands. We suggest using the command
\includegraphics
from the graphicx package. Always specify the figure width as a multiple of
the line width as in the example below using .eps graphics

or

for .pdf graphics.
See section 4.4 in the graphics bundle documentation ([http://www.ctan.org/tex-archive/macros/latex/required/graphics/grfguide.ps](http://www.ctan.org/tex-archive/macros/latex/required/graphics/grfguide.ps))

A number of width problems arise when LaTeX cannot properly hyphenate a
line. Please give LaTeX hyphenation hints using the \- command.

#### Author Contributions

If you’d like to, you may include a section for author contributions as is done
in many journals. This is optional and at the discretion of the authors.

#### Acknowledgments

Use unnumbered third level headings for the acknowledgments. All
acknowledgments, including those to funding agencies, go at the end of the paper.

## Appendix A Appendix

You may include other additional sections here.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.

Please refer to weakness section.

[1] https://arxiv.org/abs/2006.16225
[2] https://arxiv.org/abs/2103.01937
soundness: 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. 
- The paper is easy to read and comprehend.
- 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).
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. 

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

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
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.

2. What is the precise meaning of "permutable state-space memory" in the context of this work?
soundness: 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.

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

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

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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?
soundness: 3
strengths: - The paper is well-written and easy to follow. 
- 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.
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.
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. 
- 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. 

[1] Neural Production Systems https://arxiv.org/abs/2103.01937

[2] Slot state space models https://arxiv.org/abs/2406.12272

[3] Recurrent Independent Mechanisms https://arxiv.org/abs/1909.10893

[4] Benchmarking Unsupervised Object Representations for Video Sequences
 https://arxiv.org/abs/2006.07034

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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?
- 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?
-  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$?
- And, $\bar{A}^\times (Z_0, X_{t:t})$ looks 1, it needs to be mentioned I think.
- 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?
- 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).
soundness: 3
strengths: - It is a concrete paper with strong evidences in the theoretical and empirical aspects.
    - They discussed their modeling clearly, especially, showed the theoretical backgrounds and plenty of empirical results.
- It is based on a good motivation, which is to design the factorized world model with the contextually represented factors.
- 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).
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.
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.
- 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.

Meta-review: The paper proposes FACTS, a state-space framework that effectively combines graph-structured memory and routing mechanisms to achieve both permutation invariance and parallel computation. The work demonstrates strong empirical validation across time series forecasting and object-centric tasks, with clear methodology and significant computational efficiency gains. Their comprehensive experiments show competitive or superior performance compared to specialized models, and the authors thoroughly addressed reviewer concerns during rebuttal. The parallel computation capability while maintaining permutation invariance represents a meaningful advance in world modeling. This innovative combination of technical elements, backed by strong empirical results and clear practical benefits, makes it a valuable contribution worthy of acceptance.

Decision: Accept (Poster)