Reusable Slotwise Mechanisms

Summary

This paper presents an architecture (Reusable Slotwise Mechanism or RSM) for modelling the temporal dynamics of objects based on a slot-represent. The main idea is to extract a global context (Central Contextual Information or CCI) from the past frames and use this to stochastically select one of a small discrete set of (learned) mechanisms for each slot and time-step. The mechanism computes the update of a slot, based on its current representation as well as the CCI. The slots are updated sequentially. It is shown anecdotally that they specialize to particular object behaviors/dynamics (such as idle, free-fall and collision) and are reused across slots and time. RSM is compared to several strong baseline methods on several tasks including dynamics modelling, question answering, and task planning. It demonstrates significant improvements over the baselines with a comparable amount of parameters and compute.

Strengths

* The studied problem of modelling temporal dynamics from raw video is an important one, and there still is a lot of headroom. * The idea behind RSM is simple and well motivated. * RSM clearly improves upon the baselines and produces SOTA results on several relevant benchmarks. It thus presents an important step towards unsupervised learning of dynamics models from video. * The presentation is clear and the paper + supplementary contain all relevant information for reproducing the experiments.

Weaknesses

To evaluate the different parts of the RSM architecture, the paper presents 4 ablations in section 4.5. Unfortunately the only quantitative results in Fig 5. the present fell out of the main paper into the appendix. But if I understand correctly there is a more severe problem: it seems from the text that ablations RSM!2, RSM!3 and RSMk are **inference time** ablations. That means during training the model had access to CCI in steps 2 and 3 and has learned to choose a mechanism. It is thus not at all surprising that changing any part of that during inference will deteriorate the performance. Thus, all that these ablations demonstrate is that the model has not actively learned to ignore the CCI and the mechanisms. It would be much more informative to see how the model performs if these modifications are also applied during training. That could for example help answer the much more interesting question "How much does having access to CCI during step 2 help the model choose the right mechanism". For this very reason ablation 4 is much more informative. In my opinion the such training-time ablations are necessary to properly support the claims about the role and importance of CCI and the mechanisms. Other minor weaknesses: * the differences between RSM and the baselines could be made clearer. There is some information in Sec 3.2 but it doesn't clarify what exactly is different. * The citations especially in the introduction and related work seem a bit biased with many citations to Goyal et al while leaving out other relevant and influential work in the area such as: - Eslami, S. M., Heess, N. & Weber, T. Attend, infer, repeat: Fast scene understanding with generative models. Adv. Neural Inf. Process. Syst. (2016) - Burgess, C. P. et al. MONet: Unsupervised Scene Decomposition and Representation. arXiv [cs.CV] (2019) - Greff, K., Kaufman, R. L. & Kabra, R. Multi-object representation learning with iterative variational inference. International (2019) - Battaglia, P. W. et al. Relational inductive biases, deep learning, and graph networks. arXiv [cs.LG] (2018) - Greff, K., van Steenkiste, S. & Schmidhuber, J. On the Binding Problem in Artificial Neural Networks. arXiv [cs.NE] (2020) - Schölkopf, B. et al. Toward Causal Representation Learning. Proc. IEEE 109, 612–634 (2021) * Sec 2.2 feels a bit repetetive with the enumeration of the three main components of RSM and the list of four steps in RSM mostly explaining the same things.

Questions

* My main request would be training-time ablations for $RSM_{!2}$ and $RSM_{!3}$ with quantitative results, since they would help buttress the claims about the importance of the CCI. * Similarly, the paper claims in various places that the CCI has advantages over pairwise interactions. Is there evidence from a fair comparison of the two? Communication is pairwise in NPS, but is that the only relevant difference? * How big is $d_{cci}$? The paper claims that it acts as a bottleneck, yet the CCI is recomputed for every slot from a self-attention among all previous slots. So unless $d_{cci}$ is at least smaller than $d_s$ I do not see in what sense it constitutes a bottleneck. * Are all mechanisms equally important? I suspect that main advantage of RSM might be the ability to put objects firmly in an idle mode.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

The paper discusses some limitations in the supplementary. I see no necessity for a discussion of negative societal impact.

Summary

This paper presents a study on future prediction using object-centric representations and disentangled and reusable mechanisms (RSM) that govern the interactions between objects. The authors' key contribution lies in the proposal of Central Contextual Information (CCI) representation, which captures the interactions between objects observed in the past. The CCI is encoded using multi-head attention layers with all object-centric representations from previous observations and a bottleneck projection layer. It aids the model in selecting a single mechanism to predict object transitions from the current state to the next state. Experimental results demonstrate that the introduced RSM outperforms several competing baselines in tasks such as future frame prediction, visual question-answering, and action planning in relatively simple environments.

Strengths

