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

## Abstract

Abstract Compositional representations are thought to enable humans to generalize across combinatorially vast state spaces. Models with learnable object slots, which encode information about objects in separate latent codes, have shown promise for this type of generalization but rely on strong architectural priors. Models with distributed representations, on the other hand, use overlapping, potentially entangled neural codes, and their ability to support compositional generalization remains underexplored. In this paper we examine whether distributed models can develop linearly separable representations of objects, like slotted models, through unsupervised training on videos of object interactions. We show that, surprisingly, models with distributed representations often match or outperform models with object slots in downstream prediction tasks. Furthermore, we find that linearly separable object representations can emerge without object-centric priors, with auxiliary objectives like next-state prediction playing a key role. Finally, we observe that distributed models’ object representations are never fully disentangled, even if they are linearly separable: Multiple objects can be encoded through partially overlapping neural populations while still being highly separable with a linear classifier. We hypothesize that maintaining partially shared codes enables distributed models to better compress object dynamics, potentially enhancing generalization.

## 1 Introduction

Humans naturally decompose scenes, events and processes in terms of the objects that feature in them (Tenenbaum et al., ; Lake et al., ). These object-centric construals have been argued to explain humans’ ability to reason and generalize successfully (Goodman et al., ; Lake et al., ; Schulze Buschoff et al., ). It has therefore long been a chief aim in machine learning research to design models and agents that learn to represent the world compositionally, e.g. in terms of the building blocks that compose it. In computer vision, models with *object slots* learn to encode scenes into a latent, compositional code, where each object in the scene is modelled by a distinct part of the latent space. This strong architectural assumption allows the models to learn representations that improve compositional generalization (Brady et al., ; Wiedemer et al., ) and reasoning about objects (Wu et al., ).

Slotted representations are often contrasted with distributed representations. Models with distributed representations encode information about a scene, and potentially the objects that compose it, in overlapping populations of neurons. Can models with distributed representations learn to encode objects in a compositional way without supervision? And can distributed coding schemes offer advantages over *purely* object-centric coding schemes?

By compressing properties of multiple objects in a shared code, models with distributed representations could potentially gain richer representations where scene similarities are more abundant (Smola & Schölkopf, ; Lucas et al., ; Demircan et al., ; Garvert et al., ). For instance, if two objects are represented similarly, the model could use what it knows about the dynamics of one object to generalize about the dynamics of the other object. This could in turn facilitate learning, potentially at the loss of fully separable object representations.

Figure: Figure 1: Overview of the decoding analysis and datasets. A: We propose a simple test for assessing compositional object representations. After unsupervised pre-training on object videos, we evaluate the linear separability of models’ latent object representations. This is done by training a linear classifier on the absolute differences of two successive encoded frames where only one object changes. B: We evaluate the models on five datasets of dynamically interacting objects, ranging from simple depictions of blocks and sprites to realistic simulations of 3D objects.
Refer to caption: x1.png

In our study, we offer experimental evidence that models with distributed representations can learn compositional construals of objects in an unsupervised manner, when trained on sufficiently large datasets. Across 5 datasets that consist of dynamically interacting objects we see that models with distributed representations either match or outperform their slot-based counterparts in the tasks they were trained to perform (image reconstruction or dynamics prediction). Next, we define a simple metric inspired by Higgins et al. () that quantifies how accurately object identities can be linearly decoded from a model’s latent representations (see Figure 1). We see that as training data size increases, the models with distributed representations develop gradually more disentangled representations of objects. However, while object representations become separable, their properties remain encoded through *partially* overlapping populations of neurons, potentially allowing for richer generalization. Investigating the effect of training objective and loss function on the separability metric, we find that next-state prediction is a crucial component for the development of separable object representations for models without object slots.

While we see object disentanglement emerge without any supervision or regularization, this disentanglement is not absolute. Distributed models represent object properties in partially shared latent spaces. We speculate that this can facilitate generalization: When comparing models’ representations of object dynamics, we see clear clustering based on object identity (indicating separable object representations), but also clustering based on what type of transformation was applied to the object. This means that transformations such as object rotations, object scaling or object movements are represented more similarly, independently of which object they are applied to. Such compositional codes for group transformations of objects are made possible by the fact that all objects share a common latent space, instead of occupying separate ones, suggesting that there are benefits to distributed coding schemes.

## 2 Related work

Object-centric representations have been argued to improve the sample efficiency and generalization ability of vision and dynamics models in compositional domains (Elsayed et al., ; Kipf et al., ; Wiedemer et al., ; Wu et al., ; Locatello et al., ). Object-centric representations, given an appropriate model architecture, can also be learned in an unsupervised manner (Kipf et al., ; Locatello et al., ; Brady et al., ) and on real-world datasets (Seitzer et al., ). Previous studies have highlighted the importance of architectural features, like object-slots (Greff et al., ; ; Dittadi et al., ), as well as data properties, like having access to temporal information (Zadaianchuk et al., ).

There is a considerable overlap in the literature on object-centric and disentangled representations. A disentangled representation is one \saywhich separates the factors of variation, explicitly representing the important attributes of the data (Bengio et al., ; Locatello et al., ; Higgins et al., ; ). In a disentangled representation, a change in a single ground truth factor should lead to a change in a single factor in the learned representation (Locatello et al., ; Ridgeway & Mozer, ; Kim & Mnih, ). Information bottlenecks methods like $\beta$-VAEs have also been shown to be able to disentangle object features (Burgess et al., ; Higgins et al., ).

