GENESIS-V2: Inferring Unordered Object Representations without Iterative Refinement

Abstract: Advances in unsupervised learning of object-representations have culminated in the development of a broad range of methods for unsupervised object segmentation and interpretable object-centric scene generation. These methods, however, are limited to simulated and real-world datasets with limited visual complexity. Moreover, object representations are often inferred using RNNs which do not scale well to large images or iterative refinement which avoids imposing an unnatural ordering on objects in an image but requires the a priori initialisation of a fixed number of object representations. In contrast to established paradigms, this work proposes an embedding-based approach in which embeddings of pixels are clustered in a differentiable fashion using a stochastic stick-breaking process. Similar to iterative refinement, this clustering procedure also leads to randomly ordered object representations, but without the need of initialising a fixed number of clusters a priori. This is used to develop a new model, GENESIS-v2, which can infer a variable number of object representations without using RNNs or iterative refinement. We show that GENESIS-v2 performs strongly in comparison to recent baselines in terms of unsupervised image segmentation and object-centric scene generation on established synthetic datasets as well as more complex real-world datasets.


Reviews and author responses:

 Thanks for the comments.
I keep my original rating.

Additional notes:
     If you make the claim that KL regularization is the key to high segmentation quality, then you probably want to explain why the KL regularization in your baseline cannot achieve a similar effect.

 Thanks for the response. After reading the other reviews and responses, I still think this is a nice contribution and I'm keeping my score.

 Thanks for your response and clarifications. I have left my score as is.

 Thanks to the authors' efforts to address my concerns/questions in the rebuttal, and I keep my original rating of 6.

Additional notes:
- It is reasonable that the attention masks $m$ are unsegmented, and one needs a refinement step to create the final mask $\pi$. It would be informative to include the visual examples comparing $m$ vs. $\pi$ and related discussion in the revised manuscript.
- I agree that the non-"blob-like" objects are unrealistic, and it is the right direction to utilizes the inductive biases of locality and KL bottleneck, as the authors mentioned.

 We thank reviewer 9w6y for their time and their feedback.

Regarding the reviewer's comments about the segmentation of the background:
- GENESIS-v2 does not make any strong explicit assumptions about the size of objects. Indeed, it can be seen from the segmentation colour-maps (colour indicates "slot ID") for ShapeStacks, Sketchy, and APC that the background is explicitly segmented as a distinct scene component in the same fashion as the foreground objects. Results on the ObjectsRoom dataset are the only notable exception where the model always allocates the background to the last slot, corresponding to the "remainder" of the IC-SBP. While the combination of distance-based clustering and the semi-convolutional embeddings provides a weak inductive bias towards mask locality, especially at the beginning of training, this is not a "hard constraint".
- We share the sense that segmenting the entire background as a single component is not always the "correct" segmentation and that there are situations where it would be desirable to segment the floor, walls, sky, shelves, etc. separately. Particularly on ShapeStacks and APC, however, it can be seen that the baselines largely segment the background into regions of similar colour rather than meaningful entities. This also makes the segmentations sensitive to lighting conditions, reflections, etc.
- We believe that the behaviour of GENESIS-v2 arises from the fact that the background structures are of finite variety, i.e., even when viewing only a fraction of the background, it is possible to largely predict the rest of the background. We argue that the behaviour of GENESIS-v2 is the result of better conceptual compression which makes the model less sensitive to pixel colour and enables it to better model the scene as a whole. This is induced by the KL regularisation during training which encourages efficient compression and penalises redundant information between slots.
- We appreciate the reviewer's comment that this is not sufficiently discussed in the text and we will update the paper accordingly.


 We thank reviewer mGsC for their time as well as the thoughtful and constructive comments.

