To Reviewer FyWe

Q: Limited theoretical analysis.

A: Our system could potentially learn to assign the entire image to a single mask, leaving all other masks empty. In doing so, it would essentially fall back to being equivalent to a standard DDPM UNet architecture. All synthesis would occur as a result of a single run of the decoder. However, if the system learns to split the image into different regions, it gets to use multiple runs of the decoder in order to synthesize the denoised result. There is a computational advantage (using more decoder runs) that the system can leverage to better denoising, if it can learn how to partition the problem into parallel subtasks. This is the fundamental driver behind the denoising loss encouraging learning of region partitions. We agree that a quantitative analysis of these dynamics would be valuable.

Q: Scope of experiments and Incomplete comparisons; Have you explored applying this method to more complex segmentation tasks with many object categories?

A:Regarding dataset selection, our work is at the stage of a new architectural design for diffusion-based segmentation and generation, with 2 or 3 class segmentation results demonstrating the improvement across multiple datasets, scaling up to ImageNet.

Training unconditioned DDPM from scratch for complex domains, e.g., COCO or PASCAL, is challenging, even for the single goal of high-quality image generation. Simultaneous generation and segmentation introduces another level of difficulty. Our work proposes and validates new ideas on object-centric datasets. Subsequent efforts, and more computational resources, will be needed to scale up our techniques to larger and more complex datasets.

Moreover, our work is itself on the path of scaling the segmentation capabilities of DDPMs in comparison to prior work. Consider the evolution from: DatasetDDPM (ICLR'2022) introduced the idea of using a pre-trained DDPM to extract segmentations in a supervised manner, validated on the datasets up to the scale of CelebA, to: Ours (submission'2024) trains a factorized diffusion model from scratch to accomplish both segmentation and generation simultaneously in an unsupervised manner, validated on the larger scale ImageNet dataset with more complex scenes.

Regarding baseline selection, our goal is to perform both unsupervised segmentation and image generation in a unified framework. This is an entirely novel combination of capabilities; our work is an initial proof-of-concept, not an attempt to achieve state-of-the-art unsupervised segmentation results. The relevant competing baselines mainly lie in the scope of generative models for unsupervised segmentation. Our approach is able to achieve better unsupervised segmentation results across multiple datasets. We can also achieve improved generation quality (over standard DDPM) as shown in Experimental Section 4.1. Simply putting the state-of-the-art unsupervised segmentation methods in comparison with our approach ignores the more challenging setting in which we operate (requiring that we are also able to generate images), as well as the fact that our system improves generation quality.

Q: Architectural choices and ablations.

A: We do have a more systemic investigation in different architectural choices in Appendix Section 3: reorganizing our architectural design to support hierarchical mask factorization in place of a flat set of $K$ regions. Additionally, during the rebuttal phase, we enlarge the flat set of $K$ regions to 5 for CLEVR dataset, as shown in Figure 2 of the rebuttal PDF. We also show the results with K=2 were found to be less satisfactory than K=3 for both segmentation and generation in Table 1 (rebuttal PDF).

Q: Why need to generate both contents at the same time?

A: We only train the model once and achieve both targets. Doing both together actually improves generation quality, while yielding the ability to segment as a bonus. Alternatively, through the lens of a traditional computer vision task, we learn to segment in an unsupervised manner and get the ability to generate as a bonus.

Q: How sensitive to the choice K?

A: We are defining the maximum number of regions the model can use for factorization. There is no requirement that it utilize all of these mask components. For example, it could learn to only use two mask channels out of three. Or it could predict every pixel as belonging to the same mask channel, thereby collapsing back to the vanilla UNet as a base case. The fact that the model prefers to learn something nontrivial is related to the actual structure of the data; nothing forces it to use all K channels.

Q: More details on hierarchy segmentation.

A: We have detailed the investigation in hierarchy segmentation in Appendix A.3. We formulate a hierarchical factorized diffusion architecture to progressively refine segmentation results from a coarse initial prediction to a fine, detailed final segmentation. Our initial investigation into hierarchical extensions suggests a promising future path towards handling complex scenes.

Q: How computationally expensive for training and inference, compared with DDPMs?

A: Given the nice property of weight sharing scheme in parallel decoding, the only additional weights are from our mask generator. We have an efficient batching implementation for decoding, which makes training and sampling speed comparable to standard DDPMs. As for segmentation, the inference speed is the same as a single forward pass in a standard DDPM.

Q: Runtime compared to other unsupervised segmentation methods?

A: If both use the same UNet encoder-decoder architecture, the inference complexity is the same. The parallel decoding scheme and diffusion process does not affect the efficiency of segmentation inference speed.

Q: Some literature can be referred to.

A: Thanks for the suggestion. We will include a discussion in the final version.

To Reviewer 9WST

Q: Experiments only cover 2-3 class scenarios.