In recent work, Brady et al. () put forward a measure of representational object-centricness that measures \sayif there exists an invertible function between each ground-truth slot and exactly one inferred latent slot. We work with a related metric suitable for generic model classes (e.g. without image decoders) that instead measures the degree to which changes to individual objects can be predicted from changes in a model’s latent representations.

## 3 Methods

### 3.1 Models

We focus on models that learn representations of scenes in an unsupervised manner, e.g. without information about object identities provided as labels or masks. Unsupervised training regimes, such as auto-encoding (Kingma, ), denoising (He et al., ) and contrastive objectives (Chen et al., ) have shown promise as representation learning tools in many domains, ranging from image and language understanding (Radford et al., ) to reinforcement learning (Schwarzer et al., ; Gelada et al., ; Saanum et al., ). In this paper we investigate two classes of such training objectives: *i*) Reconstruction-based or auto-encoding objectives, where the goal is to encode and reconstruct images of scenes of objects. And *ii*), contrastive objectives, where the goal is to maximize embedding similarities of positive pairs and minimize similarities of negative pairs. Accordingly, the models rely on an image encoder, a Convolutional Neural Net (CNN) in our case, to map images of objects to latent representations. For auto-encoding models we additionally equip the model with a CNN decoder that maps the latent representation back to pixel-space.

$$ $\displaystyle z_{t}=e_{\theta}(x_{t})$ (1) $\displaystyle\tilde{x}_{t}=g_{\theta}(z_{t})$ (2) $$

Here $z_{t}$ is the model’s representation, and $e_{\theta}$ and $g_{\theta}$ are the CNN encoder and decoder. $x_{t}$ and $\tilde{x}_{t}$ are the image and reconstruction, respectively. We subscript image and representation variables with the time-point $t$ since our data are dynamic. Having access to this temporal information about the data, we also consider models that use future-state prediction as an auxiliary objective for representation learning. Observing how objects interact dynamically can provide the models with useful cues about object identities, and could facilitate learning systematic representations of objects (Zadaianchuk et al., ). When modelling the dynamics of the object data, we equip the model with a latent dynamics module that predicts the model’s representation at the next time point $t+1$, given the current representation, and an action $a_{t}$, if the dynamics data is accompanied by actions.

$$ $\tilde{z}_{t+1}=d_{\theta}(z_{t},a_{t})$ (3) $$

Here $d_{\theta}$ denotes the dynamics module, which is a Multi-Layer Perceptron (MLP) in the case that the dynamics are Markovian, e.g. fully predictable from the information provided in the current observation $x_{t}$ (and potentially action $a_{t}$). In non-Markovian settings we use a causal Transformer that integrates information across representations of past observations $(z_{t-n},...,z_{t-1},z_{t})$ to predict the dynamics. See Appendix B for details on the model architecture and hyperparameters.

For the auto-encoding models, we train the encoder and decoder to reconstruct the *current* frame from the current representation in the static setting, and the *next* frame from the predicted *next* latent representation in the dynamic setting. Here, we additionally train the dynamics model to minimize the distance between the predicted and actual representation of the next frame. This leaves us with the following loss functions:

$$ $\displaystyle\mathcal{L}_{\text{{AE-static}}}=||x_{t}-g_{\theta}(z_{t})||^{2}_ {2}$ (4) $\displaystyle\mathcal{L}_{\text{{AE-dynamic}}}=||x_{t+1}-g_{\theta}(\tilde{z}_ {t+1})||^{2}_{2}+||z_{t+1}-d_{\theta}(z_{t},a_{t})||^{2}_{2}$ (5) $$

We refer to these models as the *auto-encoder* and *sequential auto-encoder*, respectively.

For the contrastive models, we consider both a static and dynamic training scheme as well. In the static case, we present the model with a frame $x_{t}$ as well as a randomly augmented view of the same frame $h(x_{t})$ (Laskin et al., ; Grill et al., ). The model is then trained to minimize the embedding distance between the original and augmented view of the image, while maximizing the embedding distance between the original image and its representations of augmented views of other frames $x^{-}$ in the batch, up to a margin $\lambda$. In the dynamic setting we train the contrastive model as follows: Given an initial latent representation (and potentially action), we train the encoder and dynamics model to produce a prediction that is as close as possible to the encoded representation of the next frame $z_{t+1}$, and that is maximally far away from encoded representations of other frames $z^{-}$, up to a margin $\lambda$. The loss functions take the following form:

$$ $\displaystyle\mathcal{L}_{\text{{contrastive-static}}}=||z_{t}-e_{\theta}(h(x_ {t}))||_{2}^{2}+\max(0,\lambda-||e_{\theta}(h(x^{-}))-z_{t}||_{2}^{2})$ (6) $\displaystyle\mathcal{L}_{\text{{contrastive-dynamic}}}=||z_{t+1}-d_{\theta}(z _{t},a_{t})||_{2}^{2}+\max(0,\lambda-||z^{-}-d_{\theta}(z_{t},a_{t})||_{2}^{2})$ (7) $$

We refer to the static contrastive model as *CRL*, for Contrastive Representation Learner, and the dynamic contrastive model as *CWM*, for Contrastive World Model.

#### 3.1.1 Slotted models

We compare the auto-encoding models and CWM to baselines which attempt to learn slotted representations. As a baseline to the auto-encoding models, we implement Slot Attention (Locatello et al., ), an auto-encoder that reconstructs images as an additive composition of multiple object slots. The slots compete to represent the objects in the scene using an iterative attention mechanism, and the full model is trained with a simple auto-encoding objective as in equation 5.