Regarding individual comments/questions:
1. We also believe that time and computational resources are put to better use by pushing models like this to work on more complex datasets rather than getting them to work on Multi-dSprites and CLEVR which are largely "solved" by MONet/IODINE/SlotAttention.
2. The same GECO reconstruction target was used as in the GENESIS paper for all datasets apart from Sketchy where a smaller value was empirically found to work better. To clarify, a fixed sigma is used and GECO "anneals" beta in the loss function. Similar results can possibly be obtained by hand-tuning sigma/beta but we found GECO to be a more intuitive method for performing this type of tuning. We did not try learning sigma or making sigma input-dependent, which might be interesting to explore in the future.
3. We have indeed thought about different likelihoods. For example, Gated PixelCNN [1] uses a softmax with 256 outputs for each colour channel. This is more flexible than a Gaussian likelihood, but also computationally more expensive as 3*256 values need to be predicted at each pixel location. It also does not capture the intuition that RGB values that are close to each other are also visually similar. An improvement is proposed in PixelCNN++ [2] by using a "discretised logistic mixture likelihood". We did not conduct experiments with either of these likelihoods though. The Spatial Gaussian Mixture Model formulation as developed in previous works appears to provide a very good inductive bias for object decomposition and it is not clear whether other formulations would work similarly well or better.
4. Thanks for spotting; will fix.
5. Thanks for spotting; will fix.
6. Thanks for spotting. We are using a "maximally distinct" colour palette but will investigate whether this can be further improved.
7. If a user is only interested in predicting segmentation masks and does not care about latent representations or reconstructions, then they might prefer to only run the UNet to directly predict the segmentation masks. This would allow them to "skip" the object decoders and save compute. We will update the text and make this explicit.
8. The work on IODINE [3] used 320 images for computing segmentation metrics (equal to ten batches of size 32 or five batches of size 64) and we chose to follow their methodology.
9. We believe the authors of GENESIS meant to write that the foreground-based ARI does not penalise the 'under-segmentation' of objects, i.e., there is no penalty if an object segmentation mask covers both the object and part of the background. This might be a typo in the GENESIS paper.

[1] "Conditional Image Generation with PixelCNN Decoders"\
[2] "PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture Likelihood and Other Modifications"\
[3] "Multi-Object Representation Learning with Iterative Variational Inference"

 We would like to thank reviewer jTRB for their time and the detailed review.
- We appreciate the suggestion to make the description of the new methods "kinder for newcomers" and as well as the suggestions for extending the discussion of related works. Both will be taken into consideration for the next version of the paper.
- Shape of the attention masks: Yes, ideally the attention masks should be identical to the segmentation masks. In practice, we find that they are clearly correlated to the position of the objects and therefore the segmentation masks, but they often undersegment the object and overlap with parts of the background. This is not too surprising considering that the receptive field for each element in the feature map produced by the UNet backbone is the whole image, so information relating to a certain object can easily be "dilated" to other locations in the feature map.
- "Is it fundamentally possible to discover objects from a static image?": The results in this paper as well as other works indicate that static images can be segmented into interpretable entities without supervision by incorporating appropriate structural knowledge and inductive biases. Locality is one such inductive bias which is used in GENESIS-v2. While current methods are still limited in the visual complexity of the images that can be handled, this boundary is being shifted continuously. In general, what should constitute an "object" is arguably not uniquely defined and depends on context (or task information). For example, is a tree a single object or are all the leaves individual objects? Task or temporal information are without doubt useful for discovering objects and we believe that advances in discovering objects in static images should translate to other settings where more information is available. It would indeed be interesting to investigate how GENESIS-v2 performs on less "blob-like" objects as suggested by the reviewer. Objects with more complex topologies, though, are arguably not encountered as often in practical applications. The KL regularisation is a second important inductive bias by penalising information redundancy between object slots. Given the finite variation in terms of background appearances (e.g. the square patterns in Fig. 2), it makes sense that the whole wall is segmented as one object rather than multiple objects, as it should be possible to largely predict what the rest of the wall looks like even when "seeing" only a part of it. This goes somewhat against the typical properties of "objects" which are often to some degree independent of their surroundings: for example, if an entire block in a ShapeStacks tower is masked out, then it is impossible to tell what colour the block is.
- Future directions: We agree that both of these suggestions are promising directions to pursue in future work. We will mention them in the next version of the paper.

 We thank reviewer nSnK for their time and the thoughtful comments. Responses to individual points are below.
- Weaknesses
    - Iterative refinement (e.g. SlotAttention/IODINE): While the number of iterations in these models is independent of the number of objects, the computational cost of each iteration increases linearly with the number of instantiated object slots. The IC-SBP offers computational savings by allowing for a flexible number of slots, i.e., not requiring the number of object slots to be over-specified with some of them remaining "empty".
    - Oversegmentation of background: We expect that GENESIS-v2 would outperform all the baselines when computing the ARI/MSC for all image pixels rather than only for pixels belonging to foreground objects --- if the whole background is labelled as a single scene component. However, it is debatable what constitutes the "correct" segmentation of the background. For example, one could argue that rather than treating the background as a single scene component, it might be desirable to segment the walls/floor/sky separately. While either is more desirable than segmenting the background into non-meaningful slots, the more fine-grained labels required for such an evaluation are not available.
    - Stopping heuristic/number of slots: Comparing the number of used slots to the true number of objects is an interesting suggestion which we will investigate.
