Factorized Diffusion Architectures for Unsupervised Image Generation and Segmentation

We develop a neural network architecture which, trained in an unsupervised manner as a denoising diffusion model, simultaneously learns to both generate and segment images. Learning is driven entirely by the denoising diffusion objective, without any annotation or prior knowledge about regions during training. A computational bottleneck, built into the neural architecture, encourages the denoising network to partition an input into regions, denoise them in parallel, and combine the results. Our trained model generates both synthetic images and, by simple examination of its internal predicted partitions, a semantic segmentation of those images. Without any finetuning, we directly apply our unsupervised model to the downstream task of segmenting real images via noising and subsequently denoising them. Experiments demonstrate that our model achieves accurate unsupervised image segmentation and high-quality synthetic image generation across multiple datasets.

Paper

Similar papers

Peer review

Reviewer FyWe5/10 · confidence 5/52024-07-03

Summary

This paper proposes a novel architecture for denoising diffusion probabilistic models (DDPMs) that enables simultaneous unsupervised image generation and segmentation. The key idea is to constrain the denoising network with a structured bottleneck that factorizes the image into regions, which are then denoised in parallel. This architectural design encourages the model to learn meaningful segmentations as a byproduct of training for image generation. The authors demonstrate that their approach can generate high-quality images along with corresponding segmentations, as well as segment real images, without any supervision or additional loss terms. Experiments on multiple datasets show improvements in both image quality and segmentation accuracy compared to baselines.

Strengths

* Novel architectural design that unifies unsupervised image generation and segmentation in a principled way * Achieves strong results on both tasks without requiring additional loss terms or supervision * Provides insights into how the model learns to factorize images into regions * Explores extensions to hierarchical segmentations

Weaknesses

* Limited theoretical analysis: The paper lacks a rigorous theoretical foundation for why the proposed architecture leads to meaningful segmentations. While empirical results are strong, a deeper analysis of why factorizing the denoising process encourages semantic segmentation would strengthen the contribution. For instance, the authors could explore connections to information bottleneck principles or analyze the gradients flowing through different parts of the network. * Scope of experiments: The experiments primarily focus on datasets with relatively simple segmentations (2-3 regions), such as Flower, CUB, and face datasets. While the ImageNet results are promising, they are limited in scope. The paper lacks experiments on more challenging datasets with complex, multi-object scenes (e.g., COCO, Cityscapes) that would better demonstrate the method's scalability and generalizability. * Incomplete comparisons: The paper misses comparisons with some recent, relevant unsupervised segmentation methods, such as PiCIE (CVPR 2021) or STEGO (ICCV 2021). Including these would provide a more comprehensive evaluation of the state-of-the-art. For the image generation task, comparisons with other recent diffusion model variants (e.g., latent diffusion models) would be valuable to contextualize the improvements. * Architectural choices and ablations: The paper does not extensively explore variations in the factorized architecture. For instance, how does performance change with different numbers of parallel decoders or alternative ways of combining their outputs? More comprehensive ablation studies would help isolate the impact of different components of the proposed architecture.

Questions

* The setting of generating both the images and the segmentation maps is interesting. However, Isn't this setup a little trival? Why we need to generate both contents at the same time? How could this feature help us? * How sensitive is the method to the choice of number of regions K? Is there a way to automatically determine the optimal K for a given dataset? * The paper mentions that the method can be extended to hierarchical segmentations. Could you provide more details on how this would work and what challenges might arise? * Have you explored applying this method to more complex segmentation tasks with many object categories? What modifications might be needed? * How computationally expensive is the proposed method compared to standard DDPMs, both for training and inference? * The paper claims the method can be applied to segment novel images with just one forward pass. How does the runtime compare to other unsupervised segmentation methods? * Some literature that share the similar insights of unifying generation and segmentation can be refered to: [1] Lai Z, Duan Y, Dai J, et al. Denoising diffusion semantic segmentation with mask prior modeling[J]. arXiv preprint arXiv:2306.01721, 2023. [2] Li C, Liu X, Li W, et al. U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation[J]. arXiv preprint arXiv:2406.02918, 2024.

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have adequately addressed the limitations of their work in Section 5, acknowledging that the current results are limited to 2-3 class segmentations and discussing the need for further work on handling more complex scenes. They also mention computational costs as a potential limitation. The paper does not directly discuss potential negative societal impacts, which could be briefly addressed, though the risks seem relatively low for this type of fundamental methodological work.

Reviewer vzUs4/10 · confidence 4/52024-07-11

Summary

