Paper: Compositional Image Decomposition with Diffusion Models (ICLR-2024)

Reviewer comment (Evaluation):
As the proposed method contains a set of diffusion models, the computational cost of the proposed method and existing works should be discussed in the paper.

Author response:
**Q2) Computational cost**
Please see general response Section 4. Our method uses $K$ diffusion models, so the computational cost is $K$ times that of a normal diffusion model. In practice, the method is implemented as $1$ denoising network that conditions on $K$ latents, as opposed to $K$ individual denoising networks. One could significantly reduce computational cost by fixing the earlier part of the network and only condition on latents in the second half of the network. This would likely achieve similar results while reducing computation. In principle, we could also parallelize $K$ forward passes for each factor, thus reducing both training and inference time.  We have added this detail in the Section E.3 of the Appendix. Regarding the computational cost of other approaches, Beta-VAE, MONet, COMET, and Hessian Penalty generally used less training time until convergence, while Slot Attention and GENESIS-V2 use a comparable training time of 24 hours.

Reviewer comment (Reproducibility):
For training details in the supplemental, each model is trained on an NVIDIA V100 or an NVIDIA RTX 2080 with the same hours. I was wondering whether the model performance would be different. In addition, is the memory of NVIDIA RTX 2080 24GB or 8GB?

Author response:
**Q3) Machine specs**
Our models were trained with the same batch size of 32 images, so the machine type should not affect model performance. The memory of a NVIDIA RTX 2080 machine is 24GB.

Reviewer comment (Evaluation):
For the ablation study, why use MSE and LPIPS to evaluate the reconstruction quality, rather than the metrics used in Table 1? How about the results of the ablated versions on other datasets used in the paper?

Author response:
**Q4) Ablation study metrics**

Thank you for your question. Our method not only decomposes images into a set of factors, but also enables recombination of all inferred factors to reconstruct the input image. As a result, we used MSE and LPIPS to measure both pixel-wise errors and perceptual similarity of the original input image and reconstructed image. In addition to MSE and LPIPS, we have also added FID and KID to the ablation study in the table below and update it in the main paper accordingly (Table 2):

| Dataset | Multiple Components | Predict x_0 | MSE &darr; | LPIPS &darr; | FID &darr; | KID &darr;
|  :----------------  |  :------:  |  :------:  | :------:  | :------:  | :------:  | :------:  |
| CelebA-HQ | YES | NO | 105.003 | 0.603 | 155.46 | 0.141 
| CelebA-HQ | NO | YES | 88.551 | 0.192 | 30.10 | 0.022
| CelebA-HQ | YES | YES | **76.168** | **0.089** | **16.48** | **0.013** 
| CLEVR | YES | NO | 56.179 | 0.3601 | 42.72 | 0.033
| CLEVR | NO | YES | 26.094 | 0.2236 | 24.27 | 0.023
| CLEVR | YES | YES | **6.178** | **0.0122** | **11.54** | **0.010**

In the table above, we have also included ablation study for the CLEVR dataset. The non-ablated version of the model, with predicting $x_0$ and using multiple components (i.e., 4 components), achieves the best scores on both CelebA-HQ and CLEVR datasets across all metrics, indicating that the conclusion is consistent across different datasets.

Reviewer comment (Presentation):
How to determine the types of factors that can be inferred from the image? For example, I am not sure whether the second to the fourth columns correspond to shadow, objects, and background respectively.

Author response:
**Q5) How are the types of inferred factors determined?**
Please see general response Section 3. Since our method learns to decompose images in an unsupervised manner, there is no way to know what type of factors would be inferred beforehand. As a result, we provide a large number of example images and name inferred factors through visual inspection in order to help readers understand what each component represents. We have updated the Experiment section to clarify this.

Reviewer comment (Novelty):
The method seems to be similar to [1]

Author response:
**Q1) The method seems similar to [1]**
Though both works involve diffusion models representing concepts, the goal of this paper is to discover and decompose images into component diffusion models, while [1] focuses on composing pre-trained diffusion models. To accomplish this decomposition, we propose a new training-time objective to find each component.

