Identifiable Object-Centric Representation Learning via Probabilistic Slot Attention

We thank the reviewer for their detailed comments and constructive feedback. We appreciate the fact that our paper was considered to be clearly written and well presented, and we are glad that our results are perceived to be solid.

"However, the paper does not really have experimental results supporting this claim..."

We respectfully disagree that the results do not support this claim, PSA-NoA and PSA-Proj both include the probabilistic latent setting as proposed in our work. Notably, there is a known trade-off between identifiability and expressivity induced by the choice of decoder structure [45]. Depending on the use case, it may be beneficial to combine both latent and additive decoder structures in practice, particularly if the latter introduces useful inductive biases and/or simplifies the optimization problem. In our experiments, we observed that when using PSA in tandem with an additive decoder it is possible to outperform all other baselines. As for non-additive decoder-based experiments, PSA-NoA must be compared directly against SA-NoA for it be a fair comparison, and we observed it to perform better than SA-NoA while remaining competitive with the remaining baselines.

"Can the authors apply PSA to a better non-additive decoder? For example, the Transformer decoder proposed in [61]. Otherwise, it is hard to assess the contribution of this work."

We agree that the applicability of our framework on large-scale datasets is a crucial evaluation of our theoretical results, but that does not reduce the contribution of our work, which is primarily theoretical. The main focus of this work is to investigate the theoretical identifiability of slot representations and the conditions that ensure this property, rather than provide state-of-the-art results on large-scale datasets. To verify our theoretical claims, we first conduct detailed experiments on controlled datasets, and then extend our demonstrations to unstructured image data. We stress that the synthetic datasets we used are necessary for properly testing our identifiability hypotheses.

One of the main assumptions necessary to prove slot identifiability in our setting is weak injectivity, which is achieved when we use piece-wise linear decoders. In the case of transformer decoders, this assumption is not guaranteed to hold because of the complexity of the attention mechanism (further theory is required here). With that said, we have conducted additional empirical evaluations on the identifiability of slot representations obtained with more complex transformer decoders, which result in an SMCC of $\mathbf{0.73 \pm 0.04}$ and R2 of $\mathbf{0.55 \pm 0.06}$ on the CLEVR dataset, which is significantly better than SA and all other baselines. For more details and additional experiments please see the general comment at the top.

"What is the implementation detail of the non-additive decoder "standard convolutional decoder"? How to decode an image from K 1D slots using a CNN decoder?"

Thank you for pointing this out as missing as it escaped our attention. We simply concatenate all K slots together and upscale the resolution using four transposed convolutions until we reach the image resolution, we also use Leaky-ReLU activations. The architecture is very similar to the one used in the original SA work, we will be sure to add the remaining details to the paper, thanks again.

"How can we apply the theory proposed in this paper to more complex datasets, e.g., real-world COCO images? This is not required for a theory paper, but I'm curious about the authors' thoughts."

We have included experiments on PascalVOC2012 datasets and also tested our model with more complex decoders, please refer to the general comment above for details and additional results.

I thank the author for the rebuttal. The experiments on the Transformer-based slot decoder and the large-scale experiments on real-world datasets are extensive and strong. My main concerns are addressed. Now I recommend acceptance of the paper and have adjusted my score to Weak Accept.

We thank the reviewer for their very detailed comments and constructive feedback which helped improve our paper significantly! We also appreciate the fact that our paper was perceived to be well-written, easy to understand, and of interest to the community.

"One of the main issues I have with this work is that I do not think that the paper’s storyline ..."

We thank the reviewer for their very insightful and constructive suggestion regarding the messaging of the paper. We largely agree with the overall sentiment and will incorporate the suggested changes which will significantly improve the paper’s pitch. To clarify, it is not our intention to compete with or replace additive decoders as a model choice since they undoubtedly provide useful inductive biases for object-centric learning as ours and many other previous works show.

We further remark that before our work there was a lack of explanatory theory for why state-of-the-art results were able to be obtained using non-additive autoregressive Transformers (DINOSAUR [60]) and/or diffusion-based decoders (Slot-Diffusion [r2]). We showed that by viewing slot attention through a probabilistic graphical modelling perspective it is possible to prove slot identifiability for non-additive decoders using proof techniques from identifiable generative modelling. Given that vanilla slot attention can be seen as a simplified version of probabilistic slot attention, akin to the relationship between soft k-means and GMMs, our theoretical results suggest why non-additive decoder structures can work well given the appropriate latent structure and inference procedure are in place. With that said, there is an identifiability and expressivity trade-off [45] induced by the choice of decoder structure, so depending on the use case it may indeed be advantageous to combine both latent and additive decoder structure!

