Next state prediction gives rise to entangled, yet compositional representations of objects

I want to clarify that I agree that the paper provides robust evidence in terms of numbers, but I have very strong doubts towards the conclusion due to the baselines.

If you read the CSWM paper, you will find that their main contribution is about the training loss --- the constrasive loss, not the neural architecture. In fact, their neural architecture, while seems to be "slot based", is in fact very adhoc--- they just used the different channels of a convolutional layer as "slots". It is not surprising that their reconstruction does not perform very well.

The main selling point of the slot attention paper is about the disentangled representations, not about accurately predicting object dynamics. Which is quite understandable to me why they underperforms in the paper's experiments.

I still think this paper's baseline are too weak. Both CSWM and slot attention, at least in their original form, are TOY models designed for their specific task, they are not general purpose models like transformers. These models are NOT meant to be generalizable to other tasks and datasets. It makes no sense to me to compare with them, except you use the exact same dataset and criterion used in their paper.

I dont think you can draw any conclusions just comparing with them but not with recent scaling up versions like SAVI++ and PSL.

I maybe too negative about the paper in the beginning, I think the claim that "next token predict can result in linearly separable representations" is well supported. But this is not that new. Predicting the future is a good objective for learning video representations is a VERY well known fact among deep learning community. For example, the 2019 paper Video Representation Learning by Dense Predictive Coding used similar loss for video representation learning. I think the only novelty is to pose the problem specific to object centric case. Surprising? Not for me.

But the claim that "the proposed model can match or even better than slot based algorithms in downstream tasks" is not well supported. The reason is that the baselines are too weak.

I am raising the score to 3, but I still dont think this paper should be published in the current form.

We thank the reviewer for engaging with our work. However, we were quite puzzled about this review. Although the reviewer first states that our paper “provides robust evidence across diverse datasets”, they later write that “The claims about distributed models achieving compositional representations are not robustly supported by the experiments.“. The reviewer also claimed that “it is unfair to compare your model with small models designed for special dataset, like CSWM, on new datasets such as multi-dsprite”. We firmly disagree with this statement, as the Multi-dSprite environment we construct is very similar in character and complexity to the Cubes and 3-body Physics datasets CSWM was trained on in the original paper. Moreover, we even show in Figure 2 that CSWM can get close to a 100% test accuracy when trained on 50k transitions, suggesting that it is well-calibrated to deal with this dataset. Finally, the CWM model we compare it to only differs from the CSWM model in that it uses an MLP for dynamics prediction instead of a GNN, and uses a generic CNN encoder instead of one mapping images to slots. We therefore do not believe that a rating of 1 with confidence 5 is justified. We go through the reviewer’s comments in more detail below.

The CSWM baseline, a model specifically optimized for the datasets used in their studies, underperforms on the new datasets used in this paper.

We do not believe that CSWM underperforms. In the dSprite environment it eventually attains close to 100% test accuracy when given enough data. In other environments, like 3-body physics, our results for CSWM match those reported in the original paper - still CWM outperforms CSWM, despite not having slots.

The study relies on Slot Attention, which is designed for static images, as a baseline in dynamic scenes.

We want to clarify that we train the Slot Attention model as a static auto-encoder, reconstructing static scenes from two MOVi-based datasets. We only evaluate it on its ability to reconstruct static scenes, not to perform dynamic prediction.

The claim that next-token or next-state prediction without specifically designed inductive biases can lead to a (somewhat) disentangled representation is not new. This phenomenon has been observed and documented in previous works, making the findings here less innovative.

We would like to encourage the reviewer to reference the papers they had in mind so we can assess the novelty of our findings. We are unaware of any previous works showing that next-state prediction can produce separable representations of objects. We also note that the novelty of our work was highlighted by all other reviewers.

In Eq.6 , the t should be replaced by t+1?

This equation refers to the static Auto-Encoder, which does not perform next-state prediction. It is only trained to reconstruct static images, as is the Slot Attention model.

We thank the reviewer for their thoughtful and encouraging review of our paper. We are glad the reviewer found our paper well written and interesting. The reviewer also provided some important suggestions for improving the presentation of our results. As a consequence we have added a figure summarizing all dynamics models’ scores for our prediction accuracy metric for all datasets. We also added a figure summarizing all the reported linear separability scores into one figure for all datasets. We address the reviewer’s points in more detail below.

For the dynamic loss (L_{AE-dynamic}), have the authors tried each term separately in addition to both combined?

We did not experiment with ablating the latent consistency loss, as this loss term has been shown to be important in other works [1, 2]. Ablating the pixel reconstruction loss would lead to representational collapse, as the encoder can minimize the other loss term by mapping all observations to a constant vector [3].