+) The motivation behind the proposed CCI is clearly explained, with the authors using multi-head attention layers to capture interactions between objects in past observations. The subsequent bottlenecking of the attention layers’ output into the CCI enables the model to learn a concise representation that encapsulates meaningful information about object interactions. +) The selection process for predicting state transitions is intriguing and surprisingly effective. The use of Gumbel-softmax (assuming it is Gumbel-softmax instead of Gumbel-max as mentioned in the manuscript) allows the model to decompose the potential mechanisms governing state changes. +) The visualizations in Figures 4 and 5 provide compelling evidence of the model's ability to decompose mechanisms into discrete and reusable components. Figure 4 demonstrates the model's capability in decomposing mechanisms, while Figure 5 showcases how the CCI encodes essential information for determining outcomes following object interactions. +) The experimental results support the efficacy of the proposed RSM, as it outperforms several competing baselines, including SlotFormer, SwitchFormer, and NPS, across various tasks such as future frame prediction, visual question-answering (which relies on future visual states), and action planning.

Weaknesses

-) The assumption that there is only one mechanism responsible for the state change of an object may not always hold true. Factors such as the object's original momentum and collisions with other objects can influence its behavior. Therefore, the use of Gumbel-softmax (again, assuming it is Gumbel-softmax instead of Gumbel-max as mentioned in the manuscript) in Equation (2) might impose overly strong restrictions on modeling interaction mechanisms. -) There appears to be a misinterpretation regarding the mention of the Gumbel-max layer in Step 2. The Gumbel-max layer, without the Straight-Through (ST) Gumbel Trick, is not differentiable and is primarily used to enable the categorical sampling using unnormalized log probabilities by the reparameterization trick. The references cited at Line 122 also pertain to Gumbel-Softmax, not Gumbel-max. For more details, please refer to the PyTorch implementations of the Gumbel-Softmax layer: https://pytorch.org/docs/stable/generated/torch.nn.functional.gumbel_softmax.html#torch.nn.functional.gumbel_softmax or the TensorFlow API: https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/RelaxedOneHotCategorical. -) There is a discrepancy at Line 269 where the authors mention the visualization of disentangled mechanisms in Figure 5. However, Figure 5 is not found in the main submitted manuscript but in the supplementary material (Page 13). The mention of Figure 5 at Line 269 implies it should be present in the main paper, potentially violating the Formatting instructions. -) While the proposed RSM outperforms competing baselines on various datasets and tasks, it is worth noting that the experiments are conducted in relatively unrealistic environments with toy objects and simple interactions. It would be beneficial to demonstrate the effectiveness of the proposed approach (object-centric representations + RSM) on more realistic datasets such as MOVi-B, C, D, E, and F.

Questions

o) The usage of the non-differentiable Gumbel-max layer needs clarification. If it is indeed used, how are the gradients backpropagated to the layer $\psi$? If not, is it a typo? o) In Equ. (3), if only one mechanism is selected, why does the formulation still require a sum operation to combine all other transitions predicted by other mechanisms? o) Have the authors experimented with using the Softmax operation in Equ. (2) instead of the hard selection provided by the Gumbel-softmax? o) Can the authors demonstrate the scalability of the visualization of mechanism assignments/decompositions with more scenarios? For example, can they provide additional visualizations of mechanism decompositions with different initial states/transitions on different datasets such as OBJ3D or CLEVRER? o) In Line 101, when mentioning the projection $\phi$, is it referring to a linear layer or an MLP (multi-layer perceptron)?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

-) As mentioned in the Weaknesses, the experiments conducted in this study are limited to relatively unrealistic environments with toy objects and simple interactions. As a result, it remains uncertain whether the proposed approach would maintain its effectiveness in a more realistic environment or when applied to real-world data.

Summary

The paper proposed Reusable slot-wise mechanisms wherein the authors introduce "Central Contextual Information" -- a bottleneck that captures the global context and helps in choosing which slot to attend to. Through experiments on several simple and relatively complex datasets are shown with reconstruction, VQA and planning tasks to show the benefit of the reusable mechanisms.

Strengths

1. I find the introduction a bottleneck CCI to summarize the context of past $T$ frames to be intuitive and novel. 2. The method shows significant improvement in OOD action planning tasks which is impressive!

Weaknesses