The authors propose a model for simultaneous unsupervised semantic segmentation and image generation based on denoising diffusion probabilistic models. The methodology alters the architecture of a typical DDPM by conditioning the decoder on the outputs of a mask generator.

Strengths

In general, the paper is well-written and technically sound. The experimental results surpass the state-of-the-art on the datasets they were shown and the results seem promising.

Weaknesses

The abstract lacks context/motivation. The conclusions and experiments lack discussions of the limitations of the proposed approach. The authors propose a simple solution to a complex problem (unsupervised segmentation) through altering the architecture of a DDPM by adding an additional decoder to generate masks in an unsupervised manner and conditioning the DDPM’s original decoder on the outputs of the mask generator. However, the explanation of the model is unclear: Initially, 𝑚𝑘 is introduced as the output of the mask generator after applying the softmax activation function, for each of the k channels. As such, 𝑚𝑘 should have dimensions width x height x 1, where width and height are the dimensions of the original images. Then, it is stated that 𝑚𝑘 is concatenated with ℎenc (constituted by several outputs of the encoder provided to the decoder as skip connections) and ℎ𝑚𝑖d (latent vector) and provided as input. It is unclear how these variables of different dimensions are concatenated. In the manuscript, it says “We downsample 𝑚 accordingly to the same resolution as ℎ𝑚𝑖d and ℎenc at different stages” (line 157), however it is still unclear how exactly this downsampling step is achieved. Does 𝑚 represent the outputs of each layer of the mask generator and is it provided as input to the decoder as skip connections? If so, how is 𝑚𝑘 obtained from 𝑚? Do the outputs of all layers of the mask generator possess exactly 𝑘 channels? An example highlighting the dimensions of 𝑚𝑘, ℎ𝑚𝑖d and ℎenc and how these are concatenated would be helpful. The experiments are a limited when it comes to the datasets used to evaluate the model for semantic segmentation. The only two datasets used for over 2 classes (distinguishing more than background and foreground) were face datasets, where all images are focused on the face and present relatively low variability. It is unclear how the method would work with over 3 classes and on datasets with higher variability such that not all objects are present in all images. Furthermore, even when only two classes are used, to distinguish background from foreground, the datasets seem to be limited, with low variability in backgrounds, making it difficult to assess the model’s results. The results on the CUB dataset, whose backgrounds present more variability and the regions with objects of interest (birds) are smaller, show difficulties segmenting the birds from the backgrounds, especially in the presence of branches (Fig 8). In the ImageNet experiment, it would also be interesting to see what would happen if 5 classes (4 objects + background) were used (k=5 in the softmax layer) rather than just two. Would the network be able to learn to distinguish certain objects or would it still only be able to separate background from foreground? As it stands, it becomes difficult to assess the proposed network’s strengths and whether it is truly capable of capturing semantic information or whether it just separates background from foreground without requiring more in-depth semantic information about the nature of the objects. The paper lacks a discussion of the limitation of the methods.

Questions

see above the comments about the concatenation of ℎ𝑚𝑖d and ℎenc.

Rating

4

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

yes.

Reviewer FkB44/10 · confidence 4/52024-07-12

Summary

The paper introduces a novel neural network architecture capable of simultaneous image generation and segmentation in an unsupervised manner, eliminating the need for pre-labeled data. The core concept involves training the network to dissect an image, clean individual sections, and reassemble them. Notably, the system can decipher the image content solely through the analysis of these sections. The authors demonstrate the effectiveness of their unsupervised model for real-image segmentation. This work presents a promising approach for generating realistic images and understanding their content with minimal reliance on labeled data.

Strengths

- Clear and concise presentation - Experimental results validate the proposed method's effectiveness

Weaknesses

- Limited discussion of related work: The paper lacks a comprehensive review of relevant research, hindering the assessment of genuine novelty. It's crucial to mention existing works like DatasetGAN (utilizing inner features of trained GANs for segmentation), Diffuse, Attend, and Segment (high-quality zero-shot segmentation with attention maps), DINO, and FeatUP (representation utilization in generative models for downstream tasks).

Questions

- How does the proposed method differ from the mentioned related works? What are the key advantages it offers? - Are the results quantitatively compared with existing segmentation methods (e.g., DatasetGAN, etc.)?

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

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

Reviewer 9WST6/10 · confidence 3/52024-07-12

Summary

The paper proposed a structural modification of a DDPM which causes it to learn a decomposition of images into regions. This factorization enables unsupervised segmentation and simultaneously improves the quality of the generated images. The method is evaluated on various datasets, demonstrating its effectiveness in both tasks.

Strengths

- Evaluations on 5 datasets. - Consistently good performance relative to other unsupervised methods. - Conceptually simple but powerful design. - Ablation on encoder design details. - Zero-shot segmentation shown on 2 additional datasets.