Instead of plotting the different metrics with respect to the number of training data, it would be interesting to do the same for various sizes of latent representation. Augmenting the capacity of this latent space will affect the information contained, making it more or less distributed -- and potentially showing different scores on the metrics. Have the authors tried this?

Thanks for this interesting suggestion. We trained CWM with various latent representation sizes (5, 25, 50, 100 and 250, respectively) in the Multi-dSprite environment. While too low latent dimension sizes yielded worse accuracy, once the representational space becomes big enough (~50), accuracy saturates, suggesting that the encoder needs to be of suitable capacity to represent objects in a disentangled way (see Figure 13B, page 17).

the accuracy of both models increases with more data, but the gap seems to decrease. Is there a point with even more data where the CSWM is as good as the CWM?

Indeed, CSWM catches up with the CWM in the Cubes and Multi-dSprite environments when trained on 50k transitions. Adding more data in the other environments is therefore likely to continue to increase CSWM performance.

why choose the CSWM as the baseline for the RSA if it's not the best model? It would be interesting to perform this analysis on all the models with respect to the CWM which seems to perform best on all the metrics, to see which component plays the biggest role in explaining the quality of the representations.

Thanks for this suggestion. We added a figure where we show how well CSWM, as well as the static and dynamic auto-encoder models align with CWM. Although the dynamic auto-encoder develops representations with a similar degree of object-separability to CWM, its representations are on average less aligned than the CSWM. This suggests that model representation may still differ in important ways despite representing objects in separable subspaces. See Figure 16, page 18.

It would also be informative to add in the Appendix the complete architecture of all the models, along with their number of parameters.

We thank the reviewer for the suggestion, we report the architecture, hyperparameters and parameter counts for the models in Appendix B, page 19.

[1]Watter, Manuel, et al. "Embed to control: A locally linear latent dynamics model for control from raw images." Advances in neural information processing systems 28 (2015).

[2]Hafner, Danijar, et al. "Dream to control: Learning behaviors by latent imagination." arXiv preprint arXiv:1912.01603 (2019).

[3]Schwarzer, Max, et al. "Data-efficient reinforcement learning with self-predictive representations." arXiv preprint arXiv:2007.05929 (2020).

We would like to thank the reviewer for their thorough review of our paper. We appreciate that the reviewer found our paper well written and our findings novel and interesting. At the same time the reviewer raised a number of important points. These pertained to i) compositional generalization, ii) downstream tasks and iii) validating forward performance.

First, we conducted two new experiments targeting compositional generalization in the MOVi environment (see Figure 8, page 15): in the first experiment, we trained Transformer based CSWM and CWM to predict dynamics of up to 4 objects, and tested on dynamics of 5 to 8 objects. In the second experiment, we trained the same models to predict dynamics of two red cubes and two green spheres, and tested on the dynamics of two green cubes and two red spheres (color flip). In both experiments we see that there’s a train-test disparity that gradually diminishes with more and more training data for both CSWM and CWM. CWM retains an edge in sample-efficiency. In summary, we see that CWM (without slots) performs well on the new compositional generalization experiments when given sufficient data.

We then assessed the downstream task performance of our pretrained models. Here we constructed two new tasks - one control task and one prediction task. In the control task, we train a Soft Actor Critic (SAC) agent to manipulate a randomly sampled sprite in the Multi-dSprite environment to go to a particular location on the grid. The SAC agent receives observations that are the encodings of the scene produced by one of the pretrained models. We evaluate the agent using the embeddings of CSWM, CWM and the auto-encoder, and see that the agent trained to perform control using CSWM representations performs the best, with the CWM-based agent trailing closely behind (see Figure 9, page 15). This suggests that slotted representation could offer advantages in control tasks like this.

In the second downstream task we used the representations of the trained MOVi-A models to predict a novel quantity. We froze the encoders of CSWM and CWM and trained a linear classifier to predict the number of objects present in a scene. Here we see that both models can predict object cardinality better than chance with a simple linear classifier with only minor differences in prediction accuracy between them (see Figure 10, page 16). Moreover, prediction accuracy generally increases the more data the models were trained with in their original task.

Lastly, we followed the reviewer’s suggestion to further validate our forward accuracy metrics by training stop-gradient decoders to reconstruct images from latent states. In the Cubes and Multi-dSprite environments we trained CNN decoders for 100 epochs to reconstruct images from the representations of CSWM and CWM, respectively. We then evaluated the reconstruction accuracy with the LPIPS metric as the dynamics models predicted future states in an open-loop fashion. Matching our other prediction accuracy metric, we see that CWM retains lower LPIPS for future state predictions than CSWM in both environments. See Figure 11, page 16, and Figure 12 page 17 for example reconstructions and LPIPS scores for both environments.