The contrastive dynamics model is compared to a *structured* variant, the Contrastive Structured World Model (CSWM) (Kipf et al., ), that decomposes the scene into distinct object slots, and uses a graph neural network to predict how these object slots evolve over time. In non-Markovian settings, we replace the graph neural network with a Transformer encoder that applies spatio-temporal attention over a sequence of past object-slot representations, akin to the Slotformer architecture (Wu et al., ). Here too, the loss function exactly matches the one used to train the contrastive dynamics model with distributed representations.

### 3.2 Assessing object representations

How can we quantify the degree to which a non-slotted model has learned systematic object representations? While many metrics are possible, we propose one which is both simple and has connections to other metrics proposed to quantify representation disentanglement. If a representation of an object $o_{i}$ is disentangled from representations of other objects $o^{\prime}$, then a change to $o_{i}$ should only change one subspace of the models’ latent representation $z$. Additionally, this subspace should not be affected by changes to any of the other objects $o^{\prime}$. In other words, each object is represented across completely non-overlapping populations of neurons. Complete disentanglement is a tall order for models without the structural properties of slotted models. To get a continuous relaxation of this absolute object disentanglement metric we ask a related question of the models’ latent spaces: Given a set of changes to individual objects in a scene, how accurately can a linear classifier predict which object was changed from the resulting absolute difference in the model’s latent representations? The accuracy of this linear classifier on an evaluation set is our proposed metric. This metric is in fact a variation of the disentanglement metric proposed in Higgins et al. (), but applied to objects rather than ground truth generative features. Even if a model attains a perfect score on this metric, it does not necessarily mean that it represent objects in perfectly disjoint, non-overlapping populations of neurons. To illustrate, if a change to object $o_{i}$ always changes latent $z^{1}$ marginally and latent $z^{2}$ greatly, and a change to object $o_{j}$ changes $z^{1}$ greatly and $z^{2}$ marginally, a linear classifier can reliably separate the two objects in terms of the change in $z$, despite them having entangled representations. In other words, it is possible to use overlapping neural codes to represent objects, while still having object representations that are linearly separable. We investigate this further in Section 5.

In practice, we implement our metric by constructing datasets consisting of pairs of images $(x_{t}^{i},x_{t+1}^{i})$ from the data domain on which a model was trained. The only difference between these two images is that a single object has changed from time $t$ to $t+1$. For each such pair we associate it with a label $y^{i}$, a categorical variable indicating *which* object was altered from $t$ to $t+1$. We then extract a model’s representation of each frame in the pair (after it has been trained), and compute the vector of absolute differences between these two representations:

$$ $|\Delta^{i}|=|e_{\theta}(x_{t}^{i})-e_{\theta}(x_{t+1}^{i})|$ (8) $$

From the set of ensuing absolute difference vectors $\mathcal{X}=(|\Delta^{1}|,...,|\Delta^{n}|)$ we train a linear classifier to predict the corresponding object labels $\mathcal{Y}=(y^{1},...,y^{n})$ while minimizing the $L_{1}$ norm of the learned coefficients, as recommended by Higgins et al. (). We report the accuracy on a left out subset of $(\mathcal{X},\mathcal{Y})$ that the classifier was not trained on.

## 4 Experiments

We trained the two classes of distributed models, as well as their object-centric counterparts, on five datasets of dynamically interacting objects. Two of these datasets, cubes and 3-body physics, were introduced in Kipf et al. () to showcase how object-centric representations facilitate learning of object dynamics. Extending the evaluation, we created our own dataset of object interactions based on the dSprites environment (Matthey et al., ). This dataset consisted of four sprites with different shapes and colors traversing latent generative factors, such as $(x,y)$-coordinates, scale and orientation, on a random walk. Lastly, we trained our models on two more complex Multi-Object Video (MOVi) datasets generated using the Kubric simulator (Greff et al., ), a 3D physics engine for simulating realistic object interactions. We generated one dataset consisting of $14,000$ videos with a constant set of five cubes with fixed physical properties that interacted (initial object conditions such as directional velocities and position were randomized for each video). We refer to this dataset as MOVi (simple), due to the constant object properties. Additionally we trained models on the MOVi-A dataset, consisting of almost $10,000$ videos where the number of objects, their shapes and physical properties, such as mass and friction, are not fixed and vary across videos. The cubes and Multi-dSprites datasets had action variables that accompanied the videos, and were predictive of the way the objects would change from one frame to another. The other datasets were action-free. All models were trained for 100 epochs with five random seeds on the cubes, 3-body physics and multi-dSprites datasets, and for 125 epochs with three random seeds on the MOVi datasets. Furthermore, to assess the effect of dataset size on our evaluation metrics, we split each dataset up in different sizes.

We evaluated the slotted CSWM and distributed CWM models’ prediction abilities by measuring the accuracy with which they could predict novel object trajectories of length $n$ from an unseen evaluation set in an open loop manner. Prediction accuracy was estimated as the percentage of test trajectories where the predicted latent state $\tilde{z}_{t+n}$ at the end of the trajectory was closest in terms of Euclidean distance to the model’s representation of the last frame in the trajectory $z_{t+n}$ out of 1000 evaluation trajectories. For the cubes, 3-body physics and Multi-dSprites datasets, we conducted the evaluations with a trajectory length of $n=10$, and a trajectory length of $n=3$ in the more complex MOVi datasets.