A: Our work is at the stage of a new architectural design for diffusion-based segmentation and generation, with 2 or 3 class segmentation results demonstrating improvements across multiple datasets, scaling up to ImageNet.

Training unconditioned DDPM from scratch for complex domains, e.g., COCO or PASCAL is challenging, even if limited to the goal of high-quality image generation. We are addressing an additional challenge of simultaneously generating a latent representation of segmentation.

Our work proposes and validates new ideas on object-centric datasets. We believe there is a path toward scaling our method to more complex data; e.g., see the Appendix for our effort towards hierarchical segmentation. Fully exploring the new technique we have introduced will require follow-up papers.

Q: ImageNet results only on downsampled 64x64 images.

A: This choice is due to computational constraints and the need to validate our method's feasibility at a lower resolution before potentially scaling to higher resolutions. Training models on full-resolution (e.g., 256 $\times$ 256) ImageNet images would require significantly more computational resources and time. Resolution is an aspect of the system orthogonal to our architectural innovation.

Q: In addition to DiffuMask, please consider discussing DiffSeg, which like DiffuMask, relies on a Stable Diffusion model, but unlike it, does output masks explicitly.

A: Thanks for the suggestion. We will add the discussion about DiffuMask in the revised draft.

Q: Is the proposed scheme specific to diffusion? Could it be applied in other autoencoder settings?

A: While the details of our proposed scheme are tailored for diffusion models, our core principles could potentially be adapted to other autoencoder architectures, e.g., variational autoencoders (VAEs) or masked autoencoders (MAEs):

• Bottleneck Design: The use of a structured bottleneck could be implemented in autoencoders to encourage learning of meaningful latent representations that aid in segmentation.
• Parallel Decoding: The idea of parallel decoding for different segments could be utilized in the decoder portion of any autoencoder to encourage factorizing the reconstruction task into independent subtasks.

Q: What happens if you set K to higher values, despite not needing the additional classes? e.g., K = 4, 5 for binary segmentation?

A: $K$ is the maximum number of regions the model may use; it could learn fewer. Setting K to higher values may allow for some additional flexibility during training, as the model learns how to partition images into regions. Setting K much larger than necessary simply wastes compute as the model learns to leave some region maps empty (unused). During the rebuttal phase, we experiment with enlarging the flat set of $K$ regions to 5 for the CLEVR dataset, as shown in Figure 2 (rebuttal PDF). The model still performs foreground-background segmentation adequately by ignoring extra classes if not present.

Q: You mentioned that K = 3 was helpful for binary segmentation. How much worse were the results with K=2?

A: For binary segmentation, we found setting $K=3$ rather than $K=2$ to assist training, with learned regions emerging as foreground, background, and a contour or transition between the two, as shown in Figure 1 of the rebuttal PDF. Also, as shown in Table 1 in the rebuttal material, the results with K=2 were less satisfactory for both segmentation and generation. With the goal of handling more regions and multiscale structure, we believe a more promising future investigation is to reorganize our architectural design to support hierarchical mask factorization in place of a flat set of $K$ regions, as shown in Appendix Section A.3.

To Reviewer vzUs

Q: The abstract lacks context/motivation. The conclusions and experiments lack discussions of the limitations of the proposed approach.

A: In the current abstract, we provide a clear introduction to the challenges and motivations driving our research. The abstract begins by discussing the limitations of supervised deep learning, specifically the reliance on large amounts of annotated data, which is a well-known barrier in computer vision tasks. We explicitly mention the goal of achieving unsupervised image generation and segmentation, addressing the significant problem of high annotation costs in supervised learning. The abstract succinctly outlines our approach—a denoising diffusion model with a computational bottleneck—which is a novel strategy for tackling the dual tasks of generation and segmentation simultaneously. It concludes by highlighting the success of our model in achieving high-quality image synthesis and segmentation across multiple datasets, reinforcing the effectiveness of our approach.

The conclusion and experiment sections of the paper do acknowledge and discuss limitations, which are crucial for a balanced understanding of the proposed method. The conclusion acknowledges that the work is at an early stage of architectural design, focusing on simpler class scenarios to establish a proof of concept before tackling more complex tasks. We discuss that the experiments are performed primarily on two to three class scenarios, recognizing these as initial limitations with plans to extend the approach to higher resolutions and more complex segmentation tasks in future work. The experiments highlight the need for further exploration of scalability and generalization to more diverse datasets and tasks. This is particularly important for understanding how the model performs in real-world applications beyond the current scope. The conclusion outlines future research directions aimed at addressing these limitations, such as investigating hierarchical extensions and handling more complex scenes.

Q: The explanation of the model is unclear .