"I would have appreciated a short paragraph giving some intuition on how exactly the probabilistic ..."

We would like to point out that we do discuss the intuition of our identifiability results in depth in the appendix, just before the proofs. As opposed to [4, 5], we use a more expressive GMM latent distribution followed by a weakly injective piece-wise linear decoder which in combination ensures that the slot representations are identifiable.

"The authors use two main metrics to validate identifiability: the slot identifiability score (SIS) from ..."

SIS is the relative R2 (coefficient of determination) score-based measure, capturing the relative ratios of variance for a dependent variable explained by an independent variable with baseline model scores that are trained on the slot representations of a dynamically updating model. As for SMCC, it measures the linear/affine relationship between permuted slots and features. We provide more detailed explanations in appendix F.

"The authors experiments on image data focus primarily on simple decoders opposed ..."

Although we agree that assessing the performance of probabilistic slot attention on more complex datasets would be useful, the controlled scenarios and datasets we used are necessary for properly testing our theoretical identifiability hypotheses, which is the objective of this work.

Regarding scalability concerns of probabilistic slot attention (PSA), we stress that probabilistic slot attention PSA retains the $\mathcal{O}(TNKD)$ computational complexity of vanilla slot attention and as such enjoys the same scalability properties. We did evaluate our model on large-scale datasets and with more complex decoders, we discuss them in the general comment.

"Do the authors view their method as being a replacement for structured decoders in object-centric learning or do they view the method as being better used in tandem with structured decoders?"

From a theoretical identifiability viewpoint, latent distributional structure is a sufficient requirement as long as the decoder is piece-wise linear. However, additive decoders provide strong, useful inductive biases and may be easier to optimize relative to a probabilistic model with an arbitrary decoder. In our experiments, we observed that using our approach in tandem with an additive decoder tends to outperform other models.

"Do the authors have an explanation for the low SIS score on the toy data experiments?"

SIS was proposed by [1] and it uses the relative R2 scores, where baseline models are trained on the slot representations of a dynamically updating model. Due to this, it tends to exhibit high variability as illustrated in Appendix F (we use the official implementation on GitHub for all our analysis). We believe this instability in the baseline model, which is also pointed out by others on their GitHub repository, could be due to the validation dataset size resulting in lower scores.

"Does “R2” refer to SIS when used or is this a different score?"

R2 is the coefficient of determination, while SIS is a relative R2 score. R2 is propositional to SIS but due to the instability of SIS, we use R2 for all imaging experiments.

"Do the authors have any intuition about how well their method would perform for more complex models/datasets?"

The method can be scaled to large-scale datasets using the approaches outlined in [60], with convolutional or transformer-based decoder since the computational complexity is the same as vanilla slot attention ($\mathcal{O}(TNKD)$). Please refer to the general comment for results on complex datasets using more powerful decoders.

[r2] Wu, Z., Hu, J., Lu, W., Gilitschenski, I. and Garg, A., 2023. Slotdiffusion: Object-centric generative modeling with diffusion models. Advances in Neural Information Processing Systems, 36, pp.50932-50958.

I thank the authors for their reply.

$\newline$

"One of the main issues I have with this work is that I do not think that the paper’s storyline ..."

I appreciate the authors acknowledgement of my concerns and look forward to reading an updated version incorporating this feedback.

Regarding the connection between slot attention and Transformer’s success in object-centric learning:

If this is the messaging of the paper that the authors wish to convey, then I would encourage them to rewrite the introduction accordingly. This messaging was not clear to me from reading the paper. Moreover, I am not completely convinced by this explanation given that decoder structure does indeed play an important role empirically. Thus, I find it more likely that the success of Transformers is due to a combination of probabilistic and decoder structure (i.e. the inductive biases of the Transformer). In other words, I do not think the authors results on their own provide a complete explanation of Transformer's success in object-centric learning but may provide some evidence to this end.

$\newline$

I would have appreciated a short paragraph giving some intuition on how exactly the probabilistic …

Thank you for the paragraph reference. Based on this, I would like to make a related point on the clarity of the theory. Namely, I do not think that the authors sufficiently highlight the piece-wise linear structure needed on the decoder for their theoretical result.