To assess object-separability we created evaluation videos for all five datasets. In these evaluation sets only single objects from the respective object domain were changed while all other objects in the scene remained fixed. After training models on each of the datasets, we assessed how well one could linearly classify which object had moved using the protocol described in Section 3.2.

### 4.1 Object slots are not necessary for learning object dynamics

Figure: Figure 2: Prediction accuracies for slotted and non-slotted contrastive dynamics models. In all five datasets we see that the CWM is not only competitive, but sometimes outperforms the CSWM when it comes to predicting object dynamics. Scores are averaged over five seeds, with error bars depicting the standard error of the mean.
Refer to caption: extracted/5906950/figures/prediction_accuracy.png

Evaluating the prediction accuracy of the CSWM and CWM, we observe that object-slots are not necessary for accurately predicting object dynamics. In fact, the CWM models often outperformed their slotted counterparts (see Figure 2). As expected, we see test accuracy generally increase with training set size. This suggests that compositional generalization about objects, the ability to generalize about properties of objects in novel constellations and combinations, does not require explicit object-centric priors as provided by slotted architectures.

### 4.2 Predicting object dynamics improves object separability

Figure: Figure 3: Object decoding accuracy as a function of training set size, for contrastive models. CWM representations of objects become more linearly separable with dataset size, despite no architectural components that encourage the formation of object-centric representations. However, contrastive learning without next step prediction (CRL) does not give rise to object-centric representations, suggesting an important role for information provided by dynamic data. Scores are averaged over five seeds (three seeds in the MOVi domains), with error bars depicting standard error of the mean.
Refer to caption: extracted/5906950/figures/object_decodability_contrastive.png

Figure: Figure 4: Object decoding accuracy as a function of training set size, for auto-encoding models. The dynamic training scheme yields a monotonic increase in object separability with training set size in four out of five datasets. Scores are averaged over five seeds (three seeds in MOVi domains), with error bars depicting standard error of the mean.
Refer to caption: extracted/5906950/figures/object_decodability_autoencoder.png

If models without object-slots can successfully generalize about object dynamics in combinatorially novel scenarios, is this because they too develop separable and compositional representations of objects? We evaluated the degree to which CWM’s representations of objects were linearly separable. Here we observe that representations of objects get more and more separable with a linear decoder as the models are provided with more training examples. In simpler domains like the cubes and 3-body physics, the models attain scores close to a $100\%$ in the largest data setting (see Figure 3). In the more challenging domains like Multi-dSprites and the MOVi environments, where multiple objects are moving and interacting simultaneously, the decodability is lower, but substantially larger than chance at around $70\%$. For comparison, evaluating randomly initialized networks with the same metric only gives slightly better than chance object decodability scores, meaning that default representations for these models are strongly entangled (see Appendix C). Moreover, we see the same trend where larger training set sizes translate to better decodability. This suggests that, even for complex datasets with multiple interacting, realistically rendered objects, systematic and separable representations of objects can potentially emerge with scale.

To assess the importance of next-state prediction, we evaluate the object-separability of the CRL’s representations. Surprisingly, the CRL attains separability scores that are close to chance, suggesting that training on dynamic object data offers valuable information for learning composable representations in the contrastive setting.

Figure: Figure 5: Reconstructions and LPIPS similarity for different models on the MOVi (simple) and MOVi-A datasets. Auto-encoding models without object slots approach or match the reconstruction ability of Slot Attention on novel object configurations in the MOVi domain.
Refer to caption: x2.png

Do these trends hold for models with non-contrastive learning objectives? We evaluated the static and sequential variants of the auto-encoding models. First, we observe a significantly stronger tendency for the *sequential* auto-encoder to develop separable object representations (see Figure 4). This also suggests that providing the models with information about object dynamics in the form of a training signal can facilitate the development of composable object representations. In fact, it is only in the Multi-dSprites domain that the static autoencoder shows a monotonic increase in object separability with training set size. Comparing auto-encoding objectives to the contrastive objective, we see that object separability was generally lower for auto-encoding models in the cubes and 3-body physics datasets.

Next we assessed the reconstruction quality of the static and sequential auto-encoder, and compared them to Slot Attention. We used the LPIPS (Zhang et al., ) perceptual similarity metric to quantify reconstruction fidelity on novel object configurations in the MOVi-A and MOVi (simple) datasets. While we see that Slot Attention has a small edge on the distributed models in terms of fidelity, both the static and sequential auto-encoder approach Slot Attention with more data (see Figure 5). Lastly, the auto-encoder performs better than the sequential auto-encoder on the test set, despite attaining substantially lower object separability scores in the MOVi domain.

## 5 The benefits of partially entangled representations

Even though unsupervised training on images of objects leads to linearly decodable representations of objects, especially in the dynamic model class, the representations of objects do not ever become completely disjoint (see Figure 6). That is, the models rely on distributed codes in their latent spaces that often represent distinct objects using overlapping populations of neurons. Yet this does not seem to impact the models’ ability to perform compositional generalization, e.g. predict dynamics and reconstruct scenes of novel compositions of objects.

