SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We adopt the approach from prior works [36, 44] and construct a task success classifier as an action scoring function. This function is designed as a simple linear classifier, which considers the predicted object slots, to rank a pre-defined set of 10,000 actions from [4], executed accordingly. We utilize AUCCESS, which quantifies the success rate over the number of attempts curve's Area Under Curve
Figure 3: Action planning task in PHYRE. RSM strategically positions a red ball at step 0 to prevent the green ball from falling onto the glass by causing a collision that alters the original trajectory of the green ball and causes it to make contact with the blue floor (indicated by the arrow).
Figure 4: The underlying mechanism assignment in PHYRE. Mechanisms are assigned to each slot at  $t$  to obtain the updates at  $t + 1$ . RSM disentangles objects' dynamics into reusable mechanisms, which can be expressed as Collision (2), Moving left or right (3), Idle (4), and Falling (5).
(AUC) for the first 100 actions. We report the average score over ten folds, with the best performance among five different runs on each fold.
The first line in Table 3 indicates the action planning results of the proposed RSM and baselines in the iid setting (within-template protocol). RSM achieves the highest performance compared to baselines, indicating the critical role of efficient slots communication in complex tasks like action planning. In addition, Figure 3 shows a successful case of RSM solving the planning task by strategically placing the red object at step 0, causing a collision between the green and blue objects at the end.
4.3 Out-of-Distribution (OD) Generalization
To provide more evidence for  $\mathcal{H}_2$ , we resume analyzing the performance of the action planning task in PHYRE in two OOD scenarios:
OOD in Templates Performance on cross-template protocol is indicated in the last line in Table 3. Overall, RSM demonstrates strong generalization and has the smallest gap between iid and OOD performance compared to the baselines. The cross-template is a natural method to evaluate the OOD generalization in PHYRE since scenes in the train and test sets are in distinct templates and contain dissimilar object sizes and objects in the background [4]. We refer the reader to Appendix B for further details and visualizations.
OD in Dynamics We introduce modifications to PHYRE dataset generation that introduce distribution shifts to the dynamics of the blue objects. Particularly, the blue objects are static floors during training. However, in the test set, the object can be assigned as a blue ball, inheriting all dynamics as a red or green ball. As a result, the blue ball appearing only in the test set is considered a strong OOD case, as the model does not have a chance to construct the blue ball and its dynamics during training. Presented in Figure 5, although revealing some distortion in shape, RSM demonstrates the ability to generalize the movements of the ball-shaped object in intensive OOD cases.
4.4 The Ability to Disentangle Object Dynamics into Mechanisms
To demonstrate  $\mathcal{H}_3$ , we conduct qualitative analysis on the underlying mechanisms assignment within the four steps of Figure 3 and visualize results in Figure 4. While there is no explicit assignment
Figure 5: Ablation on OOD objects' dynamics. The blue objects are static during training (blue floor) and can move during testing time (blue balls). RSM responds correctly to the blue ball's dynamics. Pred* are predictions by a model trained with the standard training set, Pred are predictions by a model trained with the modified training data, and Mechanism are mechanisms assigned for the blue object's slot when obtaining Pred. h.Move denotes the horizontal movement.
Figure 6: Ablation studies in OBJ3D. The original design of RSM always demonstrates the superior performance and the accurately predicted future frames compared to the modified versions, including breaking the usages of the CCI in step 2 ( $\mathbf{RSM}_{!2}$ ) and step 3 ( $\mathbf{RSM}_{!3}$ ), randomly selecting mechanisms ( $\mathbf{RSM}_{\bar{\mathbf{k}}}$ ), and parallel update of slots ( $\mathbf{RSM}_{\mathbf{p}}$ ). Best viewed in video format (See our repository).
of roles to mechanisms during training, we can infer their functionality at convergence based on slot visualizations and patterns of activation as follows: Mechanism 2 handles collisions, which can be observed in the collision between the green and red balls in step 1 and that of the red ball with the right-side wall in step 5 (blue boxes). Mechanism 3 controls the horizontal movements, observed in the green ball from steps 5 to 10 (green boxes). Mechanism 4 acts as an idleness mechanism. Mechanism 5 handles downward free-fall motion, observed from steps 2 to 3 of the two balls. Mechanism 1 does not seem to be doing anything meaningful, and this could be due to the model using more capacity than it requires to model the dynamics in this environment.
The inferred functions provide the following insights into understanding the efficacy of RSM. Firstly, RSM successfully disentangles the dynamics into reusable mechanisms, as described in Section 2.1. Secondly, RSM properly assigns such mechanisms to slots throughout the rollout steps, which not only helps preserve the accuracy of prediction but also emphasizes the effectiveness of communication among slots in deciding mechanisms for one another. The automatic emergence of a null mechanism (mechanism 4) is also worth highlighting, which significantly helps reduce the error accumulation in action-free settings, such as idle objects in the background.
To provide more evidence for  $\mathcal{H}_3$ , we conduct ablations to understand the individual effects of (1) the CCI, (2) mechanisms and their disentanglement, and (3) the sequential update of slots in RSM, and visualize the results in Figure 6. Overall, the ablations confirm the superiority of the original RSM design compared to all variations, highlighting a significant disparity in the absence of CCI.
$\mathbf{RSM}_{!2}$  masks out the CCI in step 2 at inference. We observe that the lack of CCI leads to a degenerate selection of mechanisms for slots, with 4 out of 5 object slots being controlled by mechanism 3.
$\mathbf{RSM}_{!3}$  masks out the CCI in step 3 during inference. We have found that CCI not only captures comprehensive spatial-temporal information but also provides guidance regarding specific movement
details, including directions. In particular, even when the correct mechanism is assigned (e.g., moving backwards), the slots become confused about the exact direction of movement. Additionally, slots encountered a color-related issue in the subsequent steps.
$\mathbf{RSM}_{\bar{\mathbf{k}}}$  randomly assigns mechanisms to slots by a randomized mechanism index,  $k$ , to replace the distribution  $\pi$  in Equation 2. We observe that the launched ball moves in the wrong direction and exits the scene early, while other objects shake in their positions, underscoring the importance of correctly assigning mechanisms to slots.
$\mathbf{RSM}_{\mathbf{p}}$  is the parallel version of RSM that modifies the model  $\psi(\cdot)$  in Equation 2 to assign mechanisms to  $N$  slots simultaneously, in both training and inference time. We find that  $\mathbf{RSM}_{\mathbf{p}}$  stands out as a promising model, as it demonstrates a notable LPIPS performance; however, it is essential to note the partially inaccurate dynamics caused by the weaker communication among slots.
Modular dynamics models. In the domain of modular neural networks, RIMs [18] pioneered the exploration of modularity for learning dynamics and long-term dependencies in non-stationary settings. However, RIMs suffer from conflating object and mechanism concepts, limiting their effectiveness. SCOFF [16] introduced object files (OF) and schemata to address this limitation, but it struggles with generalizing out-of-distribution (OOD) scenarios. Key distinctions between RSM and SCOFF are as follows: SCOFF schemas can handle only one OF at a time, while RSM allows multiple slots to be input to reusable mechanisms using an attention bottleneck. SCOFF and RIMs use sparse communication among OFs for rare events involving multiple objects, whereas RSM leverages the CCI to activate suitable reusable mechanisms. NPS [17] incorporates sparse interactions directly into its modular architecture, eliminating the need for sparse communication among slots or OFs. Their "production rules" handle rare events involving multiple objects by taking one primary slot and a contextual slot as input. A recent benchmark [23] evaluates causal discovery models and introduces a modular architecture with dense object interactions, similar to GNN-based methods, but assigns separate mechanisms to each object. Among the class of algorithms with less modularity in their dynamics model, R-NEM [41] was a pioneering unsupervised method for modeling dynamics using a learned object-centric latent space through iterative inference (see also [13]) which along similar approaches ([6], [5]) used Graph Neural Networks (GNNs) to model pairwise interactions and differ from our work in two key aspects. Firstly, we do not rely on GNNs to model interactions, as dense interactions are not always realistic in many environments. Secondly, we focus on learning a set of simple and reusable mechanisms that can be applied flexibly to different scenarios, rather than compressing information through shared node and edge updates in a GNN.
Unsupervised learning of object-centric representations. To decompose a scene into meaningful sub-parts, there have been lots of recent works on unsupervised learning of object-centric representations from static images [48, 20, 31, 47], videos [35, 27, 25, 12, 44], and theoretical results on the identifiability of such representations [34, 1, 9]. Although lots of these methods work very well in practice, we decided to proceed with slot attention [31] to be consistent with the baselines.
In this study, we developed RSM, a novel framework that leverages an efficient communication protocol among slots to model object dynamics. RSM comprises a set of reusable mechanisms that take as input slot representations passed through a bottleneck, the Central Contextual Information (CCI), which are then processed sequentially to obtain slot updates. Through comprehensive empirical evaluations and analysis, we show RSM's advantages over the baselines in various tasks, including video prediction, visual question answering, and action planning tasks, especially in OOD settings. Our results suggest the importance of CCI, which integrates and coordinates knowledge from different slots for both mechanism assignment and predicting slot updates. We believe there is a promise for future research endeavors in exploring more sophisticated stochastic attention mechanisms for information integration, aligning with the principles of higher-level cognition, coping with a large number of object slots, and enabling uncertainty quantification in the predictions. In Appendix F, we discuss the limitations of this work and future directions.
We gratefully acknowledge the support received for this research, including from Samsung, CIFAR and NSERC. This research was partly enabled by computational resources provided by Mila - Quebec AI Institute and The Digital Research Alliance of Canada (formerly known as Compute Canada). Each member involved in this research is funded by their primary institution. Amin Mansouri acknowledges support from Microsoft. Additionally, we sincerely thank the anonymous reviewers for their valuable comments and contributions to improve this work.
[1] Kartik Ahuja, Jason Hartford, and Yoshua Bengio. Weakly supervised representation learning with sparse perturbations. arXiv preprint arXiv: Arxiv-2206.01101, 2022.
[2] Bernard J. Baars. Global workspace theory of consciousness: toward a cognitive neuroscience of human experience. In Steven Laureys, editor, The Boundaries of Consciousness: Neurobiology and Neuropathology, volume 150 of Progress in Brain Research, pages 45-53. Elsevier, 2005. doi: https://doi.org/10.1016/S0079-6123(05)50004-9. URL https://www.sciencedirect.com/science/article/pii/S0079612305500049.
[3] Bernard J. Baars. The Global Workspace Theory of Consciousness. John Wiley and Sons, Ltd, 2017. ISBN 9781119132363.
[4] Anton Bakhtin, Laurens van der Maaten, Justin Johnson, Laura Gustafson, and Ross Girshick. Phyre: A new benchmark for physical reasoning. 2019.
[5] Peter W. Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, and Koray Kavukcuoglu. Interaction networks for learning about objects, relations and physics. In NIPS, 2016.
[6] Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Flores Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulçehre, H. Francis Song, Andrew J. Ballard, Justin Gilmer, George E. Dahl, Ashish Vaswani, Kelsey R. Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matthew M. Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. CoRR, abs/1806.01261, 2018. URL http://arxiv.org/abs/1806.01261.
[7] Daniel M. Bear, Elias Wang, Damian Mrowca, Felix J. Binder, Hsiau-Yu Fish Tung, R. T. Pramod, Cameron Holdaway, Sirui Tao, Kevin A. Smith, Fan-Yun Sun, Li Fei-Fei, Nancy Kanwisher, Joshua B. Tenenbaum, Daniel L. K. Yamins, and Judith E. Fan. Physion: Evaluating physical prediction from vision in humans and machines. CoRR, abs/2106.08261, 2021. URL https://arxiv.org/abs/2106.08261.
[8] Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. In ICLR'2020, arXiv:1901.10912, 2019.
[9] Jack Brady, Roland S. Zimmermann, Yash Sharma, Bernhard Schölkopf, Julius von Kügelgen, and Wieland Brendel. Provably learning object-centric representations. arXiv preprint arXiv: 2305.14229, 2023.
[10] Christopher P. Burgess, Loic Matthews, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matthew M. Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representation. CoRR, 2019.
[11] David Ding, Felix Hill, Adam Santoro, Malcolm Reynolds, and Matthew Botvinick. Attention over learned object embeddings enables complex visual reasoning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https://openreview.net/forum?id=1HmhW2zmVN.
[46] R. Zhang, P. Isola, A. A. Efros, E. Shechtman, and O. Wang. The unreasonable effectiveness of deep features as a perceptual metric. In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 586-595, Los Alamitos, CA, USA, jun 2018. IEEE Computer Society. doi: 10.1109/CVPR.2018.00068. URL https://doi.ieeecomputersociety.org/10.1109/CVPR.2018.00068.
[47] Yan Zhang, David W. Zhang, S. Lacoste-Julien, G. Burghouts, and Cees G. M. Snoek. Unlocking slot attention by changing optimal transport costs. ARXIV.ORG, 2023. doi: 10.48550/arXiv.2301.13197.
[48] Mingde Zhao, Zhen Liu, Sitao Luan, Shuyuan Zhang, Doina Precup, and Yoshua Bengio. A consciousness-inspired planning agent for model-based reinforcement learning. Advances in neural information processing systems, 34:1569-1581, 2021.
Each experiment is trained on 4 V100-GPUs with 12 CPUs, using a distributed data-parallel training strategy. The number of parameters and the average training time over 5 runs are summarized in Table 4. In addition, Table 5 provides the necessary configurations to reproduce our work, including the setting related to datasets and the training process that follows the prior work [44], and RSM's design that achieves the best tuning results.
Regarding the rollout frames  $K$  and the video length  $V$  in Table 5, we predict and consider  $K$  future frames from the last burn-in steps for training, whereas, we produce the total of  $V$  frames in the inference time, including  $T$  burn-in and  $V - T$  rollout steps. In other words,  $V$  equals  $T$  plus the actual rollout frames at inference time. Regarding the training process, we employ the Adam optimizer with an initial learning rate of  $2 \times 10^{-4}$  and employ a cosine decay schedule to 0. We further discuss the RSM design in Appendix C.
Table 4: Summary of the number of parameters and training duration. "M" stands for millions. "h" stands for GPU hours.
RSM	NPS	SwitchFormer	SlotFormer
OBJ3D
Num. Params	0.76M	0.99M	0.82M	0.82M
Training Duration	21h	25h	20h	21h
CLEVRER
Num. Params	3.1M	4.06M	3.22M	3.22M
Training Duration	82h	94h	72h	86h
PHYRE
Num. Params	5.13M	5.98M	6.38M	6.38M
Training Duration	29h	33h	28h	30h
Physion
Num. Params	5.61M	6.7M	6.41M	6.41M
Training Duration	32h	41h	30h	34h
Table 5: Summary of experiments' configuration, including the configuration of datasets, training process, and RSM.
OBJ3D	CLEVREER	Physion	PHYRE
Frame Size	64 × 64	64 × 64	128 × 128	128 × 128
Num. Slots N	6	7	6	8
Slot Size ds	128	128	192	128
Burn-in Frames T	6	15	15	1
Temporal Window τ	6	15	15	6
Rollout Steps K	10	10	10	10
Video Length V	6+44	15+42	15+35	1+14
Batch Size	128	128	128	64
Num. Epochs	200	80	25	50
Object-centric Model	SAVi	SAVi	STEVE	SAVi
Loss Weight λ	1.0	1.0	0.0	0.0
Num. Mechanisms M	7	7	5	5
Num. Layers of ψ(·)	1	2	2	2
Num. Layers of Mechanism	1	3	3	3
B Dataset Collection
We follow the collection and pre-processing of the dataset, including video length and the dataset splits, as done by Wu et al. [44]. In addition, we provide datasets visualizations in Figure 7 and the visualization of templates in PHYRE in Figure 8.
Algorithm 1 Reusable Slotwise Mechanisms
$N$  : number of slots  $M$  : number of mechani  $\tau$  : number of history let  $T$  : number of burn-in ft  $K$  : number of rollout st
$d_{s}$  : slot dimension  $d_{ccc}$  : central contextual information dimension
Variables in RSM
Input:  $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}\} \in \mathbb{R}^{(\tau \times N) \times d_{s}}$  with  $\tau^{*} = t - \tau + 1$ : unrolled  $N$  slots of the previous  $\tau$  steps to time  $t$   
Output:  $s_{T + 1:T + K}^{1:N}$ : predicted  $N$  slots in the next  $K$  steps from time  $T + 1$  to  $T + K$
- cci  $\in \mathbb{R}^{d_{cci}}$ : the Central Contextual Information (CCI).
-  $s_t^n \in \mathbb{R}^{d_s}$ : the slot of interest, which is slot  $n^{th}$  at time  $t$ .
-  $p \in \mathbb{R}^M$ : the Gumbel distribution over  $M$  choices of selecting mechanisms.
-  $\Delta s_t^n \in \mathbb{R}^{d_s}$ : changes of the  $n$ -th slot from time  $t$  to  $t + 1$ .
-  $\mathbf{W}_{\mathbf{q}}, \mathbf{W}_{\mathbf{k}}, \mathbf{W}_{\mathbf{v}}: \mathbb{R}^{d_s} \to \mathbb{R}^{d_s}$  denote query, key, and value projection layers transforming the unrolled  $s_{\tau^{*},t}^{1:N}$  in the attention mechanism.
- MultiheadAttention  $(\cdot): \mathbb{R}^{((\tau + 1) \times N) \times d_s} \to \mathbb{R}^{d_s}$ : apply self-attention on the  $s_{\tau^*,t}^{1:N}$
-  $\phi(\cdot): \mathbb{R}^{d_s} \to \mathbb{R}^{d_{\mathrm{cci}}}$  computes the central contextual information by passing the outputs of the MultiheadAttention(\cdot) through a nonlinear transformation (MLP).
-  $\psi (\cdot):\mathbb{R}^{d_{\mathrm{cci}} + d_s}\to \mathbb{R}^M$  computes the unnormalized probability of selecting a mechanism from  $M$  possible choices by taking the CCI and slot of interest as input and feeding that to an MLP.
- Set of  $M$  mechanisms  $g_{j}(\cdot):\mathbb{R}^{d_{cci} + d_{s}}\to \mathbb{R}^{d_{s}},j\in \{1\dots M\}$  : predict the changes of each slot based on the CCI and current state of the slot. These are also realized with MLPs.
Components in RSM:
for each  $t$  in  $[T\ldots t + K)$  do
Step 0: Prepare slots buffer  $s_{\tau^{*}:t}^{1:N} = \mathrm{concat}(s_{\tau^{*}:t}^{1:N}, s_t^{1:N}) \in \mathbb{R}^{((\tau + 1) \times N) \times d_s}$
for each  $s_t^n$  in  $s_t^{1:N}$  with  $n \in 1 \dots N$  do  
Step 1: Compute the central context  
 $\mathsf{cci} = \phi(\mathbf{MultiheadAttention}(\mathbf{W_q}(s_{\tau^*:t+1}^{1:N}), \mathbf{W_k}(s_{\tau^*:t+1}^{1:N}), \mathbf{W_v}(s_{\tau^*:t+1}^{1:N})))$
Step 2: Select a mechanism for slot  $s_t^n$ $p = \mathrm{Gumbel - max}(\psi (\mathrm{concat}(\mathrm{cci},s_t^n))$
Step 3: Apply the selected mechanism to slot  $s_t^n$ . Note that  $p$  is one-hot-like distribution.
$\Delta s_{t}^{n,j} = g_{j}(\mathrm{concat}(\mathsf{cci},s_{t}^{n}))*p^{j}$ $\forall j\in \{1,\ldots ,M\}$ $\Delta s_t^n = \sum_{j = 1}^M\Delta s_t^{n,j}$
Step 4: Update the slots buffer with the new value of  $s_{t+1}^n$ $s_{t+1}^n = s_t^n + \Delta s_t^n$
end for end for return  $s_{T + 1:T + K}^{1:N}$
Figure 7: Dataset visualization over steps.
Figure 8: Pairs of PHYRE scenes in the same template with similar objects in the background and differences in objects' positions.
OBJ3D We collect the OBJ3D dataset from the official GitHub repository<sup>4</sup>.
CLEVRER We directly download the CLEVRER dataset from the official website<sup>5</sup>.
PHYRE We use the PHYRE-1B version, which sets the number of red balls to 1. The PHYRE dataset is generated based on the instructions provided on the official GitHub page<sup>6</sup>.
Physion We directly download the Physion dataset from the official GitHub page<sup>7</sup>
C.1 Loss Functions
The following is the training objective that follows the prior work [44].
We employ slot reconstruction loss to train the rolled out future frames prediction, as described in Equation 5 where  $n$  is the slot index,  $s_{T + k}^{n}$  is the predicted rollout slot, and  $s_{T + k}^{*n}$  is the pre-trained slot (that is used as the target slot).
$$
\mathcal {L} _ {S} = \frac {1}{K \cdot N} \sum_ {k = 1} ^ {K} \sum_ {n = 1} ^ {N} \left\| s _ {T + k} ^ {n} - s _ {T + k} ^ {* n} \right\| ^ {2} \tag {5}
$$
Experiments using SAVi as the object-centric model also use the image reconstruction loss, as described in Equation 6 where  $f_{dec}$  is a frozen decoder, and  $x_{T + k}$  is the ground truth image.
$^{4}$ github.com/zhixuan-lin/G-SWM#datasets
5clevrer.csail.mit.edu/
$^{6}$ github.com/facebookresearch/phyre
$^{7}$ github.com/cogtoolslab/physics-benchmarking-neurips2021#downloading-the-physion-dataset
Experiments using STEVE as the object-centric model can still employ the image reconstruction loss; however, we do not conduct such experiments with image reconstruction loss due to the dramatically extended training time. In PHYRE, we do not utilize the image reconstruction loss  $\mathcal{L}_I$  due to the large image size that could affect the training time and the simplicity of PHYRE's objects compared to other environments in this work.
$$
\mathcal {L} _ {I} = \frac {1}{K} \sum_ {k = 1} ^ {K} \| f _ {\text {d e c}} \left(s _ {T + k}\right) - x _ {T + k} \| ^ {2} \tag {6}
$$
The overall objective function is the weighted sum of the above losses, as presented in Equation 7.
$$
\mathcal {L} = \mathcal {L} _ {S} + \lambda \mathcal {L} _ {I} \tag {7}
$$
Note that when using slot and image reconstruction losses as presented in Equation 5 and Equation 6, the model will fail in the end-to-end training since the overall objective function is constrained to pre-trained slots  $(s_{T + k}^{*n})$  and well-trained decoder  $(f_{dec}(\cdot))$ . Motivated by this observation, in Appendix E, we provide the end-to-end training objective to compare RSM and other approaches' abilities in objects' dynamics modeling from scratch.
C.2 Models Architecture
The distinction among methods is in the process of predicting the next state based on the input of past states. In this section, we provide a detailed description of how both the baselines and the proposed RSM approach the task of the next state prediction. Specifically, we examine two key aspects: (1) how slots communicate with each other and (2) how the predictions are computed, as revealed through the design specifics of each method.
C.2.1 Baselines