SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Trang Nguyen*
FPT Software AI Center
trang.nguyen@mila.quebec
Amin Mansouri
amin.mansouri@mila.quebec
Kanika Madan
madankan@mila.quebec
Khuong Nguyen
khuongnd6@fpt.com
Dianbo Liu
dianbo.liu@mila.quebec
Yoshua Bengio
Université de Montréal, CIFAR AI Chair
yoshua.bengio@mila.quebec
Agents with the ability to comprehend and reason about the dynamics of objects would be expected to exhibit improved robustness and generalization in novel scenarios. However, achieving this capability necessitates not only an effective scene representation but also an understanding of the mechanisms governing interactions among object subsets. Recent studies have made significant progress in representing scenes using object slots. In this work, we introduce Reusable Slotwise Mechanisms, or RSM, a framework that models object dynamics by leveraging communication among slots along with a modular architecture capable of dynamically selecting reusable mechanisms for predicting the future states of each object slot. Crucially, RSM leverages the Central Contextual Information (CCI), enabling selected mechanisms to access the remaining slots through a bottleneck, effectively allowing for modeling of higher order and complex interactions that might require a sparse subset of objects. Experimental results demonstrate the superior performance of RSM compared to state-of-the-art methods across various future prediction and related downstream tasks, including Visual Question Answering and action planning. Furthermore, we showcase RSM's Out-of-Distribution generalization ability to handle scenes in intricate scenarios.
Accurate prediction of future frames and reasoning over objects is crucial in various computer vision tasks. These capabilities are essential for constructing comprehensive world models in applications like autonomous driving and reinforcement learning for robots. Traditional deep learning-based representation learning methods compress entire scenes into monolithic representations, lacking
*Work done under an internship at Mila - Quebec AI Institute and FPT Software AI Residency Program. Corresponding author.
compositionality and object-centric understanding. As a result, these representations struggle with systematic generalization, interpretability, and capturing interactions between objects. This limitation leads to poor generalization performance as causal variables become entangled in non-trivial ways.
There has been growing interest in slot-based and modular representations that decompose scenes into individual entities, deviating from fixed-size monolithic feature vector representations [19, 38, 16, 17, 15, 18, 32, 21, 28, 37, 39, 48, 30]. These novel approaches offer significantly more flexibility when dealing with environments that comprise multiple objects. By employing an encoder that segments a scene into its independent constituent entities instead of compressing information into a fixed-size representation, these methods allow for greater flexibility and parameter sharing when learning object-centric representations, and their compositional nature enables better generalization. Compositional and object-centric representations can be effectively utilized alongside complex world models that accurately capture the interactions and dynamics of different entities in a scene.
These world models, when presented with proper representations, in principle, can model the transition functions that relate latent causal factors across consecutive time steps of a rollout. While monolithic blocks are still used occasionally with object-centric methods [44], recent attempts have incorporated similar inductive biases related to the object-centricity of images in modeling interactions. Structured world models and representations seem truly promising for systematically generalizing to novel scenes. Structured world models would ideally decompose the description of the evolution of a scene into causal and independent sub-modules, making it easy to recombine and repurpose those mechanisms in novel ways to solve challenges in unseen scenarios. Such separation of dynamics modeling makes structured world models more adaptable to distribution shifts, as only the parameters of a few mechanisms that have changed in a new environment would have to be retrained, and not all of the parameters in the case of a monolithic model [8].
A major class of such structured world models aims at baking in some inductive bias about the nature of object interactions. On one extreme, there have been studies that employ Graph Neural Networks (GNNs) to capture object dependencies through dense connections, while on the other hand, there has been contrasting work aiming at modeling the dynamics through only pairwise interactions. We believe, however, that ideally, an agent should be able to learn, select, and reuse a set of prediction rules based on contextual information and the characteristics of each object.
In this work, we argue that the assumptions made in previous attempts at learning the dynamics among slots may be insufficient in more realistic domains. To address these limitations, we propose Reusable Slotwise Mechanisms (RSM), a novel modular architecture incorporating a set of deep neural networks representing reusable mechanisms on top of slotwise representations [31, 10]. Inspired by the Global Workspace Theory (GWT) in the cognitive neuroscience of working memory [2, 3], we introduce the concept of Central Contextual Information (CCI), which allows each reusable mechanism, i.e., a possible explanation of state evolution, to access information from all other slots through a bottleneck, enabling accurate predictions of the next state for a specific slot. The CCI's bottleneck amounts to a relaxed inductive bias compared to the extreme cases of pairwise or fully dense interactions among slots. Finally, comprehensive experiments demonstrate that RSM outperforms the state-of-the-art in various next-step prediction tasks, including independent and identically distributed (iid) and Out-of-Distribution (OOD) scenarios.
In summary, the presented work makes the following contributions:
1. RSM: A modular dynamics model comprising a set of reusable mechanisms that take as input slot representations through an attention bottleneck and sequential slot updates.
2. RSM strengthens communication among slots in dynamics modeling by relaxing inductive biases in dynamics modeling to be not too dense or sparse but depend on a specific context.
3. RSM outperforms baselines in multiple tasks, including next frames prediction, Visual Question Answering, and action planning in both iid and OOD cases.
2 Proposed Method: RSM - Reusable Slotwise Mechanisms
2.1 RSM Overview
We introduce RSM, a modular architecture consisting of a set of  $M$  Multilayer Perceptrons (MLPs) that act as reusable mechanisms, operating on slotwise representations to predict changes in the slots.
(a) Spatial-temporal based Future Prediction Pipeline
(b) RSM Intuition
(c) Computational Flow in RSM from step  $t$  to  $t + 1$
Figure 1: The future prediction pipeline (Figure 1(a)), RSM Intuition (Figure 1(b)), and Computational Flow (Figure 1(c)). Colored circles represent slots, with the dashed border denoting changes in a slot. The Central Contextual Information (CCI) is derived from all object slots as context, assists in selecting a mechanism for slots, and acts as an input of mechanisms. Slots are sequentially updated in four steps: (1) Compute CCI by unrolling slots (updated and non-updated) over the past  $\tau$  steps, (2) Select a mechanism based on CCI and the slot of interest, (3) Predict the next state by the selected mechanism's dynamics, and (4) Update the predicted slot and prepare for the next object's turn.
What sets RSM apart from other architectures is incorporating the CCI buffer, which enhances its adaptability when dealing with environments characterized by varying levels of interaction complexity among objects by allowing the propagation of information about all other slots. Unlike previous approaches [18, 16, 17, 23], we enable a sparse subset of slots to transmit contextual information through a bottleneck for updating each slot. This inductive bias becomes helpful in environments where higher-order complex interactions need to be captured by reusable mechanisms as the CCI effectively modulates the complexity of the mechanisms and accommodates those that require one or two slots, as well as those that rely on a larger subset of slots.
The training pipeline of RSM, which is visualized in Figure 1(a), is designed to predict  $K$  rollout steps of  $N$  object slots, based on  $T$  burn-in frames with a temporal window of maximum  $\tau$  history steps for each prediction step. Prediction of the rollout begins by processing the given  $T$  burn-in frames to obtain the  $N$  slot representations for each of the  $T$  burn-in steps using an object-centric model. For any rollout step after the burn-in frames, the slots in a window of  $\tau \leq T$  previous steps will be fed to the model as additional context to predict the state (slots) at  $t + 1$ . The model takes this input and updates the slots sequentially to output the slots at  $t + 1$ . The sequential updating of slots means that updates from (according to some ordering) slots can influence the prediction of later slots within a prediction step, as illustrated in Figure 1(b). It is worth highlighting that the sequential way of updating slots breaks the symmetries in mechanism selection, and also allows for a more expressive transition function, similar to how autoregressive models enable encoding of rich distributions. The prediction process is repeated until the slots for  $K$  rollout steps are predicted.
2.2 Computational Flow in RSM
This section describes the computational flow of RSM in more detail along with a 4-step process that will be repeated for all slots within a time-step  $t$ , in a sequential manner, as illustrated in Figure 1(c) and Algorithm 1 in the Appendix. The following are the main components of the architecture, where  $d_{s}$  and  $d_{cci}$  denote the dimension of slots and the CCI, respectively:
1. MultiheadAttention( $\cdot$ ):  $\mathbb{R}^{((\tau + 1) \times N) \times d_s} \to \mathbb{R}^{d_s}$  followed by a projection  $\phi(\cdot): \mathbb{R}^{d_s} \to \mathbb{R}^{d_{cci}}$  that computes the CCI, denoted as  $\mathsf{cci} \in \mathbb{R}^{d_{cci}}$ , from all of the  $N$  slots in the past  $T$  steps concatenated. Keys, queries, and values all come from slots, so the CCI is not affected by the order of slot representations.
2. The set of  $M$  reusable mechanisms  $\{g_1, \ldots, g_M\}$  where  $g_i(\cdot): \mathbb{R}^{d_{ci} + d_s} \to \mathbb{R}^{d_s}$  are represented by independently parametrized MLPs implicitly trained to specialize in explaining different transitions. Each such  $g_i(\cdot)$  takes as input one slot concatenated with the CCI.
3.  $\psi (\cdot):\mathbb{R}^{d_{\mathrm{cci}} + d_s}\to \mathbb{R}^M$  that takes as input the CCI and the slot of interest  $s_t^i$ . It computes a categorical distribution over the possible choices of mechanisms for  $s_t^i$ , and outputs a sample of that distribution to be used for updating  $s_t^i$ .
Considering the  $N$  slots per each of the  $\tau$  steps in the temporal window before  $t$ ,  $s_{\tau^{*}:t}^{1:N} = \{s_{t - \tau +1}^{1},s_{t - \tau +1}^{2},\ldots ,s_{t - \tau +1}^{N},\ldots ,s_{t}^{1},s_{t}^{2},\ldots ,s_{t}^{N}\}$  with  $\tau^{*} = t - \tau +1$ , RSM predicts the next state of slots, denoted as  $s_{t + 1}^{\prime 1:N}$ , using the following 4-step process, which is sequentially applied to each of the slots. Suppose an ordering has been fixed over the slots for a rollout, and according to that ordering, for some  $0 < n\leq N$ , we have that  $n - 1$  slots have been updated to their predicted values at  $t + 1$  and are denoted by  $s_{t + 1}^{\prime 1},\ldots ,s_{t + 1}^{\prime n - 1}$ . Below, we explain the process of computing the next state for  $s_t^n$  (e.g. the blue slot in Figure 1(c)).
Step 1. Compute the CCI: We first append  $s_t^{t'1:N}$  to the current temporal window to achieve  $s_{\tau^{*}:t+1}^{1:N}$ . The duplicated  $s_t^{t'1:N}$  serves as a placeholder, which will be overwritten with the predicted values in subsequent steps. Subsequently, as presented in Equation 1, a MultiheadAttention( $\cdot$ ) takes  $s_{\tau^{*}:t+1}^{1:N}$  as input before passing through  $\phi(\cdot)$  to produce the central contextual information cci.
$$
\operatorname {c c i} = \phi (\text {M u l t i h e a d A t t e n t i o n} \left(s _ {\tau^ {*}: t + 1} ^ {1: N}\right) \tag {1}
$$
Step 2. Select one mechanism for  $s_t^n$ :  $\psi(\cdot)$  takes two arguments as inputs, cci from Step 1 and  $s_t^n$ , to output the logits of a categorical distribution  $\pi_{1:M}$  over  $M$  possible choices of mechanisms. We employ a hard Gumbel-softmax layer with Straight-through Trick [33, 22] on top of  $\psi$ 's outputs, as described in Equation 2, to select the mechanism.
$$
\pi_ {1: M} = \text {G u m b e l - s o f t m a x} \left(\psi \left(\operatorname {c c i}, s _ {t} ^ {n}\right)\right) \tag {2}
$$
Step 3. Predict the changes of  $s_t^n$ : Let  $\Delta s_t^n$  denote the change of  $s_t^n$  from  $t$  to  $t + 1$ .  $\Delta s_t^n$  is presented in Equation 3, where  $k = \operatorname{argmax}(\pi_{1:M})$ .
$$
\Delta s _ {t} ^ {n} = g _ {k} \left(\operatorname {c c i}, s _ {t} ^ {n}\right) \tag {3}
$$
Step 4. Update and sync  $s_{t+1}^{n}$ :  $s_{t+1}^{\prime n}$  is then computed by adding the predicted transformation from the previous step and replaces the value of  $s_{t+1}^{n}$  in the slots buffer, as described in Equation 4.
$$
s _ {t + 1} ^ {\prime n} = s _ {t} ^ {n} + \Delta s _ {t} ^ {n} \tag {4}
$$
The process above is repeated for all slots at time  $t$ , to obtain the next state (slots) prediction  $s_{t+1}^{1:N}$ .
3 Experiments Setup
This study evaluates RSM's dynamics modeling and generalization capabilities through video prediction, VQA, and action planning tasks. We aim to provide empirical evidence supporting the underlying hypotheses that guided the architectural design of RSM.
-  $\mathcal{H}_1$ : Slots communication through the CCI and reusable mechanisms reduces information loss during prediction, resulting in accurate prediction of future rollout frames (Section 4.1).
-  $\mathcal{H}_2$ : RSM effectively handles novel scenarios in the downstream tasks (Section 4.2), especially enhancing OOD generalization (Section 4.3).
-  $\mathcal{H}_3$ : The synergy between the CCI and the disentanglement of objects dynamics into reusable mechanisms is essential to RSM (analyses in Section 4.4 and ablations in Section 4.5).
In the following subsections, we describe the experiments focusing on the transition of slots over rollout steps with pre-trained object-centric models. Additionally, Appendix E provides experiments and analyses with an end-to-end training pipeline.
3.1 Environments
OBJ3D [29] contains dynamic scenes of a sphere colliding with static objects. Following Lin et al. [29], Wu et al. [44], we use 3 to 5 static objects and one launched sphere for interaction.
CLEVRER [45] shares similarities with OBJ3D, but additionally has multiple moving objects in various directions throughout the scene. For the VQA downstream task, CLEVRER offers four question types: descriptive, explanatory, predictive, and counterfactual, among which, in the spirit of improving video prediction, we focus on boosting the performance on answering predictive questions which require an understanding of future object interactions.
PHYRE [4] is a 2D physics puzzle platform where the goal is strategically placing red objects such that the green object touches the blue or purple object. Bakhtin et al. [4] design templates that describe such tasks with varying initial states. Subsequently, they define (1) within-template protocol where the test set contains the same templates as in training, and (2) cross-template protocol that tests on different templates than those in training. We report results both on within-template as iid and on cross-template to obtain the OOD evaluation.
Physion [7] is a VQA dataset that assesses a model's capability in predicting objects' movement and interaction in realistic simulated 3D environments in eight physical phenomena.
For further details and data visualization, we refer the readers to Appendix B.
We compare RSM against three main baselines: SlotFormer [44], Switch Transformer [14] denoted as SwitchFormer, and NPS [17]. We show the efficacy of relaxing the inductive bias on communication density among slots in RSM by contrasting with dense communication methods (SlotFormer and SwitchFormer) and a pair-wise communication method (NPS). Likewise, comparing RSM to SlotFormer highlights the role of disentangling objects' dynamics into mechanisms, while comparing to SwitchFormer and NPS emphasizes the vital role of communication among slots via the CCI. Additionally, we compare to SAVi-Dyn [44], which is the SOTA on CLEVRER. In other experiments, we compare to SlotFormer (current SOTA).
In the tables, we present our reproduced SlotFormer (marked by " $\dagger$ ) and our re-implemented SwitchFormer and NPS (marked by " $\ast$ "), alongside SAVi-Dyn reported by Wu et al. [44].<sup>2</sup>
Following Wu et al. [44], we focus on the transition of slots and take advantage of the pre-trained object-centric encoder-decoder pair that convert input frames into slots and vice versa. We use the pre-trained weights of SAVi and STEVE provided by Wu et al. [44], including SAVi [26] for OBJ3D, CLEVRER, and PHYRE; and STEVE [40] for Physion.
We present the best validation set configuration of RSM for each dataset, along with fine-tuning results and model size scaling in Appendix D. In summary, (1) OBJ3D and CLEVRER include 7 mechanisms, while PHYRE and Physion use 5, and (2) the number of parameters in RSM is slightly lower than that of SlotFormer in corresponding experiments. Additionally, we re-implemented SwitchFormer and NPS with a similar number of parameters as in RSM and SlotFormer.
2We adapt the code for Switch Transformer and NPS to ensure consistency of experimental setups and evaluations with SlotFormer and RSM (See Appendix C).
Table 1: Future frame prediction quality on OBJ3D and CLEVRER. Bold scores indicate the best performance, with the RSM consistently outperforming baselines by a remarkable margin.
Method	OBJ3D	CLEVRER
SSIM↑	LPIPS×100↓	SSIM↑	LPIPS×100↓	ARI↑	FG-ARI↑	FG-mIoU↑
SAVi-Dyn	0.91	12.00	0.89	19.00	8.64	64.32	18.25
NPS*	0.90±0.2	8.24±0.2	0.89±0.2	12.51±0.0	62.84±0.2	64.62±0.3	30.39±0.2
SwitchFormer*	0.91±0.2	8.09±0.3	0.88±0.3	14.28±0.1	60.61±0.4	59.32±0.3	28.94±0.2
SlotFormer†	0.90±0.2	8.32±0.2	0.88±0.2	13.09±0.1	63.38±0.3	62.91±0.2	29.68±0.3
RSM (Ours)	0.92±0.1	7.88±0.1	0.91±0.1	11.96±0.1	67.72±0.2	66.15±0.2	32.73±0.2
Figure 2: Comparison of rollout frames in OBJ3D. RSM renders frames with precise dynamics and holds visual quality, even complex actions in steps 20 to 30. In contrast, the baselines produce ones with artifacts (yellow boxes) and inaccurate dynamics (red boxes). Best viewed in video format.
We report mean and standard deviation across 5 different runs. Video visualizations of our experiments are provided in our repository  $^{3}$ . See also Appendix A on the reproducibility of our results.
4.1 Video Prediction Quality
To demonstrate  $\mathcal{H}_1$ , we provide the video prediction quality on OBJ3D and CLEVRER in Table 1 and Figure 2. In general, RSM demonstrates its effectiveness in handling object dynamics in the long-term future prediction over baselines.
Evaluation Metrics We evaluate the quality and similarity of the predicted rollout frames using SSIM [42] and LPIPS [46] metrics. Since the range of LPIPS metric is small, we report the actual values times 100, denoted as  $\mathbf{LPIPS}_{\times 100}$ , to facilitate comparisons among the methods. Additionally, we also assess the performance using ARI, FG-ARI, and FG-mIoU metrics, which measure clustering similarity and foreground segmentation accuracy of object bounding boxes and segmentation masks. We evaluate the model's performance averaged over unrolling 44 steps on OBJ3D and 42 steps on CLEVRER with six burn-in frames in both datasets.
Table 1 exhibits that RSM outperforms other approaches and achieves the highest scores across evaluation metrics for both datasets. Notably, compared to SlotFormer, RSM improves  $\mathrm{LPIPS}_{\times 100}$  by 0.46 points in OBJ3D, 1.12 points in CLEVRER, and increases FG-mIoU by 3.05 points in CLEVRER. Following RSM, NPS consistently ranks second in performance among the baselines.
These results are supported by Figure 2, which illustrates the rollout frames in OBJ3D. RSM's outputs' accurate rollout predictions with high visual fidelity demonstrate the efficacy of slot communication by having less error accumulated along time steps than any of the baselines. It is worth emphasizing that RSM excels in handling a significant series of complex movements in steps 20-30, particularly
$^{3}$ github.com/trangnnp/RSM
Table 2: VQA performance on CLEVRER and Physion. Despite not surpassing human performance in Physion, RSM outperforms baselines in both datasets. All scores are in percentage.
Method	CLEVRER	Physion
per opt. ↑	per ques. ↑	Obs. ↑	Dyn. ↑	Gap ↑
Human	-	-	74.7	-	-
NPS*	95.3±0.3	93.8±0.2	65.6±0.3	+0.6
SwitchFormer*	92.8±0.3	90.4±0.2	65.0	66.2±0.1	+1.2
SlotFormert	96.1±0.2	93.3±0.1	66.9±0.2	+1.9
RSM (Ours)	96.8±0.1	94.3±0.0	68.1±0.0	+3.1
Table 3: Action planning task in PHYRE. RSM outperforms all baselines in both iid and OOD.
PHYRE-B	NPS*	SwitchFormer*	SlotFormer†	RSM
iid (within-template)	80.52±1.0	78.27±1.9	76.4±1.1	82.89±0.6
OOD (cross-template)	42.63±1.3	48.36±1.4	42.46±1.7	57.37±1.4
during a sequence of rapid collision of object pairs. In contrast, we find that the baselines struggle with complex object movements during this period, leading to inaccuracies in predicting the dynamics towards the end. RSM effectively maintains visual quality by predicting the changes of slots instead of predicting the next state of a slot. This approach allows RSM to handle action-free scenarios well and significantly reduce error accumulation by facilitating null transitions that preserve slot integrity.
Furthermore, RSM demonstrates flexible slot communication with relaxed inductive biases on interaction density, enabling it to adapt to environments comprising mechanisms with varying levels of complexity. In contrast, we find that NPS with sparse interactions faces difficulties with close-by objects and rapid collisions (seen in OBJ3D), while SwitchFormer with dense communication struggles with distant objects (as in CLEVRER).
4.2 Downstream Tasks: Visual Question Answering and Action Planning
4.2.1 Visual Question Answering task
To demonstrate  $\mathcal{H}_2$ , we validate the performance of future frames generated by the models on the VQA task in CLEVRER and Physion, and the action planning task in PHYRE in the next section. The general pipeline is to solve the VQA task with the predicted rollout frames from the given input frames. Specifically, we employ Aloe [11] as the base reasoning model on top of the unrolled frames in CLEVRER. Likewise, in Physion, we adhere to the official protocol by training a linear model on the predicted future slots to detect objects' contact. In Physion, we also include the results obtained from human participants [7] for reference. Likewise, we collect the results from observed frames (Obs.), which are obtained by training the VQA model on top of pre-trained burn-in slots and compare them to the performance of rollout slots (Dyn.).
In Table 2, RSM consistently outperforms all three baselines in VQA for CLEVRER and Physion. On CLEVRER, RSM achieves the highest scores of  $96.8\%$  (per option) and  $94.3\%$  (per question), surpassing SlotFormer and NPS. In Physion, RSM shows notable improvement, with a  $3.1\%$  increase from  $65.0\%$  in Obs. to  $68.1\%$  in Dyn., outperforming all other baselines, indicating the benefit of enhancing the dynamics modeling to improve the downstream tasks. However, RSM is still far from human performance in Physion, showing room for further research into this class of algorithms.
4.2.2 Action Planning task