Up until page 7, the paper gives the idea that no decoder structure is needed for the theory. For example, the authors state in Section 2 that their theoretical contribution falls into the category of “(iii) imposing structure in the latent space via distributional assumptions.” This is not exactly correct, as the piece-wise linear structure is indeed decoder structure. I understand that the authors may view this structure as more easily implemented than e.g. additivity, and thus possibly less noteworthy. However, this piecewise linear structure is a key aspect of the theoretical contribution of this paper, and moreover, is important for contextualizing the authors theoretical contribution relative to prior identifiability results. Thus, I think this structure should be mentioned a bit more transparently in the introduction, and discussed with more precision in the related work. Otherwise, the messaging of the paper once again feels misleading.

$\newline$

"The authors use two main metrics to validate identifiability: the slot identifiability score (SIS) from ..."

I appreciate the authors reply on this point. I am curious, however, what the procedure was to resolve the matching problem when computing SMCC. Specifically, slots are locally permutation invariant opposed to globally due to the local permutation invariance of the decoder. Thus, how did the authors resolve this local permutation invariance when computing SMCC?

We thank the reviewer for engaging with us and for providing feedback.

“If this is the messaging of the paper that the authors wish to convey, then I would encourage them to rewrite the introduction accordingly…”

We are glad to hear that and look forward to incorporating the suggested changes, thanks again.

To clarify, our discussion of slot attention and Transformers is in response to the reviewers' requests and does not represent a change in the core message of the paper. We also acknowledge that further theory is required when using non-additive Transformer-based decoders as technically the weak injectivity property our proofs rely upon is not known/guaranteed to hold for Transformers due to the complexity of the attention mechanism (c.f. response to Reviewer HxYt). We believe that extending our theoretical results by relaxing the weak injectivity decoder assumption offers a promising direction for future research.

“I find it more likely that the success of Transformers is due to a combination of probabilistic and decoder structure (i.e. the inductive biases of the Transformer).”

We generally agree as there is a known trade-off between identifiability and expressivity induced by the choice of decoder structure [45]. As such it may be beneficial to combine both latent and decoder structures, particularly if the latter introduces useful inductive biases and/or simplifies the optimization problem. In our experiments, we observe that the combination of both typically yields better results.

“I do not think that the authors sufficiently highlight the piece-wise linear structure needed on the decoder for their theoretical result ”

We understand the reviewer's concern but respectfully disagree with their conclusion. Our theoretical results show that the decoder additivity constraint is not required if the decoder is piecewise linear and the latent space is GMM distributed. Although any decoder possesses some structure in terms of its architecture, an MLP decoder with LeakyReLU activations (satisfying weak injectivity) does not impose structure in the same sense as an additive MLP decoder, as the latter is a stronger restriction on the functional class and departs from the standard MLPs commonly used outside of object-centric learning. We will emphasize the weak injectivity assumption and the implications of piecewise decoders earlier in the introduction.

“How did the authors resolve this local permutation invariance when computing SMCC?”

As detailed in Appendix F, we used Hungarian matching to resolve this when computing SMCC.

Additionally, we believe our new experiments on PascalVOC address the reviewer’s concerns regarding the scalability of probabilistic slot attention.

Thank you for the reply!

I would encourage the authors to include these details in the appendix opposed to just a code implementation. As far as I know, matching based on Euclidean distance is non-standard. In future iterations of this work, the authors can consider including experiments to test the effectiveness of this matching protocol compared to other methods such as matching based on slot-wise mask, or matching based on R2 score (determined in an online fashion).

We thank the reviewer for their thoughtful, detailed comments and constructive feedback. We greatly appreciate the positive outlook and the fact that our work was found to be well-written, well-motivated and novel.

"The paper states that the framework allows for tractably sampling from the aggregate posterior ..."

We indeed show that it is possible (Lemma 1) to use the aggregate posterior for compositional tasks, but we felt that it does not meaningfully add to the paper's main theoretical identifiability contributions. We have now included some preliminary results for demonstration purposes (see pdf). Performing a more comprehensive compositional analysis of the aggregate posterior is certainly valuable and warrants dedicated investigation in future work.

"Experiments only validate the theory on simple synthetic datasets. Testing on more diverse and realistic ..."

As rightly pointed out by the reviewer, this work aims to study theoretical identifiability of slot representations and the conditions that ensure this property, rather than provide state-of-the-art results on large-scale datasets. To verify our theoretical claims, we first conduct detailed experiments on controlled datasets and then extend our demonstrations to unstructured image data. We stress that the synthetic datasets we used are necessary for properly testing our identifiability hypotheses.