- Minor comments
    - "Last slot tends to be the background": Both behaviours described by the reviewer can be observed. On ShapeStacks, attention and segmentation masks are fairly similar, so the application of the IC-SBP to the attention masks means that the first slot often corresponds to the background as suggested by the reviewer. On ObjectsRoom, however, the attention masks belonging to the foreground objects are often  "blurred out" across large parts of the image. As a result, when selecting a pixel at random the model "hooks" onto one of the foreground objects until only the background is left.
    - Figure 2/Table 2: Nine object slots were used for ShapeStacks, but many slots remain empty and can therefore not be seen in Figure 2.
    - KL term/ordering of posterior: As can be seen by the colour-coding of the segmentation masks, the inferred posterior of the object representations is not ordered (except on ObjectsRoom in Figure 12 where the last slot always corresponds to the background).

This paper proposes a new object-centric generative model and demonstrates its effectiveness on a number of datasets (notably including 2 real-world datasets, unlike most previous work) compared to 3 SOTA baselines, both in terms of unsupervised segmentation and image generation. The model is reminiscent of GENESIS in terms of generative model and Slot Attention in its approach to approximate inference. However, thanks to a novel differentiable clustering method based on a stick-breaking process, this model does away with recurrent attention mechanisms and iterative refinement methods common in previous work.  The proposed method is novel, interesting, significant for the community, and proven to be effective by a sufficiently extensive experimental section. Weaknesses of the model are clearly shown. Overall this is a solid contribution to the field.

The paper is generally very well-written. However, the paragraph at lines 157-171 is not as easy to follow, and the points therein could be made clearer. Similarly, the approach outlined in Sec. 3.2, based on kernel distance and adding spatial information, is a nice way of including an inductive bias, but the argument with the "delta vector" is not very clear.