Reviewer comment (Evaluation):
What is the computational cost? It may takes more space and computational resources with K diffusion models

Author response:
**Q2) What is the computational cost?**
Thank you for your question. Our method computes $K$ conditional scores for a total of $K$ latents, so the computational cost is $K$ times that of a normal diffusion model. In practice, the method is implemented as $1$ denoising network that conditions on $K$ latents, as conditioning on different latents can make arbitrary functions in principle. One could significantly reduce computational cost by fixing the earlier part of the network since latents are only conditioned in the second half of the network. In addition, in practice, we parallelize $K$ forward passes for each factor, thus keeping training and inference time constant. 

We have added this detail in the Section E.3 of the Appendix.

Reviewer comment (Reproducibility):
Is the learned encoder Encθ(x) pre-trained or trained with diffusion model jointly?

Author response:
**Q3) How is the encoder learned?**
Our method jointly learns the latent encoder and the denoising network jointly. We have also updated Section 3.2 in the paper to clarify this choice.

Reviewer comment (Theory):
I did not find in the paper an explanation about the encoder z = Enc_\theta(x). How Enc_() is learnt? What ensures that the decomposition is at the local (ie objects, things) or global (ie, illumination, stuffs) level? What ensures the disentanglement of the decomposition? (no additional constraints are enforced during the learning stage). Those aspects might have been discussed in the original paper COMET (I did not read it), however, it is worth to discuss them again in the current paper since it is key for the understanding and analysis of the success/failures of the proposed approach.

Author response:
Thank you for your detailed comments and positive evaluation. We have added a description of the encoder training and explanation of how the approach learns to disentangle latents in the text. 

**Q1) How is the encoder learned?**
Thank you for your question. As addressed in the general response 1, our encoder and denoising network are learned jointly using the denoising loss. The main reason behind this design choice is that we can naturally enable the model to learn different types of latent representations from various datasets. 

We have also clarified this in the method section accordingly.

**Q2) What ensures the type of decomposition?**
Since our method is an unsupervised approach to decompose images into a set of latents, we aren’t able to know the types of inferred concepts those learned latents would be ahead of time. Based off visual inspection, we provided a 'psuedolabel' for each component based off the visualization of the component. To more clearly display each component, we provide a set of visualizations in our experiments and appendix to help readers understand inferred concepts. 

We have also updated this in the Experiment section to avoid any further confusion.

**Q3) How is disentanglement ensured?** 
Please see our general response 2. Disentanglement is primarily ensured by information bottleneck. The encoder encodes the information in $x$ into a compact low-dimensional latent representation $z$ consisting of $K$ sub-latents $z_k$. To most effectively denoise an image given these low-dimensional latents, each function $\epsilon_\theta(x_i^t, t, z_k)$ learns to focus on distinct information in an image. This enforces diversity and disentanglement in the learned functions across $K$ sub-latent representations. While we explored additional constraints to enforce independence in components, we found that it gave similar performance to just using a low dimensional latent. 

We have added this information in Section 3.2.

Reviewer comment (Theory):
The authors argue that they 'learn a set of different denoising functions to recover an image x_i' (page 4). However, the denoising function \epsilon_\theta is not, in eq.9, parameterized by k. The only dependance to k is in the input latent variable z_k. It would imply that there is a single denoising function, but with different input argument (in particular the z_k). Please clarify.

Author response:
**Q3) Why is one single denoising network used instead of $K$**
Thank you for your question. In principle, a model can learn arbitrary functions by conditioning on different latent representations. Thus one single denoising network that conditions on $K$ different latents should achieve similar performance as $K$ denoising networks. In addition, our method is more memory-efficient as it only requires one network instead of $K$.

Reviewer comment (Writing):
The encoder Enc_\theta() and the denoising function \epsilon_\theta, are both parameterized by \theta. This is probably a typo, the two networks being parameterized by two sets of independent weights, \theta_1 and \theta_2. Please correct as needed in the paper.

Author response:
**Q4) Typos**
Thank you for pointing that out. We have updated the encoder parametrization notation to be $Enc_\phi$ instead of $Enc_\theta$. We will make sure that our notations are consistent in our next version of the paper.