A: This question is about Eq.(5), which is not our proposed design. In line 158-160, "However, such a design significantly modifies (e.g., channel sizes) the original U-Net decoder architecture. Moreover, conditioning with the whole mask representation may also result in a trivial solution that simply ignores region masks." As such, there's no concatenation of mask and $h_enc$ in our proposed parallel decoding scheme. Refer to Eq.(6) and Eq. (7) for the description of our design.

We again clarify the design as depicted in Figure 2: $h_{mid}$ (latent features) is directly passed to the decoders following conventional DDPM without concatenation. $m_k$ denotes the $k-th$ channel of the mask generator final output $m$. The mask $m_k$ is used at each decoding layer as part of skip connections. It provides region-specific information that helps the decoder refine outputs by distinguishing between different semantic regions. As shown in Figure 2, we adopt channel-wise masking: Each channel of $m$, denoted as $m_k$, is applied to $h_{enc}$ through element-wise multiplication. We downsample $m_k$ using bilinear interpolation to align its spatial dimensions with $h_{enc}$ at different stages. Only element-wise multiplication is used with consistent spatial dimensions, allowing integration across the network's architecture without significantly changing the channel dimension of conventional DDPM.

Q: The experiments are a limited when it comes to the datasets used to evaluate the model for semantic segmentation.

A: Training unconditioned DDPM from scratch for complex domains, e.g., COCO or PASCAL, is challenging, even for the single goal of high-quality image generation. Simultaneous generation and segmentation introduces another level of difficulty. Our work proposes and validates new ideas on object-centric datasets. Subsequent efforts, and more computational resources, will be needed to scale up our techniques to larger and more complex datasets.

Moreover, our work is itself on the path of scaling the segmentation capabilities of DDPMs in comparison to prior work. Consider the evolution from: DatasetDDPM (ICLR'2022) introduced the idea of using a pre-trained DDPM to extract segmentations in a supervised manner, validated on the datasets up to the scale of CelebA, to: Ours (submission'2024) trains a factorized diffusion model from scratch to accomplish both segmentation and generation simultaneously in an unsupervised manner, validated on the larger scale ImageNet dataset with more complex scenes.

As for the segmentation results, we provide a supervised result as a reference point against which to compare unsupervised methods using the same architecture. One would not expect any unsupervised method, limited to training on the same data, to be able to match the supervised method's performance. Our method, which is unsupervised, outperforms the unsupervised baseline (DatasetDDPM-unsup) based on the same UNet architecture.

It is the case that we want to explore much larger K for ImageNet in the future, as doing so will be necessary for segmenting complex scenes. A challenge is implementing K parallel pathways efficiently for large K; note that training is already expensive as everything takes place within a diffusion model.

To Reviewer FkB4

Q: Limited discussion of related work; How does the proposed method differ from the mentioned related works? What are the key advantages it offers?

A: Thanks for the suggestion. We will incorporate this discussion into our related work section to provide a clear picture of how our approach diverges from existing methods:

• DatasetGAN has utilized intermediate GAN features to perform segmentation, demonstrating the potential of leveraging generative models for downstream tasks. However, our approach differs by directly integrating segmentation within the diffusion process, eliminating the need for separate feature extraction and potentially offering more coherent results.

• Similarly, Diffuse, Attend, and Segment employs attention mechanisms for zero-shot segmentation, providing high-quality segmentation through attention-driven methods. Our factorized diffusion architecture offers an alternative by producing segmentation masks as an intrinsic part of the denoising process, potentially yielding more integrated and coherent segmentation results.

• In the realm of self-supervised learning, DINO showcases the power of representation learning through knowledge distillation, applicable to tasks like segmentation. Our approach aligns with this goal of versatile representation learning but achieves segmentation directly within the generative model, eliminating the need for additional supervision or distillation processes.

• FeatUP highlights the adaptability of generative model representations for downstream tasks. While FeatUP focuses on broad applicability, our method targets segmentation with a specifically designed architecture, achieving simultaneous image generation and segmentation in a unified framework.

In comparison with prior work, we are offering an entirely new approach to solving an end task (segmentation) using unsupervised generative learning: (1) Specify an architecture whose bottleneck representation encodes the solution to the end task, and whose decoder's computational structure constrains synthesis of data from that representation. (2) Train the system end-to-end, from scratch, for generation alone. (3) Read off the bottleneck features as the solution.

Q: Are the results quantitatively compared with existing segmentation methods (e.g., DatasetGAN, etc.)?

A: Both DatasetGAN and DatasetDDPM are not unsupervised approaches like ours. They require annotations and train in a few shot manner. In the paper, we compared with DatasetDDPM, which is a more suitable baseline. As shown in Table 3 and 4, our method outperforms DatasetDDPM by a large margin under the same unsupervised setting.

Q: The ethical implications of using generative models to create harmful content should be addressed.

A: We are committed to ensuring that the development and application of generative models are guided by ethical principles. We will incorporate the discussion in the final version.