Compositional Image Decomposition with Diffusion Models

Thank you for your constructive comments. We have added quantitative metrics for local factor decomposition and cross-dataset generalization as well as for the ablation study. In addition, we answer questions about the computational cost, machine information, and inferring factors in more detail.

Q1) Quantitative metrics for local factors and cross-dataset generalization

To quantitatively evaluate our method on local factors, we compare our method with several baselines for local factors (CLEVR) in the table below (we have also updated it in Table 1 in our main paper). Our method achieves the best performance across all three metrics, achieving the best performance on local factor decomposition.

For cross-dataset evaluation (CLEVR combined with CLEVR Toy), because there is no ground truth for recombined images, we computed FID and KID scores of generated recombination images against the original CLEVR dataset as well as the original CLEVR Toy dataset. We report in the table below and have added it in the Appendix. Our approach achieves better scores for both original datasets compared to COMET, which suggests that our generations are more successful in recombining objects from the original datasets.

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.

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.

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):

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.

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.

Thank you for the feedback. We address the comparison to a related work, and answer other questions about the computational cost, encoder, and inferred factors in more detail.

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.

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.

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.

Q4) How does the model learn to decompose images into different factors? We have described how the model learns to disentangle components in general response Section 2, and explained how we name our inferred factors in general response Section 3. To further clarify, the model learns to extract meaningful latents primarily through the information bottleneck principle. The encoder, which is jointly trained, compresses the inputs $x$ into a latent representation $z$ consisting of $K$ low-dimensional sub-latents $z_k$. Since these sub-latents are very low-dimensional, to most effectively denoise images, each of the different functions $\epsilon_\theta(x_i^t, t, z_k)$, focuses on different aspects of an image (as if any two functions focus on the same aspect, then the information capacity between the two latents is not being effectively used). This encourages disentanglement in the learned functions across $K$ sub-latent representations, where each sub-latent captures specific, independent aspects of the input images $x$. An interesting direction of future work is to add other constraints to enforce orthogonality between different sub-latents, though we had limited success when trying this. We have added this information in Section 3.2.

Since our method learns to decompose images in an unsupervised manner, there is no way of knowing what type of factors would be inferred beforehand. As a result, we provide a large set 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.

Hi, yes in Figure XVIII of the original paper you can see a visual comparison of disentanglement of primitives with baselines. We've also just added a revision to the paper to provide additional visual comparisons of disentanglement with MONet in Figure XIX (screenshot here) and recombination in Figure XX (screenshot here).

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.

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

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.