Regarding scalability concerns of probabilistic slot attention (PSA), we emphasise that PSA retains the $\mathcal{O}(TNKD)$ computational complexity of vanilla slot attention, where $T$ denotes the number of attention iterations, $N$ the number of input vectors, $K$ the number of slots and $D$ the slot/input dimension. The additional operations we introduce for calculating slot mixing coefficients and slot variances (under diagonal slot covariance structure) have complexities of $\mathcal{O}(NK)$ and $\mathcal{O}(NKD)$ respectively, which do not alter the dominant term. Furthermore, when used in conjunction with additive decoder-based models, PSA can reduce computational complexity by pruning inactive slots via automatic relevance determination (ARD) as outlined in Section 4.

Finally, we have now demonstrated the applicability of PSA to transformer-based decoders - please refer to the general comment above for details.

"Table 1 lists -disentanglement and weak injectivity as core assumptions. While the former ..."

We have included a discussion on the weak injectivity assumption in the remark just below it - we’ll also include a similar discussion in the main text. In summary, weak injectivity ensures that a mixing function $f_d$: (i) in a small neighbourhood around a specific point $x_0 \in \mathcal{X}$ is injective – meaning each point in this neighbourhood maps to exactly one point in the latent space $\mathcal{Z}$; and (ii) while $f_d$ may not be globally injective, the set of points in $\mathcal{X}$ that map back to an infinite number of points in $\mathcal{Z}$ (non-injective points) is almost non-existent in terms of the Lebesgue measure on the image of $\mathcal{Z}$ under $f_d$. This assumption is generally satisfied when using Leaky-ReLU networks with randomly initialized weights (Appendix C).

"In fig. 3, it is not clear what experiment the latents (x-axis) correspond to."

Our apologies for the confusion. Figure 3 is a simple illustrative example of an aggregate Gaussian mixture density, it is there to provide the reader with a conceptual intuition and does not correspond to an experimental setting.

"Adequate details about the encoder and decoder in the synthetic modelling scenario have not been provided."

We thank the reviewer for pointing this out as it escaped our attention. We have now added the architectural details for all our models in the appendix - all the code will be made available also.

"In Algorithm 1, line 6, where attention A is calculated, mean is calculated as $W_q\mu$ but variance is not calculated as $W_q \sigma^2$?"

This suggestion is a valid design choice if either the weights $W_q$ are constrained or we use an activation function to ensure all entries in $W_q \sigma^2$ remain positive. However, since $\sigma(t)^2$, at attention iteration $t$, already has an indirect dependency on $W_q$ through $\mu(t-1)$ we omit this style of projection of the variances for simplicity.

"The work explicitly mentions that having additive decoders is not a necessity of the current work, but any ..."

This is unfortunately incorrect as our "NoA" model variants do in fact correspond to a non-additive convolutional decoder as stated in Section 6. Switching these with more powerful autoregressive transformer-based decoders would possibly improve the results but would constitute an unfair comparison with our baselines. For new experiments please refer to the general comment.

"R2 score is reported in Sec 6 without definition. I assume this is the correlation? This is defined as MCC in the paper."

R2 score is a coefficient of determination, it is proportional to correlation, which can be empirically observed even in our experiments. SIS score as introduced in [5] is a relative measure of R2.

"Should we not return \pi at the end of the algorithm?"

Yes thanks for pointing this out, we have now corrected it.

"Could you elaborate on assumption 6 and why it is not necessary in the current work?"

Thanks, we will add a remark explaining this in the paper. Object sufficiency is crucial when learning grounded object representations [43]. Here we do not focus on grounding so strict object sufficiency is technically not required.

[r1] Yao, W., Sun, Y., Ho, A., Sun, C. and Zhang, K., 2021. Learning temporally causal latent processes from general temporal data. arXiv preprint arXiv:2110.05428.

I thank the authors for the responses, and the new results in the common response. Please find my responses below:

• In the new results in the common response, it appears that the proposed method uses the training strategies of SPOT. The baselines could also benefit from this, after all. Shouldn't this be the fair comparison?
• Since PSA Transformer is included, a natural comparison is that of SA Transformer. Is there a reason why this was not included?
• I appreciate the qualitative results on PASCAL VOC.
• (Minor) In the attached pdf, compositional generation is shown using PSA, it would have been nice to see qualitative comparison with other SA-based efforts that allow compositional generation. I understand this is nitpicky, considering the limited space. But this would have been useful for completeness.
• While I understand the computational complexity discussion, I would have liked to see wall-clock times of training with PSA as opposed to vanilla SA, at least in approx values.