To get a better qualitative understanding for the degree of object separability, we analysed the trained CWM’s and their slotted counterparts’ representations and the degree to which they showed systematic similarities. We obtained object representations of $300$ initial frames $x_{t}$ and successor frames $x_{t+1}$ where only one object changed in one aspect from $t$ to $t+1$ in the cubes and Multi-dSprites domains. We chose these domains since they had actions that accompanied the dynamics. Earlier, we used the absolute difference between these representations $|\Delta|$ to get a sense of how objects were represented. However, one could also use these absolute differences to get a sense of how *transformations* or *actions* that acted on objects were represented. For instance, pushing the red cube along the $y$-axis on the grid might cause a similar change in representation as pushing the *blue* cube along the $y$-axis, even though the same action is applied to different objects. Analogously, in the Multi-dSprites domain, shrinking or rotating the heart sprite could induce similar representational changes as shrinking and rotating the square sprite. We do not expect to see this for slotted models, as they are more likely to represent the properties of different objects in orthogonal sub-spaces.

Figure: Figure 6: A: Representational similarity matrices showing cosine similarity of state transitions $|\Delta|$ for the CSWM and CWM on the cubes (left) and Multi-dSprites datasets (right). The cosine similarities are either ordered according to which object changed, or which action was performend on one the objects. In both cases, clusters are visible, though object clusters are more prominent in the slotted models, and action clusters more prominent in the distributed models. The cosine similarities are averaged over five seeds. B: CSWM intra-object similarity is significantly higher than its intra-action similarity, since object dynamics are isolated in separate subspaces. On the other hand, the CWM’s intra-action similarities are much closer to the intra-object similarities, allowing for richer generalization while preserving object separability.
Refer to caption: x3.png

For both model types and datasets we obtained the absolute representational change $|\Delta|$ for all $300$ frame pairs and computed pairwise similarity matrices using cosine similarity as our metric (see Figure 6A). These matrices contained information about which *transitions* were represented similarly for both distributed and slotted models. First we sorted these similarity matrices according to *which object* changed. Here we see five clear object clusters for both models in the cubes domain, and four clusters in the Multi-dSprites domain, albeit to a lesser degree for the distributed models.

Next, we sorted the similarity matrices according to *which transformation or action* was applied to the objects. While action clusters are identifiable for the CSWM, intra-action similarities were significantly lower than for the CWM (see Figure 6B). Representing object properties in a shared representational space not only allows for systematic representations of objects, but can also give rise to systematic representations of *transformations* that act on objects.

### 5.1 Distributed and slotted representational alignment increases with more data

Figure: Figure 7: Alignment with slot models as determinted by representational similarity alignment (RSA) (Kriegeskorte et al., ; Kornblith et al., ). The representations of the contrastive model (CWM) become more aligned with its’ slotted counterpart with more data.
Refer to caption: extracted/5906950/figures/rsa.png

It has recently been argued that deep neural network models’ representations tend to grow more and more similar as training data size and model sizes increase (Huh et al., ). Do we see a similar convergence in latent representations of scenes composed of objects?

We extracted representations from all five seeds and all dataset sizes for CWM and the Auto-encoding models in the cubes, 3-body physics and Multi-dSprites domains, and compared them to the corresponding representations of the object-centric CSWM models. Next we computed the Euclidean distance between all pairs of representations for all models in the three different domains, and then calculated the degree of correlation between these distance matrices of the different models. If two models’ distance matrices were highly positively correlated, they perceived the same pairs of images as similar and dissimilar. In other words, their representations show a high degree of *alignment*.

In all three domains we see that the CWM’s representations grow more and more aligned to those of the CSWM with more data, reaching a score of around $0.8$ on average in the largest data setting (see Figure 7). This indicates that, while alignment can increase substantially with enough data, simple architectural features can leave gaps, as for instance shown in Section 5. We did not observe this trend with the auto-encoder, which showed a significantly lower level of alignment with the CSWM. This suggests that, while data can drive alignment, this has to be paired with an appropriate loss function.

## 6 Discussion

### 6.1 Limitations

Our study has focused on unsupervised representation learning models, paired with static and dynamic prediction objectives. However, the space of unsupervised learning techniques is vast. Future work should investigate object-separability in *self-supervised* representation learning settings (Grill et al., ; Zbontar et al., ; Schwarzer et al., ). Moreover, other model architectures like Vision Transformers (Dosovitskiy et al., ) are promising, as their attention patterns have been shown to match segmentation masks of natural images of objects (Caron et al., ).

The degree to which these properties scale to naturalistic and real-world video datasets is unclear. A natural next step is to compare larger slotted architectures, such as VideoSAUR (Zadaianchuk et al., ), to distributed models on naturalistic videos. It is possible that in these complex, real world domains, relying on richer, more entangled representations can facilitate generalization. Lastly, regularization and information bottleneck methods may significantly aid in learning separable object representations as well (Alemi et al., ; Shamir et al., ).

### 6.2 Conclusion

We have shown that models without object slots can learn object representations that are disentangled enough to be *separable*, but entangled enough to support generalization about *transformations* of objects. Furthermore, training models to predict object dynamics significantly improved object separability. We believe our findings are important because they highlight multiple ways in which a representation can be beneficial for generalization: Slotted models can seamlessly decompose the world into its constituent objects, facilitating compositional generalization. Models with simpler, unconstrained latent spaces can decompose the world in ways that also separate objects, while allowing information about one object’s dynamics and properties to permeate to others.

## Acknowledgements

We thank the HCAI lab for valuable feedback and comments throughout the project. This work was supported by the Institute for Human-Centered AI at the Helmholtz Center for Computational Health, the Volkswagen Foundation, the Max Planck Society, the German Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 01IS18039A, and funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Germany’s Excellence Strategy–EXC2064/1–390727645.15/18.

## Appendix A Appendix