Other comments/questions:
1. Since GENESIS does not seem perform particularly well on standard datasets like multi-dSprites and CLEVR, it would have been interesting to see if v2 closes this gap with other methods like MONet and slot attention. On the other hand, the additional compute required to run these experiments was arguably put to better use by the authors by testing the models on 2 real-world datasets.
1. I assume the reconstruction target for GECO is chosen empirically, starting from an educated guess of possible good values (e.g. taken from GENESIS which is also trained with GECO). Is this "annealing schedule" necessary, or could one empirically choose a reasonable sigma for the likelihood components? Have you tried learning sigma per-component or even per-pixel (probably wouldn't work at first, but maybe after an initial training phase with fixed sigma)?
1. Have you considered or tried different likelihoods? Since these datasets are probably discrete (uint8), the right thing to do would be to use a discrete distribution (although it might not matter in practice unless we want to measure ELBO or likelihood) with support [0,1]. Using Gaussians makes sense considering prior work, but I'm curious if you have insights on this.
1. lines 25-26: many models for object-centric representation learning are in fact not generative, but they are cited as generative here
1. line 105: not all cited works use a SGMM as likelihood, e.g. slot attention uses MSE as loss, which is equivalent to using a Gaussian log likelihood where the mean is a linear (convex) combination of pixel values from all slot reconstructions.
1. the colors of segmentation mask are sometimes hard to tell apart (e.g. Fig 12 in appendix)
1. lines 208-209: which applications are you referring to?
1. why measure segmentation metrics on 320 images only? (also, is there any reason for this particular number?)
1. line 237: as far as I remember, ARI does not penalize _over_-segmentation, i.e. when an object is segmented in multiple clusters, as much as MSC. Could you double-check this? I think this is also mentioned in the GENESIS paper. The societal impact is addressed.

The paper proposes an object-centric generative model that combines the best of previous approaches that use RNNs or iterative refinement. The model can infer a variable number of object slots at test time as in RNNs, and the inferred object slots are randomly ordered as in iterative refinement. This is achieved by a novel differentiable iterative clustering algorithm. At each iteration, a random pixel (not yet assigned to any cluster) is chosen as the cluster seed, and a new cluster is formed by computing similarity between the per-pixel embeddings and the seed embedding. At test time, a stopping heuristic can be applied so that the number of clusters is flexible. On two simulated datasets, the model achieves better segmentation and generation quality than MONET and GENESIS, and comparable or slightly worse segmentation quality than Slot Attention (which cannot do generation). The model also shows promising results on two real-world datasets. In particular, on the APC dataset, the model shows significantly better segmentation quality than all the baselines.  **STRENGTHS**

- The model unifies two desirable properties of object-centric representations: flexible number and unordered slots.
- The method is novel. It is inspired by supervised instance segmentation literature, with improvements that made it suitable for unsupervised, end-to-end segmentation.
- The results on real-world datasets look promising, especially on the APC dataset where previous methods completely failed.

**WEAKNESSES**

- While it is conceptually desirable to be able to infer a flexible number of object slots, it is not clear from the paper what computational benefits this can bring. Specifically,
    - In iterative refinement, the number of iterations is usually small and independent of the number of objects. However, in the proposed model, the number of iterations increases with the number of objects, which can be inefficient when there are many objects.
    - Using a flexible number of slots only showed success on one of the datasets, and it required additional loss to match the performance when using a fixed number of slots.
    - Perhaps one advantage of using a flexible number of slots is that it can reduce oversegmentation, which seems to be the case in the qualitative results. However, this is not supported by quantitative evaluation. Maybe the authors could include background ARI, and see whether it is improved when using a flexible number of slots.
    - It is shown in Table 2 that using the stopping heuristic can reduce the number of slots. However, it is not clear how close this is to the true number of objects.

**MINOR COMMENTS**

- It is mentioned (L175-176) that the last slot tends to be the background. I am curious how this emerged. In the datasets considered, most pixels seem to belong to the background, so when a cluster seed is randomly sampled, it should have high probability to be a background cluster. This seems to suggest the opposite, that the first few slots are more likely to be the background.
- How many slots are used in Figure 2? Table 2 suggests that you used 9 slots for this dataset, but the number of clusters seems to be far less than 9 in Figure 2.
- The paper could be strengthened by showing generalization to more objects at test time, like what was shown in Slot Attention. This may also demonstrate the benfit of having a flexible number of slots.
- Does the KL term introduce some implicit ordering to the posterior? The ability to predict a flexible number of object slots at test time seems limited, and the model may have scalability issues when the scene contains many objects.

Object-centric generative models aim to decompose a scene image $x$ into multiple object-wise (and background) latent $z_{1:K}$. To extract the *set* of latent vectors, prior work considered a fixed-order recurrent approach (e.g., GENESIS) or an iterative refinement approach (e.g., IODINE, DSPN, SetTransformer, and Slot Attention). However, the former may impose an unnatural ordering and suffer from the forgetting issue for long sequences, while the latter requires a predefined fixed set cardinality. Instead, the paper suggests iteratively cluster the local pixels to form an object, which does not suffer from the ordering issue (except the last background) and can naturally decide the set cardinality (by a stopping heuristic). The proposed method produces more visually plausible scene decomposition results in standard synthetic and newly demonstrated real (e.g., robot arm) datasets.  The paper proposes a new idea for object-centric generative models and outperforms prior work, including more challenging robot arm and Amazon picking datasets. This would be a nice contribution stepping towards more realistic scenarios, and thus I recommend weak accept. However, I have some concerns and questions, which I expect to be clarified in the revised paper.


**Basically a local clustering method**

At first glance, it was hard to catch the core idea. After reading the paper carefully, I found that the proposed method is pretty simple (well described in Algorithm 1). Specifically, it repeats:
1. Choose a random location (i,j) from the (soft) remaining area.
2. Cluster the pixels which are similar to the one of (i,j).
3. (Softly) remove the chosen cluster and goto 1.

Despite its simplicity, I think the current description in Section 3 is a little confusing. For example, when I was reading lines 142-144, I could not catch what is the scope $s$ and seed score $c$ because the technical details for soft clustering assignment (i.e., scope) and random permutation (i.e., seed score) suddenly appeared before I understand the entire picture. Also, I was curious why one should use the semi-convolutional pixel embeddings in line 125, but found that it is just a positional encoding after reading the original paper. The paper could be kinder for newcomers by (a) presenting the high-level idea first, then explain the technical details, and (b) making contents self-contained. I think the description could be more concise, and there would be enough space for being self-contained.


**How does the attention masks $m$ shaped?**

Ideally, the attention masks $m$ should be identical to the final mask $\pi$, and there is no reason to generate $\pi$ again from the latent $z$ (may with some thresholding trick if needed). What is the reality? The auxiliary mask loss results on ObjectsRoom give negative signals. It would be interesting to analyze how the attention masks $m$ changes during training; check if it converges to the one-hot mask as training goes.


**Discussion on steak-breaking processes**

The paper named its method as an instance colouring stick-breaking process. Aside from being a *stick-breaking process* formally (e.g., the remaining background is not symmetric), the paper could discuss the related literature that inspired this work, e.g., how finite realization of the Dirichlet process can be implemented. Also, there have been some trials to define a stochastic process prior to VAEs [1,2], although they do not consider the spatial structure of latent.

[1] Nalisnick and Smyth. Stick-Breaking Variational Autoencoders. ICLR 2017.\
[2] Joo et al. Dirichlet Variational Autoencoder. arXiv 2019.


**Is it fundamentally possible to discover objects from a static image?**

The proposed method heavily relies on the heuristic of the *locality* (by distance kernel) and *large size* (by stopping criteria) of the objects. As a result, it shows more visually plausible results on the demonstrated benchmark datasets. However, it may not work well for *non-Gaussian-like* objects. Also, is it possible to distinguish the pattern of background and objects (e.g., square patterns in Figure 2)? I suspect that the current success of GENESIS-v2 comes from the dataset assumptions: (a) objects are Gaussian-like, and (b) backgrounds are relatively simple. It would be interesting to check the results on more complex objects.


**Possible future directions**

While it is out-of-the-scope of the paper, there would be several directions to improve the proposed model. First, to address the fundamental limitations of object discovery, the model would utilize the additional information, e.g., sequential movements of objects like SQAIR. Second, recent work on VAEs found that hierarchical structure is the key to high-resolution image generation (e.g., [3]). It would be interesting to combine current object-centric structures with hierarchical VAEs.

[3] Vahdat and Kautz. NVAE: A Deep Hierarchical Variational Autoencoder. NeurIPS 2020. The paper discussed its limitation, e.g., in lines 250-255. The paper also discussed the potential negative societal impacts but in Appendix. Following the suggestion of NeurIPS style guide, the revised paper could contain the explicit limitations and societal impacts section in the main text.

Paper proposed a novel unsupervised pixel-level clustering algorithm that detects meaningful object representations on multiple datasets.
The proposed model also supports novel data sampling.  Originality:
While GENESIS-V2 shares similar components with GENESIS, the pixel-level clustering mechanism is of great novelty.

Clarity:
The Paper is well written. I found it easy to follow.

Quality:

The generative model proposed is sound.
Evaluation on multiple datasets (including real-world datasets) is provided and the proposed model shows superior performance on all of them.
Images sampled from the proposed model are of high quality.

I do have one major concern about the proposed model.
GENESIS-V2 has an inductive bias that ignores background structures.
Do the kernel functions impose a shape prior or implicit limits on the size of the objects?
The author, in the APC case, set "object" to be the one individual small item in images, excluding any background structures.
Actually, outstanding background structures in the APC dataset are ignored by GENESIS-V2 but picked out by all other baselines.
In Fig 12, floors, walls, and sky are segmented as one single object by GENESIS-V2.
According to my memory, almost all previous works on the Object-Room dataset treat those structures separately.

Please correct me if I am wrong but I don't think this inductive bias is mentioned in the main paper.
I am not against such inductive bias at all and actually find it rather interesting.
However, IMHO, a more detailed discussion on such strong inductive bias should be added.

Significance:
After the physical scene graph, GENESIS-V2 is another strong evidence showing that the clustering-based algorithm has its unique potential and deserves more attention from the community.
 Societal impact is not addressed and there is not much to address I think.

Meta-review: - The proposed method is tackling an important problem. The reviewers found the approach is reasonable and has some novelty. Demonstration of the method on the real world dataset is a step forward in the line of research. 
- The major concerns from the reviewers are addressed well enough by the rebuttal. 
- The clarity is fair enough but some clarification as pointed by the reviewers would make the paper better. 

Decision: accept