Having said the above, I do see the strengths of the paper mentioned in my original review, and stay with WA as my decision at this time.

We thank the reviewer for engaging with us and for the feedback.

"In the new results in the common response, it appears that the proposed method uses the training strategies of SPOT. The baselines could also benefit from this, after all. Shouldn't this be the fair comparison?"

We do have a fair baseline comparison as all the methods we trained used the same strategies, please refer to rows SA MLP (w/DINO) and SA MLP (w/DINO)$^\ddagger$.

"Since PSA Transformer is included, a natural comparison is that of SA Transformer. Is there a reason why this was not included?"

As explained in the general comment, we were previously not able to complete the PSA Transformer training run (\approx 15K steps) due to time constraints so it would not have been fair to compare directly with a fully trained SA Transformer (250K steps [60]). The main point was to show that PSA is scalable and using a more powerful Transformer decoder outperforms the MLP variants. Please find the updated results for SA and PSA Transformers on the PascalVOC dataset (SA Transformer results are based on our reimplementation of the DINOSAUR strategy).

"I appreciate the qualitative results on PASCAL VOC."

We are glad to hear that!

"(Minor) In the attached pdf, compositional generation is shown using PSA, it would have been nice to see qualitative comparison with other SA-based efforts that allow compositional generation. I understand this is nitpick, considering the limited space. But this would have been useful for completeness."

We emphasize that compositional generation is not the focus of the paper but is a byproduct of our theoretical framework which would be interesting to explore further in future work.

"While I understand the computational complexity discussion, I would have liked to see wall-clock times of training with PSA as opposed to vanilla SA, at least in approx values."

No problem, please find below the training iteration speeds for both models on PascalVOC with a single RTX 3090:

• SA runs at 2.31 iterations per second
• PSA runs at 2.23 iterations per second

This results in PSA being approximately 3 seconds slower per training epoch, which is quite negligible.

We thank the reviewer for their effort and overall positive outlook on our paper. We are encouraged to read that our work is perceived as solid, novel, and well-written. Please see our responses to the questions raised below.

"the experiments are only take on two toy datasets."

We kindly remind the reviewer that the objective of our work is to study the theoretical identifiability of slot representations and the conditions that ensure this property, rather than to pursue state-of-the-art empirical results. Understanding when object-centric representations can theoretically be identified is crucial for scaling slot-based methods to high-dimensional images with correctness guarantees. To verify our theoretical results, we first conduct detailed experiments on controlled datasets and then extend our demonstrations to unstructured image data. We stress that the synthetic datasets we used are necessary for properly testing our identifiability hypotheses.

Before our work, there was a lack of explanatory theory for why state-of-the-art results were able to be obtained using non-additive autoregressive Transformers (DINOSAUR [60]) and/or diffusion-based decoders (Slot-Diffusion [r2]). We showed that by viewing slot attention through a probabilistic graphical modelling perspective it is possible to prove slot identifiability for non-additive decoders using proof techniques from the identifiable generative modelling literature. Given that vanilla slot attention can be seen as a simplified version of probabilistic slot attention, akin to the relationship between soft k-means and GMMs, our theoretical results suggest why non-additive decoder structures can work well given the appropriate latent structure and inference procedure are in place.

Nonetheless, we have now evaluated our method on real-world large-scale datasets and using more powerful decoders to demonstrate that our method also scales well - please find the details of our experiments in the general comment at the top.

"How much computational complexity has increased?"

We emphasise that probabilistic slot attention (PSA) retains the $\mathcal{O}(TNKD)$ computational complexity of vanilla slot attention, where $T$ denotes the number of attention iterations, $N$ the number of input vectors, $K$ the number of slots and $D$ the slot/input dimension. The additional operations we introduce for calculating slot mixing coefficients and slot variances (under diagonal slot covariance structure) have complexities of $\mathcal{O}(NK)$ and $\mathcal{O}(NKD)$ respectively, which do not alter the dominant term. Furthermore, when used in conjunction with additive decoder-based models, PSA can reduce computational complexity by pruning inactive slots via automatic relevance determination (ARD) as outlined in Section 4.