3. **CCI**: Have authors tried treating each slot independently as input to the transformer with positional encoding instead of concatenating past $T$ time slots? 4. In Fig 1(c), is should the $t$ in blue circle (drawn adjacent to the CCI bubble) be $s_t$. If so please clarify in the figure as $t + \delta s_t $ in Step 3 is confusing as one represents time and another represents the predicted change in slot representation. 5. **Memory Constraints**: The CCI (Step 1) has a huge memory constraint of $(\tau+1) \times N \times d_s$ which is $7 \times 6 \times 128 = 5376$ for OBJ3D dataset (which has the smallest memory requirements), and $16\times 6 \times 192 = 18432$ which is huge. This appears to be a huge factor given that object-centric models are already memory expensive when it comes to applications such as RL. 6. **OOD in dynamics**: How does RSM behave when it sees an OOD dynamics for a particular object. For example, in the OBJ3D dataset (or any other simplistic dataset with enough variations in objects and their dynamics) if a green object is static throughout during the training, but during testing scenario if the green object moves would RSM reuse the mechanism corresponding to the motion of object? If time permits, I would like to see a simple experiment on thise on any of the datasets (which ever is faster to train/test). Please reach out to me if you need more clarification on this. 7. **Inference Time**: Can the authors provide comparison between all the baselines and RSM with regards to inference time on the different datasets evaluated? **[Minor comments which I have not considered for rating of the paper]** -- The authors need not reply to these 2 comments below. 8. I think reconstruction results shown using SSIM or MSE on most simplistic datasets such as OBJ3D or CLEVERER where the values are in the 90s aren't really indicative of how useful the representations are for downstream tasks. It appears that VQA task in CLEVERER is also pretty simplistic and has reached mid-90s -- it would be better the community in general moves to addressing complex tasks of planning, RL and hard-VQA (such as Physion VQA task as shown in the paper). 9. **Readability of methodology section**: In section 2.1 (RSM Overview), it would be helpful to re-iterate what the purpose of having CCI is. The only mention of CCI is in the last section of introduction, so explicitly having it as a part of the RSM overview would be helpful. ---- **Rationale for rating** I find the contributions of this work to be novel and experimentation to be through and hence I lean towards accepting the paper. However, I would like the authors to address my questions before I finalize my decision. In addition to this, if authors have clarifying questions regarding comment (6), I would request them to reach out to me early on during the rebuttal phase.

Questions

See *Weaknesses*

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

NA

Summary

This work presents a novel framework to model object dynamics by leveraging communication slots in a modular architecture. The method uses Central Contextual Information (CCI) to allow information exchange among existing slots. The authors demonstrate the efficacy of their method's superior empirical performance on VQA, action planning, and future frame prediction.

Strengths

1. The idea of using central contextual information is novel and interesting and seems to provide clear empirical benefits over prior work. 2. The authors have clearly outlined the primary contributions and have carefully designed ablations to support their claim.

Weaknesses

One question I have for the authors is how this method would be extended to a real scene where the entities that are interacting cannot be clearly separated. For instance, in CLEVRER or Obj3D, it's known ahead of time that the 3-5 objects would be moving around and interacting with each other. But in a real scene, finding that might get very challenging. I would be interesting in hearing author's perspective about that.

Questions

See weaknesses.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

See weaknesses.

Answer to Rebuttal

I would like to thank the authors for their comprehensive rebuttal. I appreciate the additional experiments run. It is surprising to me that the training time ablations generally perform worse than the corresponding test-time ablations. I would have expected the model to perform better when training and test settings are identical. Overall I think the additional results and clarifications do improve the paper, and I reaffirm my recommendation of a clear accept.

The authors' responses effectively address all my concerns

I would like to thank the authors for providing detailed responses to address my concerns and confusing points, especially regarding the confusion surrounding the gumbel-softmax issue. I'm content with this matter as long as the final version clarifies that the backpropagation is facilitated by the Straight-Through Gumbel Trick. Additionally, I appreciate the authors' clarification on Equation (3) and Line 101. I appreciate the authors for incorporating new experimental results with the selection mechanism utilizing the softmax operation. While it's interesting that this model slightly underperforms compared to the RSM with gumbel-softmax, I agree with the authors' reasoning that "more realistic scenes might benefit from such modifications". Lastly, I would like to thank the authors for introducing new experiments involving the MOVi-C dataset, along with additional visualizations of mechanism assignment in OBJ3D and CLEVRER. After considering other reviewers' comments, I don't have further concerns at this point. Given the authors' commitment to addressing the major points of confusion in the final version, I have increased to raise my score to Weak Accept.

Thanks for the rebuttal

Thank you to the authors for their rebuttal and showing some more visualizations on OOD as requested. The Fig 4 in the posted pdf during rebuttal sheds potentially important light on what object-centric dynamics models are (and are not) capable of in their current form. The oddly predicted shape of the blue ball during the OOD rollout is interesting example. Also, based on reviewers DY7c and 8CU7's comments -- I appreciate adding the **MOVi-C** dataset. As the authors are aware of it -- the model does not seem to be tune correctly yet as the visual results are blurry. > Lastly, we would appreciate it if the reviewer could help us realize in which settings our method could constitute a potential memory-related problem since even with batches as big as 512,... Sure. My concern was if RSM-like model is used in say model-based RL framework such as Dreamer -- because the policy learning happens in imagination (via rollout of the world model) -- it might be additionally expensive to do both `forward` and `backward` passes on the model. But given the authors response on memory and the inference speed (forward pass only) -- I think my concerns have been alleviated. Thanks to the authors once again. I will change my score post-reviewer + AC discussion period.