Questions:

The overview of object-centric representations is focused entirely on slot-based models.

We thank the reviewer for this suggestion and have added the following paragraph in the related works section: "While we focus on slot-based models, there are different approaches to learning object-centric representations. Patch-based models such as SPACE (Lin et al., 2020) and MarioNette (Smirnov et al., 2021) decompose scenes into disentangled representations by reconstructing the input from patches. Keypoint-based models such as DLP (Daniel & Tamar, 2022) build on representations as sets of geometrical points as an alternative to single-vector representations."

What does it mean for the prediction to be “closest to the last frame”, closest compared to what?

We evaluate whether the predicted state is the closest to the last frame in the sequence out of all predicted states in the. This means that, in a test set of a 1000 videos, we compare the final encoded state with the predicted final states for all 1000 videos. The prediction is deemed correct if the final state is the closest (in terms of L2 distance) to the predicted final state in the corresponding video, and incorrect if it is closer to any of the other 999 predictions. We compute this for all videos in the test set and report the percentage of correct predictions. We added the following clarification on line 268:

"In other words, a prediction is deemed correct if the final encoded state is the closest in $L_2$ distance to the predicted final state in the corresponding video, and incorrect if it is closer to any of the other 999 predictions."

We thank the reviewer for their review and engaging with our work. We are happy that the reviewer found our paradigm interesting. The reviewer also highlighted some important points, especially regarding the definition of compositionality and whether our metric is trivially maximized by models with high-dimensional latent spaces. We sought to address these concerns by showing that CWM can generalize compositionally about i) scenes with more objects than it was originally trained on, and ii) about objects with different combinations of attributes than trained on (such as objects with swapped colors). We further show that our object-separability metric cannot be maximized simply by having large feature spaces. A randomly initialized network with 2000 latent dimensions shows only slightly better than random object separability, whereas a trained model with 50 latent dimensions is at ceiling. We believe that addressing these points has greatly improved our manuscript. We address each point in more detail below.

There is a lot of work on compositional representations in machine vision. From early papers focused on compositionality of attributes We thank the reviewer for highlighting this.

We have expanded the related work section with the following paragraph to include earlier papers on compositionality of attributes, as well as additional model types suggested by Reviewer xykj:

"Atzmon et al. (2016) and Johnson et al. (2017) introduced datasets to test compositional generalization in machine learning models. While we focus on slot-based models, there are different approaches to learning object-centric representations. Patch-based models such as SPACE (Lin et al., 2020) and MarioNette (Smirnov et al., 2021) decompose scenes into disentangled representations by reconstructing the input from patches. Keypoint-based models such as DLP (Daniel & Tamar, 2022) build on representations as sets of geometrical points as an alternative to single-vector representations."

In attribute compositionality, it was argued that discriminative models entangle peoperties that are correlated intheir trainin data, and cannot generalize to new combinations. But here, objects are not entangked dugin training

We thank the reviewer for raising this point. We conducted a new experiment in which CSWM and CWM were trained on object dynamics where objects were entangled with their attributes (cubes were always red and spheres always green), and tested on videos of objects where these attributes were flipped. Again we see that the models can generalize about the novel test combination of objects and attributes when given enough data (see Figure 8, page 15).

In high enough dimensions, it is easy to get linear separability, depending on the number of classes.

This is a good point. In high enough dimensions linear separability could be trivial. This is why we install an L1 norm penalty on the linear classifier weights when we perform the linear separability analysis. We further tested whether models with high-dimensional latent spaces trivially attained high separability scores. We randomly initialized image encoders with 50, 100, 500, 1000, and 2000 latent dimensions, respectively. Without training, these models perform only slightly better than chance levels, whereas a trained CWM model with 50 latent dimensions is close to ceiling performance on our separability metric (see Figure 13, page 17).

What specific properties of the architecture or training procedure lead to linear separability effect?

We find two factors important for separability: Training set size, and next-state prediction. We further see that the dynamic contrastive models (CWM) generally attain higher object decodability scores than auto-encoding models. We show the combined scores for all models in Figure 15, page 18.

How general are the results in the paper? Would they generalize to other object-centric approaches beside CWM / CSWM?

We additionally show in the paper that models trained with auto-encoding objectives develop separable representations of objects when trained to predict dynamics. This suggests that this phenomenon might also generalize to other architectures and training setups.