## Appendix B Architecture and hyperparameters

### B.1 Convolutional Neural Networks

For our image encoders we rely on a standard CNN architecture used in other works such as Yarats et al. (; ). For the 3-body physics dataset we stacked two frames, as in Kipf et al. () to provide information about directional velocity.

This network is followed by an MLP with $2$ hidden layers and $512$ hidden units. All models used ReLU (Nair & Hinton, ) activation functions and were optimized using Adam Kingma & Ba (). The CSWM and Slot Attention models were trained using the hyperparameters and encoder architectures provided in Kipf et al. () and Locatello et al. (), respectively.

### B.2 Transformer

We use the causal Transformer architecture of GPT-2 (Radford et al., ), building on the implementation in the nanoGPT repository(^1^11See https://github.com/karpathy/nanoGPT). In the slotted dynamics model, the transformer applied attention over the sequence of slots per time-step. Assuming $K$ slots and $T$ time-steps, the transformer applied attention over a sequence of $K\times T$ data-points. In MOVi (simple) we trained CSWM with $6$ object slots, allowing each slot to model one object plus a background slot. In MOVi-A we trained CSWM with $11$ object slots, allowing for $10$ separate object representations plus the background. The Slot Attention model was trained with the same number of slots as CSWM. The transformers were trained with the following hyperparameters:

**Table 1: Transformer hyperparameters.**
| Hyperparameter | Value |
| --- | --- |
| MLP Hidden units | 2048 |
| Transformer blocks | 2 |
| Context length | 4 |
| Heads | 8 |

### B.3 Hyperparameters

Below are specific hyperparameters for the different distributed model classes.

**Table 2: Contrastive model hyperparameters.**
| Hyperparameter | Value |
| --- | --- |
| Hidden units | 512 |
| Batch size | 512 (1024 for MOVi) |
| MLP hidden layers | 2 |
| Latent dimensions $|z_{t}|$ | 50 (500 for MOVi) |
| Margin $\lambda$ | 1 (100 for MOVi) |
| Learning rate | 0.001 (0.0004 for MOVi) |

For auto-encoding models we use the transpose of the encoder networks presented above. These models are trained with the following hyperparameters:

**Table 3: Auto-encoder hyperparameters.**
| Hyperparameter | Value |
| --- | --- |
| Hidden units | 512 |
| Batch size | 512 (64 for MOVi) |
| MLP hidden layers | 2 |
| Latent dimensions $|z_{t}|$ | 50 (500 for MOVi) |
| Learning rate | 0.001 (0.0004 for MOVi) |

**Table 4: Sequential auto-encoder hyperparameters.**
| Hyperparameter | Value |
| --- | --- |
| Hidden units | 512 |
| Batch size | 512 (124 for MOVi) |
| MLP hidden layers | 2 |
| Latent dimensions $|z_{t}|$ | 50 (500 for MOVi) |
| Learning rate | 0.001 (0.0004 for MOVi) |

## Appendix C Object decodability baseline

Figure: Figure 8: When initializing the image encoders randomly, the linear decodability of objects is higher than chance, but substantially lower than the level of trained models.
Refer to caption: extracted/5906950/figures/object_decodability_random.png

### C.1 Slotted representations are recoverable from distributed representations

We assessed the degree to which CWMs representations could be mapped to discrete object slots in the cubes dataset. We trained a slot decoder network with the same architecture as the Slot Attention model to reconstruct images from the representations of CWM. Assuming $K$ object slots, the decoder was trained to reconstruct the original frame as an additive composition of $K$ individual objects. Crucially, we froze CWM’s encoder, meaning that the decoder could only use information learned through the contrastive dynamic training to recompose the scene. We see that the slot decoder can not only learn to reconstruct the scenes with high fidelity, but also reconstruct the scenes by individually reconstructing the cubes and composing them.

Figure: Figure 9: Training a slot-decoder to reconstruct scenes from CWM’s learned representations leads to scene re-compositions that track the original ground truth objects.
Refer to caption: extracted/5906950/figures/contrastive2slots.png

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: Q1. What specific properties of the architecture or training procedure lead to linear separability effect?

Q2: How general are the results in the paper? Would they generalize to other object-centric approaches beside CWM / CSWM?
soundness: 3
strengths: -- Studying inductive biases for neural architecture is interesting. 
Its great to see a paper that is about more than just improving a few points on a known task. 

-- Interesting experiments.
summary: The paper studies object disentanglement in the representation of visual scenes. Specifically, object-centric representations enforce an inductive bias for disentangling representation of different objects.  The claims that (1) distributed representations (not object disentangled) may outperform object-centric  ones in down-stream tasks; (2) Limited object disentanglement (linear separability)  may arise even without forcing it through the architecture.

The paper compares the representations learning in structured and non-structured representations and analyzes their linear separability.
weaknesses: W1. There is a lot of work on compositional representations in machine vision. 
From early papers focused on compositionality of attributes 
 [Yuval Atzmon et al 2016, Justin Jhonson 2017]. 

W2. The main higlighted claim is that, ith sufficient data, models can learn representation that disentangle objects. 
I am missing an argument that explains why that is surprising or unsexepected. Does theory sugests otherwise? 
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. So why would disentaglement be a problem? 

W3. The paper measures linear separability. This is a different concept that compositionality, and the paper should make the distinction super clear and explicit, already in the title.  In the ML literature, compositionality usually means that one can represent new things using combination of concepts. 

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

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: For major questions and requests, see weaknesses.

Additional minor questions/requests:
- What are the actions that can be taken in the cubes and Multi-dSprites environments?
- Figure 6A: could the authors add labels of the objects/actions to the similarity matrices? It would be easier to interpret them knowing this information.
soundness: 2
strengths: Summarized points:
- Well written paper
- Perform an interesting and seemingly novel study
- Revisits and challenges an important basic assumption that incorporating explicit object-centric structure in representation-learning models is beneficial/necessary for representing object-centric scenes/dynamics
summary: This work performs a study comparing the learned representation spaces of unsupervised object-centric slot-based models and their unstructured counterparts trained on various datasets of dynamically interacting objects. The study focuses on the disentanglement of these representations in terms of objects and their capacity to capture underlying transformations that act on objects, regardless of their identity. The experiments include models trained on both reconstruction-based and contrastive objectives and test the learned representation’s efficacy in image reconstruction as well as multi-step dynamics prediction. Object separability is measured using a novel metric based on the ability of a linear classifier to recognize which object changed between two consecutive frame representations. Based on their experiments and analysis, the authors suggest that explicit object-level factorization is not necessary for learning these tasks and achieving compositional generalization, more so when the training objective includes next-state prediction. This is explained by the fact that the learned representations, although not object-centric, converge to being partially disentangled in terms of objects (as quantified by the object separability metric) while maintaining a level of entanglement that facilitates sharing of information about transformations that act on objects.
weaknesses: Summarized points:
- Some conclusions are not convincingly backed-up by experiments
- There is room for more systematic fine-grained evaluation of performance with respect to different types of compositional generalization
- The implications of this study to downstream task performance, a major use-case/motivation for unsupervised representation-learning, are not assessed in experiments and mostly not discussed

**Related Work**

The overview of object-centric representations is focused entirely on slot-based models. Although these might be the most dominant in the literature, I believe this section can benefit from including additional types such as patch-based (e.g. [SPACE](https://arxiv.org/abs/2001.02407), [MarioNette](https://arxiv.org/abs/2104.14553)) and keypoint-based (e.g. [Jakab et al.](https://arxiv.org/abs/1806.07823), [DLP](https://arxiv.org/abs/2205.15821)).

**Experiments - Model Forward Prediction Ability**

*Prediction Accuracy*: The evaluation metric described in the second paragraph of the Experiment Section is not clear to me. What does it mean for the prediction to be “closest to the last frame”, closest compared to what?

In the case of the object-centric dynamics model, how do you compute the distance between two latent representations that are not from the same trajectory (e.g. contrastive examples or encoded ground-truth future states)? The latents are sets of vectors and thus lack ordering such that a simple L2 loss would not necessarily compute distances between aligned slots.

I find it surprising that CWM outperforms CSWM in object dynamics prediction. The original CSWM paper seems to have results on the cubes dataset that contradict your results (see Table 1 -> 3D Blocks -> 10 Steps -> CSWM vs. -factored states). Do you have an explanation for this discrepancy? Is there an essential difference in the setting or metrics they used compared to yours? 

Another question that arises here is why did you choose to compare the prediction abilities in latent space? The latent spaces are not as comparable as reconstructions of images in my opinion since they possess different structure. I suggest training a “stop-gradient decoder” on reconstructing the contrastive latents so you can obtain image reconstructions without affecting the contrastive model’s optimization process. You can then use perceptual metrics to quantify prediction performance (e.g. LPIPS, which you used in the autoencoder study).

Why do the authors not study object dynamics prediction with the autoencoder models in addition to the contrastive models in section 4.1? This study could potentially strengthen your conclusions. In any case, I believe it should be part of your experiments.

**Experiments - Reconstruction Quality**

Figure 5 presents similarity metrics using LPIPS. For better qualitative assessment I would suggest adding additional metrics such as pixelwise MSE, SSIM, FID.

“the auto-encoder performs better than the sequential auto-encoder on the test set, despite attaining substantially lower object separability scores” - Is it possible that this is due to the fact that the auto-encoder was trained on that task precisely while the sequential auto-encoder was trained on next-frame prediction? The distribution of latent representations produced by the latent dynamics model is most likely different from latent representations directly encoded from images.

**Experiments - Compositional Generalization**

Can the authors clarify why high accuracy in the dynamics prediction tasks suggests compositional generalization capabilities? Is there a systematic separation between the train and test data in terms of compositions of objects and/or their properties such that performing well on the test data would necessarily require compositional generalization?
Put more simply: can you say for certain or with high probability that the test data does in fact contain novel constellations and combinations of objects?
If so, I request that the authors add details about the train-test split with respect to the relevant factors of variation in each dataset. 

The above questions are also relevant to the reconstruction experiments.

**Discussion**

As I see it, the two major claims about non-object-centric unsupervised representation learning models are:
- They can learn representations with some form of disentanglement between objects (as quantified by the proposed separability score)
- The learned representation space captures some notion of underlying transformations that act on objects, regardless of their identity

I believe the results that show that the object-separability score does not translate to improved task performance when comparing the static autoencoder with the dynamic one is indicative of a broader question that should be asked, discussed and maybe answered in the context of this work: *What are the implications of the two major findings in the paper?*

*I find this question not sufficiently answered*. I suggest looking into two main aspects:

1. *Downstream Task Performance*: Questions that should be answered in my opinion:
- Are the two qualities of the learned representation indicative of downstream task performance?
- Are they efficiently leverageable by downstream models?

I would argue that in this study, downstream task performance was not assessed. Both dynamics prediction and image reconstruction are exactly what the respective models were trained on.
Downstream tasks in the actionable environments could be sequential decision-making while on the uncontrollable dynamics, one could infer underlying properties of the scene or trajectory such as the number of objects that remain in the scene by the end of the trajectory in the MOVi-A dataset.

2. *Compositional Generalization*: the models’ ability to generalize its representations to scenes with novel compositions of objects as well as novel compositions of objects and the transformations that are applied to them.

Here I argue that this type of generalization was not systematically assessed.

In order to do so, I suggest first defining the factors of variation of interest.
Some examples:
- Number of objects in the scene
- Combinations of individual object static properties such as color, shape, size, mass, friction (most relevant to the MOVi-A dataset)
- Dynamic transformations in static object properties (most relevant to Multi-dSprites)

Based on these factors one should create a train-test split such that performance of the same task on the test data would be indicative of the model’s representation’s capacity to facilitate compositional generalization of that type.
Examples of train-test splits:
- Train includes up to x objects -> Test includes between x and x+y objects
- Train includes only red cylinders and blue cubes -> Test includes blue cylinders and red cubes
- Train includes sequences where all but circular objects dynamically change scale -> Test includes sequences where circular objects do change scale

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - For the dynamic loss (L_{AE-dynamic}), have the authors tried each term separately in addition to both combined? Does it change the conclusions for dynamic AE models?

- 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?

- Even though Fig. 6 seems to show that object identity is encoded in distributed models, is it possible to perform a linear to decode only object identity for the first or last frame of the videos? Slot models should be able to decode all the objects, but maybe the distributed models will only be able to decode the ones that are moving?

- In Fig. 2, 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?


- Fig. 7, 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.
soundness: 2
strengths: - The paper is overall clearly written and well presented. 

- The question tackled by the authors is very interesting and has not been extensively investigated in the past. 

- The method proposed to investigate the content of the latent representation by systematically controlling which object moves from one frame to the other and the nature of the movement is a very elegant way to test the information encoded in the latent representation. 

- The authors propose various ways to evaluate the content of the representations with different metrics, simple yet elegant and informative.

- The number of datasets and their variety in content and difficulty solidifies the authors' conclusions.
summary: This paper investigates the quality and content of the latent representation of different models trained using unsupervised learning on dynamic scenes composed of multiple objects. The authors compare slot-based models against distributed models and investigate their ability to encode object identity and object dynamics (movement in the scene). They find that distributed models have a better ability to encode identity *and* movement than slot-based models.
They also show that the models trained with more specific losses, such as contrastive loss and next-state prediction, help to improve the quality of these representations through linear readout.
weaknesses: The results part should be improved. Each experiment compares only two models (either in terms of objective or model type -- slot-based vs. distributed), but we miss a big picture comparing the baseline (the CSWM) with all the models with their characteristics (contrastive, next-state prediction, ...). 
It would be easier to read and draw conclusions if all the models appeared on all the bar charts, clearly highlighting the benefit of each objective depending on the task, assuming that all the models are comparable in terms of number of parameters and architectural design. 
Maybe this could be done for one of the datasets, and the other results can appear in the appendix. 
It would also be informative to add in the Appendix the complete architecture of all the models, along with their number of parameters. 
I will be willing to increase my rating if this part is improved with a figure showcasing the totality of the results, for all models on all metrics, backing the claims of the authors.

---

confidence: 5
contribution: 2
initial_rating: 1
presentation: 3
questions: In Eq.6 , the t should be replaced by t+1?
soundness: 1
strengths: The paper’s strengths lie in its idea of challenging the necessity of object slots, trying to show distributed models can effectively generalize with next-state prediction alone. It provides robust evidence across diverse datasets and offers new insights into how partial neural overlap supports transformation generalization, broadening representation learning approaches.
summary: This paper explores how distributed representation models can develop compositional, linearly separable object representations without object-centric architectures/inductive biases. Through next-state prediction, the authors claim that these models match or exceed the performance of slot-based models in predicting object dynamics, even without dedicated object slots. The study finds that partially overlapping neural codes in distributed models enable effective generalization and object separability, making them a viable alternative for tasks involving dynamic object interactions.
weaknesses: I want to remind the authors that, it is unfair to compare your model with small models designed for a special dataset, like CSWM, on new datasets such as multi-dsprites. 

Thus, The paper exhibits several weaknesses in supporting its claims:

1. The claims about distributed models achieving compositional representations are not robustly supported by the experiments. The CSWM baseline, a model specifically optimized for the datasets used in their studies, underperforms on the new datasets used in this paper. This performance drop is predictable and weakens the validity of the comparisons.

2.  The study relies on Slot Attention, which is designed for static images, as a baseline in dynamic scenes. For a fair comparison in dynamic settings, it should have considered more relevant recent work, such as Parallelized Spatiotemporal Binding, which better aligns with the dynamical nature of the datasets.

3. 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.

Overall, the paper lacks sufficient experimental depth, relevant benchmarks for dynamic contexts, and originality in its claims, suggesting it may not yet be ready for acceptance in its current form.

Meta-review: This paper studies structured (object-centric) vs. unstructured visual representations in the context of next-state prediction. While the paper is well written and tackles a very interesting problem, the reviewers voiced mixed opinions about the strength of the experimental evaluation of the claims made in the paper. The AC agrees with these concerns and encourages the authors to resubmit an improved version of the paper to a later venue.

Decision: Reject