Weaknesses

- Experiments only cover 2-3 class scenarios. - ImageNet results only on downsampled 64x64 images.

Questions

- In addition to DiffuMask, please consider discussing DiffSeg (https://arxiv.org/pdf/2308.12469), which like DiffuMask, relies on a Stable Diffusion model, but unlike it, does output masks explicitly. - Is the proposed scheme specific to diffusion? Could it be applied in other autoencoder settings? - What happens if you set K to higher values, despite not needing the additional classes? e.g. K = 4, 5 for binary segmentation? - You mentioned that K = 3 was helpful for binary segmentation. How much worse were the results with K=2?

Rating

6

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

Limitations are discussed explicitly in the paper. No concerns about negative societal impact.

Area Chair afhb2024-08-09

Rebuttal is online - please respond

Dear Reviewers, Authors carefully prepared their rebuttal - trying to address the concerns you have raised. Please check the rebuttals and join the discussion about the paper. Regards,

Reviewer FyWe2024-08-10

> 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. Regarding this point, is there any evidence (theoretically or empirically) that could support this claim? It seems not that straightforward. Thank you! Best,

Authorsrebuttal2024-08-13

To Reviewer FyWe

**Q: Is there any evidence (theoretically or empirically) that could support this claim [doing both together actually improves generation quality]?** A: Yes, Table 5 (top) gives empirical evidence for precisely this claim. Our system, which is a generation architecture that internally performs segmentation, generates higher quality images (lower FID) than a standard DDPM baseline. For each dataset in Table 5, we train from scratch each system (ours and the DDPM baseline) in a fully unsupervised manner driven by the denoising objective alone. Our system and the DDPM baseline have comparable design motifs and similar parameter counts, with the distinction that, as described in Sections 3.1 and 3.2, ours adds a mask generator module and partitions synthesis into parallel decoding pathways. This structural change to the architecture significantly improves FID across all datasets (e.g., 13.35 to 10.79 for FFHQ-128 and 7.02 to 6.54 for ImageNet-64). Thus, our system is a denoising diffusion model with **improved generation quality** over a baseline DDPM, and is trained in the exact same *fully unsupervised manner*. If someone were only interested in high quality generation, our system would be preferable to the standard DDPM. Beyond this, our system **produces segmentation as a bonus**; segmentation can be read from the internal bottleneck state (region masks) of our architecture. Alternatively, someone interested in learning to segment in an unsupervised manner could view the image generation capability of our system as a bonus. As shown in Figure 1 (a)(b), we only train our model once and achieve both targets (image generation, segmentation). The fact that our architecture improves generation quality makes it a candidate to serve as the basis for future diffusion-based foundation models. Viewed in a broader context, it is perhaps not too surprising that network architecture design can have a strong influence on learning and on prediction or output quality. Convolutional neural networks and attention mechanisms in transformers are two examples of imbuing neural architectures with domain-relevant structure. We are imbuing domain-relevant structure at a more macro scale, in the form of parallel synthesis pathways that are a match to a compositional model of images.

Reviewer 9WST2024-08-12

Thank you for thoroughly responding to my questions. I understand the concerns about computational resources -- mentioning this explicitly in the paper would be helpful. Similarly, the fact that the general idea is potentially applicable to other architectures might warrant a short comment in the discussion section. Regarding K=2, I was surprised to see the magnitude of the impact on the metrics. Do you have an intuitive explanation for this? Is it the case that with K=2 the shapes of the masks do not adhere well to the contours of the underlying objects, which K>2 helps to alleviate?

Authorsrebuttal2024-08-13

To Reviewer 9WST

We will update our discussion to include computational requirements and potential applicability to other autoencoder frameworks. **Q: Regarding K=2 and impact on the metrics** The choice of K not only defines an architectural structure for inference, but one that comprises the model during training; K is the maximum number of components that the model may utilize at any point in training or inference. We train from scratch, so at the start of training, the assignment of pixels to components will be a result of random initialization. Training via gradient descent must find a path from this random configuration to a parameter configuration which yields structured region masks. It might be helpful to have more components (more masks) to utilize during training in order to smooth the optimization landscape, even if by the completion of training not all masks are needed. An analogy is overparameterization benefiting neural network training, even though networks are subsequently amenable to pruning. Consistent with this hypothesis is the observation that our K = 3 experiments for objects vs background datasets show the third component emerging as a contour or transition between the object and background (Figure 1 in rebuttal PDF). To provide further analysis, we will add visualization of region mask evolution over the duration of training to the Appendix.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC