SlowGuess/NeurIPS_2023_OCR huggingface dataset:
SlotFormer [44] SlotFormer consists of 3 main parts: (1) the Multi-Layer Perceptron (MLP) input projection layer, (2) a Transformer architecture layer, and (3) the MLP output projection layer. First, the unrolled past states (the sequence of  $\tau \times N$  slots) are passed through the input projection layer before being processed by the Transformer model. Afterward, the output projection produces the next state from the Transformer's output. Through this process, slots densely communicate with each other in all three parts of SlotFormer. In addition, the entire next state of slots is directly generated by the models. In this work, we utilize the publicly available implementation<sup>8</sup> of SlotFormer.
SwitchFormer from Switch Transformer [14]. The Switch Transformer is a variant of the Transformer architecture designed to improve efficiency and scalability. It achieves this through the usage of dynamic routing and adaptive computation. This architecture employs a "switch" module that intelligently routes tokens to different layers based on their content.
In this work, we create SwithFormer that integrates the Switch Transformer implementation into the SlotFormer codebase and replace the vanilla Transformer by Switch Transformer. In this way, SwitchFormer follows the same strategy as in SlotFormer, which conducts the dense communication among slots and directly predicts the entire next state of slots.
NPS [17] NPS is a framework that combines neural networks and production systems for object modeling that integrates neural networks into the production system. In traditional production systems, rules are used to represent knowledge and guide the system's behavior. In the case of NPS, they conduct a set of rules to handle the pair-wise interaction of slots. The two slots involved in an interaction, which are called the primary and contextual slots, are selected through attention mechanisms. In addition, the official design of NPS for object dynamics' modeling, inspired by Kipf et al. [24], predicts the changes of the primary slot within a time step instead of the entire slots. Afterward, the predicted next state of slots is the sum of the current state and the predicted slots changes.
In this work, we integrate the official NPS<sup>10</sup> to the SlotFormer's codebase for consistency in the training pipeline and sharing the pre-trained object-centric model. Besides, NPS includes an MLP
github.com/pairlab/SlotFormer
9nn.labml.ai/transformers/swi tch/index.html
$^{10}$ github.com/anjirudh9119/neural生产和 systems
slot encoder that requires a fixed input size. However, the temporal window size in PHYRE increases from 1 to 6, leading to the difference in the input size of the starting and later steps. Therefore, at the beginning of the rollout prediction process of PHYRE, we duplicate the burn-in frame to have a fixed six steps window size along the rollout process.
C.2.2 RSM
We design RSM as a framework for dynamics modeling with a relaxed inductive bias in the communication density of slots that enables a subset of slots involved in communication, based on a particular context through the CCI. RSM consists of three main elements as described in Section 2.1: (1) the multi-head self-attention that computes the CCI, (2) the  $\psi(\cdot)$  that estimates the suitable mechanism, and (3) a list of mechanisms. In terms of the multi-head self-attention, We employ a 4-head architecture for multi-head self-attention, where the hidden size of the Feed-forward Networks is set to  $2 \times d_s$ . We design  $\psi(\cdot)$  as MLP with the number of hidden layers being tuned using the validation set (See Appendix D.3). Similarly, each individual mechanism is designed as MLP layers with a tuned number of hidden layers. All mechanisms share the same architecture but have separate weights. In addition, the total number of parameters when considering all mechanisms is constrained to be the same across when the number of mechanisms varies, meaning that as the number of allocated mechanisms increases, the number of parameters per each mechanism decreases (further investigated in Appendix D.3). Like NPS, the mechanism predicts the changes in a slot within two consecutive steps instead of the entire slot. Last but not least, we (and NPS) omit the input and output projection layers as SlotFormer and SwitchFormer.
C.2.3 Downstream tasks
CLEVRER VQA Inheriting from the baseline of Wu et al. [44], we employ Aloe as the VQA model that concatenates the predicted rollout slots and the processed question (represented as language tokens) before passing them through a stack of Aloe Transformer encoder to predict the answers.
Physion VQA In the VQA task of Physion, the objective is to determine whether the red object will come into contact with the yellow object once the dynamics of all the objects have been completed. Since the task does not involve any language processing, we construct an MLP model that takes the rollout slots as input. The MLP processes these slots and produces a binary prediction, indicating whether the red object and the yellow object "touched" or "did not touch".
PHYRE Readout In the PHYRE action planning task, an action involves determining the size and position of the red ball. In our approach, we utilize the set of 10,000 predefined actions introduced by [4] and train a readout model to determine if a given action can solve the task. To construct the readout model, we draw inspiration from Wu et al. [44] that design a 2-layer MLP model on top of the encoded states. The readout model takes the predicted rollout states as input. To process these states, we employ an encoder, which differs depending on the specific model variant used. In SlotFormer, a Transformer is used, while in SwitchFormer, a Switch Transformer is employed. An MLP is operated as the encoder in the NPS model, and a Multi-head Self-Attention mechanism is utilized in the RSM model. Once the encoder has processed the states, the classifier generates a binary output indicating whether the task has been solved. This output serves as an inference for the task's solvability.
D Further Discussion on Experiment Results
D.1 Rollout Future Prediction
Figure 9 depicts the rollout frames generated by RSM alongside the baselines. Notably, RSM excels in producing robust future frames that accurately capture the dynamics of objects while maintaining visual fidelity. Nevertheless, we have encountered a challenge in generating objects with sharpness within the CLEVRER dataset. SlotFormer marks partially incorrect object's dynamics in this case.
D.2 Discussion on the Action Planning task in PHYRE
We encountered difficulties in reproducing the action planning results of SlotFormer, even when using their provided checkpoints. In Table 3, we report the iid result of 76.4 for SlotFormer, which is 5.6 points lower than the officially reported value. To investigate this issue, we explored potential reasons
Figure 9: Comparison of rollout frames in CLEVRER.
Figure 10: The scaling of RSM's parameters with different amounts of mechanisms.
and identified the following possible factors: (1) Instability of results: The official GitHub page of SlotFormer $^{11}$  acknowledges that the results in this specific task may exhibit instability. This suggests that achieving consistent and reproducible results with SlotFormer can be challenging and (2) data discrepancies: The PHYRE dataset, being regenerated rather than downloaded from a common source, introduces variations in the computing configuration. These data collection and processing differences may contribute to the disparities observed between our results and the reported values in Wu et al. [43].
D.3 Finetuning Results in RSM
In Figure 11, we present the results of hyperparameter fine-tuning on CLEVRER and Physion datasets. This section explores the impact of the number of mechanisms, the expansion of  $\psi(\cdot)$  parameters, and the structure of the mechanism models. The term number of layers in this analysis refers to the hidden layers within the MLP structure that maps the input dimension to the output dimension in the respective models.
In terms of the number of mechanisms, we have observed that having either a large or a small value for  $M$  significantly worsens the results and leads to high variance. However, we have also discovered that the tasks can be solved with relatively few mechanisms, even when dealing with diverse object movements. To validate the design principle stated in Appendix C.2.2, which proposes an inverse relationship between the number of the mechanism's parameters and the number of mechanisms, we conducted experiments denoted as  $10^*$ , which replicates mechanisms to achieve a total of 10 mechanisms without reducing the number of parameters, as compared to the best configuration (the first boxes). Our findings indicate that the replicated configuration ( $10^*$ ) achieves a slightly lower score than the best configuration, and the additional mechanisms are not selected by the  $\psi(\cdot)$  models.
In terms of the number of layers in  $\psi(\cdot)$ , the challenge is to map a  $2 \times d_s$  vector to a compact vector of size  $M$ . Our findings suggest that using 1 or 2 hidden layers yields favorable results in terms of achieving a high score. However, we find a complication in identifying consistent patterns for fine-tuning  $\psi(\cdot)$  across different datasets.
$^{11}$ github.com/pairlab/SlotFormer
(a) CLEVRER. The lower score indicates better performance.
(b) Physion. The higher score indicates better performance.
Figure 11: The finetuning results in CLEVRER and Physion. The first box of each subplot is the configuration that achieves the best performance, whereas other boxes are performance that has a different number of mechanisms, layers of the  $\psi$  model, or layers of each mechanism from the best setting. The plotted values are the mean and standard deviation over 5 different runs corresponding to each configuration. See text for more details.
Lastly, when considering the mechanism structure, we have noticed that increasing the number of parameters allows a single mechanism to achieve a moderately decent score, although not a high score. Consequently, the  $\psi(\cdot)$  model tends to choose only one mechanism for all cases, resulting in a lower score when assigning 5 or 7 layers to the mechanism.
E Experiments on End-to-end Training Pipeline
We propose an additional experiment to verify the methods' ability to model objects' dynamics from scratch, additionally, in action-conditioned environments. We establish an end-to-end training pipeline that comprehensively assesses the models' effectiveness in the entire process of extracting slots from frames, handling the objects' action-conditioned dynamics, and finally decoding slots back into frames.
RSM generally demonstrates a superior ability to model objects' dynamics and produce meaningful slots compared to the baselines.
E.1 Experiment Setup
Environment: This dataset consists of objects arranged in a  $5 \times 5$  grid. At each time step, a single frame and an action are provided. The action specifies one object and one manipulation from the set of UP, RIGHT, DOWN, and LEFT. The challenge of this task is to determine the feasibility of the given action and predict the resulting frame. For example, an action of moving an object to the LEFT is executable if no objects are obstructing the left side of the target object. Otherwise, the frame remains unchanged.
Encoder and Decoder architecture In this experiment, we follow the encoder proposed by Kipf et al. [24] that contains a simple CNN-based Object Extractor to extract  $N$  feature maps from the input frame that is followed by an MLP-based Object Encoder to encode feature maps to vector representations of objects, i.e. the object slots. Afterward, we exploit a decoder architecture for slot-based visual prediction, consisting of  $N$  slot decoders with separate weights. Each slot is decoded into an RGB reconstruction, and the final frame reconstruction is obtained by summing the reconstructions of all slots.
Training Objective. This experiment follows the Contrastive loss setup as Kipf et al. [24] that uses the prediction result of the transition model to form the positive hypothesis, whereas, sampling random input states in the same batch forms the opposing hypothesis. The target slots,  $s_{t+1}^{*1:N}$ , are obtained by passing the target next frame through the Encoder,  $s_{t+1}^{1:N}$  denotes the predicted slots, and  $\tilde{s}_t^{1:N}$  indicates the random slots in the training batch.
$$
H = \mathsf {M S E} (s _ {t + 1} ^ {1: N}, s _ {t + 1} ^ {* 1: N}), \quad \tilde {H} = \mathsf {M S E} (\tilde {s} _ {t} ^ {1: N}, s _ {t + 1} ^ {* 1: N})
$$
$$
\text {C o n t r a s t i v e L o s s}: \mathcal {L} = H + \max  (0, 1 - \tilde {H}) \tag {8}
$$
We consider the sum of two BCE loss terms in training Decoder,  $\mathcal{L}_1$  and  $\mathcal{L}_2$ .  $\mathcal{L}_1$  is applied on  $x_{t}^{\prime}$  obtained by passing  $s_t^{1:N}$  through Decoder, which is expected to be close to the input frame  $x_{t}$ .  $\mathcal{L}_2$  is applied on  $x_{t+1}^{\prime}$ , obtained by passing  $s_{t+1}^{1:N}$  through Decoder to achieve the prediction of next frames reconstruction, which is desired to be close to the target next frame  $x_{t+1}$ .
$$
x _ {t} ^ {\prime} = \operatorname {D e c o d e r} \left(s _ {t} ^ {1: N}\right), \quad x _ {t + 1} ^ {\prime 1: N} = \operatorname {D e c o d e r} \left(s _ {t + 1} ^ {1: N}\right)
$$
$$
\mathcal {L} _ {1} = \operatorname {B C E} \left(x _ {t} ^ {\prime}, x _ {t}\right), \quad \mathcal {L} _ {2} = \operatorname {B C E} \left(x _ {t + 1} ^ {\prime}, x _ {t + 1}\right) \tag {9}
$$
$$
B C E L o s s: \mathcal {L} _ {D e c o d e r} = \mathcal {L} _ {1} + \mathcal {L} _ {2}
$$
E.2 Experimental and Analytical Results
Figure 12: Observation of mechanisms assignment and performance in 5 rollout steps in 2D Shapes. In the last three rows, we present the reconstruction by only changing 1 slot with 1 mechanism applied to that slot over 5 steps, while all other slots are untouched.
Disentangling objects transition dynamics to mechanisms. In Figure 12, we study the role of each mechanism in RSM. The analysis shows that RSM produces a reasonable reconstruction compared to the ground-truth frame and encourages the mechanisms to distinguish themselves in their roles. In this sample, we can infer the 5 slots corresponding to: 1: red circle, 2: blue triangle, 3: green square, 4: purple circle, and 5: yellow triangle. Likewise, we observe the mechanism assignment in each step and list the role of each mechanism (5 mechanisms in this experiment) specialized as the following: Mechanism 1: RIGHT, Mechanism 2: UP, Mechanism 3: LEFT, Mechanism 4: DO NOT MOVE, and Mechanism 5: DOWN.
Looking deeper into the reconstruction results, RSM takes advantage of the CCI to assign a suitable mechanism for each slot in all scenarios, considers the particular situation, and reacts accordingly to the same action given the specifics of each context. For instance, we observe that with the same action UP given in step 1 on the green rectangle and step 3 on the red round, RSM recognizes the situations in which the object is allowed to move or is blocked by the upper wall, respectively, then applies the movement at step 1 while not modifying objects at step 3 and generates the correct reconstruction in both cases. A similar example is shown in the row corresponding to mechanism 2 at step  $3\rightarrow 4$  when the green object does not move UP and remains at the same position since the red object blocks it.
Figure 13: Comparison of extracted feature maps from a scene and reconstructions in 3D Cubes. RSM deals with object slots better than baseline in both slot extraction and slot decoding phases.
Figure 14: Reconstruction comparison on 3D Cubes dataset
The ability to decompose a frame into slots. We analyze RSM's slot-centric representation learning ability. In Figure 13, we illustrate a comparison of the extracted feature maps with a size of  $10 \times 10$ , which are constructed by the Encoder model and the reconstructed slots and the frame with a size of  $3 \times 50 \times 50$ , acquired by the SlotDecoder models that receive the input as the predicted next state. We find that RSM decomposes the input frame into separated slots and keeps each object in the same slot until the decoding phase. In contrast, the baselines do not capture all objects but produce noisy feature maps (SlotFormer and SwitchFormer), or put the same object in two slots and identify two objects in another slot (NPS).
We observe that combining the CCI with the sequential updates that encourage slots to observe the modification of each other benefits from recognizing the overlap of objects. More specifically, RSM only generates a part of the object in case that object is covered by another one (e.g., the green and blue cubes in slot 2 and slot 4 are partially covered by the yellow cube). On the other hand, the baselines overlook the condition of executability of action and generate overlapped objects in some cases (e.g., the green, blue, and yellow cubes in SlotFotmer overlapping one another). Lastly, SwitchFormer produces blurry objects and incorrectly predicts objects' dynamics.
Reconstruction in 3D Cubes. One of the challenges of generating reconstructions in 3D Cubes is to recognize the visibility order of objects. Figure 14 illustrates RSM's strength in communication among slots to obtain the order information, as well as its ability to generate the proper movement of slots and produce an accurate reconstruction compared to the ground truth. In contrast, SlotFotmer misses that kind of information from the beginning steps and renders the blue and green objects inside each other. Besides, other methods lose the information about some objects and produce inaccurate reconstructions at the end, witnessing a huge gap from the preceding steps up until step 10.
While RSM has demonstrated robustness for modeling objects' dynamics in various tasks in both iid and OOD settings, there is a lot of room to expand this work to overcome the following limitations:
1. Sensitivity to Hyperparameters: RSM requires tuning the number and size of mechanisms. Future research could explore automated methods for determining optimal values and enhancing RSM's adaptability across tasks and scenarios.
2. Time complexity: The sequential update of slots can become computationally expensive in larger systems with many slots (e.g., exceeding 10,000). Although this problem does not affect our work (see Appendix A), future research should also explore parallelization techniques to improve computational efficiency by assigning mechanisms and predicting the next states simultaneously for large-scale systems.
3. Environments in this study are all observable. Future work should explore a wider range of observable and unobservable environments to gather deeper insights into the working mechanism of RSM in such domains.