Factorized Diffusion Architectures for Unsupervised Image Generation and Segmentation

Authors: xin  yuan
†  google
michael  maire
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 06:07:55

Abstract: 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, semantic segmentations of those images.  Without fine-tuning, 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.

Main paper: 
 introduction: 
 Supervised deep learning yields powerful discriminative representations, and has fundamentally advanced many computer vision tasks, including image classification [13, 58, 21, 28] , object detection [18, 49, 41] , and semantic and instance segmentation [42, 22, 33] . Yet, annotation efforts [13] , especially those involving fine-grained labeling for tasks such as segmentation [39] , can become prohibitively expensive to scale with increasing dataset size. This motivates an ongoing revolution in self-supervised methods for visual representation learning, which do not require any annotated data during a large-scale pre-training phase [7, 15, 67, 35, 23, 10, 12] . However, many of these approaches, including those in the particularly successful contrastive learning paradigm [23, 10, 12] , still require supervised fine-tuning (e.g., linear probing) on labeled data to adapt networks to downstream tasks such as classification [23, 10] or segmentation [8, 69] .In parallel with the development of self-supervised deep learning, rapid progress on a variety of frameworks for deep generative models [32, 19, 65, 66, 60, 37, 27, 59, 50] has lead to new systems for high-quality image synthesis. This progress inspires efforts to explore representation learning within generative models, with recent results suggesting that image generation can serve as a good proxy task for capturing high-level semantic information, while also enabling realistic image synthesis.Building upon generative adversarial networks (GANs) [19] or variational autoencoders (VAEs) [32] , InfoGAN [11] and Deep InfoMax [26] demonstrate that generative models can perform image classification without any supervision. PerturbGAN [4] focuses on a more complex task, unsupervised image segmentation, by forcing an encoder to map an image to the input of a pre-trained generator so that it synthesizes a composite image that matches the original input image. However, here training is conducted in two stages and mask generation relies on knowledge of predefined object classes. We design a denoising diffusion model with a specific architecture that couples region prediction with spatially-masked diffusion over predicted regions, thereby generating both simultaneously. (b) An additional byproduct of running our trained denoising model on an arbitrary input image is a segmentation of that image. Using a model trained on FFHQ [31] , we achieve both high quality synthesis of images and corresponding semantic segmentations (c-d), as well as the ability to accurately segment images of real faces (e-f).Segmenting a real image is fast, requiring only one forward pass (one denoising step).Denoising diffusion probabilistic models (DDPMs) [27] also achieve impressive performance in generating realistic images. DatasetDDPM [2] investigates the intermediate activations from the pre-trained U-Net [51] network that approximates the Markov step of the reverse diffusion process in DDPM, and proposes a simple semantic segmentation pipeline fine-tuned on a few labeled images.In spite of this usage of labels, DatasetDDPM demonstrates that high-level semantic information, which is valuable for downstream vision tasks, can be extracted from pre-trained DDPM U-Net. Diff-AE [48] and PADE [70] are recently proposed methods for representation learning by reconstructing images in the DDPM framework. However, their learned representations are in the form of a latent vector containing information applicable for image classification.In contrast to all of these methods, we demonstrate a fundamentally new paradigm for unsupervised visual representation learning with generative models: constrain the architecture of the model with a structured bottleneck that provides an interpretable view of the generation process, and from which one can simply read off desired latent information. This structured bottleneck does not exist in isolation, but rather is co-designed alongside the network architecture preceding and following it. The computational layout of these pieces must work together in a manner that forces the network, when trained from scratch for generation alone, to populate the bottleneck data structure with an interpretable visual representation.We demonstrate this concept in the scenario of a DDPM for image generation and semantic segmentation as the interpretable representation to be read from the bottleneck. Thus, we frame unsupervised image segmentation and generation in a unified system. Moreover, experiments demonstrate that domain-specific bottleneck design not only allows us to accomplish an end task (segmentation) for free, but also boosts the quality of generated samples. This challenges the assumption that generic architectures (e.g., Transformers [61] ) alone suffice; we find synergy by organizing such generic building blocks into a factorized architecture which generates different image regions in parallel.Figure 1 provides an overview of our setting alongside example results, while Figure 2 illustrates the details of our DDPM architecture which are fully presented in Section 3. This architecture constrains the computational resources available for denoising in a manner that encourages learning of a factorized model of the data. Specifically, each step of the DDPM has the ability to utilize additional inference passes through multiple copies of a subnetwork if it is willing to decompose the denoising task into parallel subproblems. The specific decomposition strategy itself must be learned, but, by design, is structured in a manner that reveals the solution to our target task of image segmentation. We summarize our contributions as three-fold: • Unified learning of generation and segmentation. We train our new DDPM architecture once, obtaining a model directly applicable to two different tasks with zero modification or fine-tuning: image generation and image segmentation. Segmenting a novel input image is fast, comparable in speed to any system using a single forward pass of a U-Net [51] like architecture. • Unsupervised segmentation for free. Our method automatically learns meaningful regions (e.g., foreground and background), guided only by the DDPM denoising objective; no extra regularization terms, no use of labels. • Higher quality image synthesis. Our model generates higher-quality images than the baseline DDPM, as well as their corresponding segmentations simultaneously. We achieve excellent quantitative and qualitative results under common evaluation protocols (Section 4).Beyond improvements to image generation and segmentation, our work is a case study of a new paradigm for using generation as a learning objective, in combination with model architecture as a constraint. Rather than viewing a pre-trained generative model as a source from which to extract and repurpose features for downstream tasks, design the model architecture in the first place so that, as a byproduct of training from scratch to generate, it also learns to perform the desired task. 
related work: 
 Image Segmentation. Generic segmentation, which seeks to partition an image into meaningful regions without prior knowledge about object categories present in the scene, is a longstanding challenge for computer vision. Early methods rely on combinations of hand-crafted features based on intensity, color, and texture cues [6, 44] , clustering algorithms [57] , and a duality between closed contours and the regions they bound [1] . Deep learning modernized the feature representations used in these pipelines, yielding systems which, trained with supervision from annotated regions [43] , reach near human-level accuracy on predicting and localizing region boundaries [3, 56, 64, 34] .Semantic segmentation, which assigns a category label to each pixel location in image, has been similarly revolutionized by deep learning. Here, the development of specific architectures [42, 51, 20] enabled porting of approaches for image classification to the task of semantic segmentation.Recent research has refocused on the challenge of learning to segment without reliance on detailed annotation for training. Hwang et al. [29] combine two sequential clustering modules for both pixel-level and segment-level to perform this task. Ji et al. [30] and Ouali et al. [46] follow the concept of mutual information maximization to partition pixels into two segments. Savarese et al. [54] further propose a learning-free adversarial method from the information theoretic perspective, with the goal of minimizing predictability among different pixel subsets. Note that even completely unsupervised foreground/background segmentation is a non-trivial task. Liu et al. [40] , a recent advance in this regime, produces similar region mask output, yet depends entirely upon motion cues from video for training. We achieve such unsupervised learning from static images alone.Learning Segmentation in Generative Models. Previous generative model-based approaches learn semantic segmentation by perturbing [4] or redrawing [9] generated foreground and background masks. Despite good performance, these methods apply only to two-class partitions and require extra loss terms based upon object priors in training datasets.Denoising diffusion probabilistic models (DDPMs) [27] achieve state-of-the-art performance in generating realistic images. Their noise schedule in training may offer advantages for scaling up models in a stable manner. Recent works [2, 48, 70] explore representation learning capability in DDPMs. DatasetDDPM [2] examines few-shot segmentation with pre-trained diffusion models, but requires human labels to train a linear classifier. With the default U-Net architecture [51] , it loses the efficiency and flexibility of generating image and masks in a single-stage manner. Diff-AE [48] and PADE [70] perform representation learning driven by a reconstruction objective in the DDPM framework. Unfortunately, their learned latent vectors are not applicable to more challenging segmentation tasks and they require a pre-trained interpreter to perform downstream image classification.DiffuMask [63] takes a pre-trained Stable Diffusion model [50] , which is built using large-scale text-to-image datasets (and thus solves a far less challenging problem), and conducts a post-hoc investigation on how to extract segmentation from its attention maps. Neither our system, nor the baseline DDPM to which we compare, makes use of such additional information. Furthermore, DiffuMask does not directly output segmentation; it is basically a dataset generator, which produces generated images and pseudo labels, which are subsequently used to train a separate segmentation model. Our method, in contrast, is both completely unsupervised and provides an end-to-end solution by specifying an architectural design in which training to generate reveals segmentations as a bonus.MAGE [38] shares with us a similar motivation of framing generation and representation learning in a unified framework. However, our approach is distinct in terms of both (1) task: we tackle a more complex unsupervised segmentation task (without fine-tuning) instead of image classification (with downstream fine-tuning), and (2) design: 'masks' play a fundamentally different role in our system. MAGE adopts an MAE [24] -like masking scheme on input data, in order to provide a proxy reconstruction objective for self-supervised representation learning. Our use of region masks serves a different purpose, as they are integral components of the model being learned and facilitate factorization of the image generation process into parallel synthesis of different segments.BlobGAN [16] is a generative model for creating images with fine-grained control over the spatial arrangement of content. It leverages blob-like components instead of accurate region masks as basic building blocks for the synthesis process, allowing for intuitive content manipulation. In the generative modeling space, BlobGAN serves a different purpose than our method: BlobGAN excels in scenarios requiring explicit spatial control and interactive editing, while our factorized diffusion approach provides a framework for learning high-quality image generation and segmentation.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: [NA] Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
factorized diffusion models: 
 Figure 2 illustrates the overall architecture of our system, which partitions the denoising network within a diffusion model into an unsupervised region mask generator and parallel per-region decoders.To simultaneously learn representations for both image generation and unsupervised segmentation, we first design the region mask generator based on the first half (encoder) of a standard DDPM U-Net. We obtain input x t , a noised version of x 0 , via forward diffusion:EQUATIONwhere α t = 1 -β t , ᾱt = t s=1 α t . In addition to the encoder half of the U-Net, we instantiate a middle block consisting of layers operating on lower spatial resolution features. Parameterizing these subnetworks as θ enc and θ mid , we extract latent representations:EQUATIONEQUATIONwhere h enc encapsulates features at all internal layers of θ enc , for subsequent use as inputs, via skip connections, to corresponding layers of decoder-style networks (second half of a standard U-Net).We instantiate a mask generator, θ mask , as one such decoder-style subnetwork. A softmax layer produces an output tensor with K channels, representing K different regions in image x 0 :EQUATIONFollowing a U-Net architecture, h enc feeds into θ mask through skip-connections.We aim to extend a standard DDPM U-Net decoder θ dec to consider region structure during generation. One simple design is to condition on m = {m 0 , m 1 , ...} by concatenating it with input h mid and h enc along the channel dimension:EQUATIONwhere h mid and h enc are generated from Eq. 2 and Eq. 3. We downsample m accordingly to the same resolution as h mid and h enc at different stages. 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.To address these issues, we separate the decoding scheme into multiple parallel branches of weightshared U-Net decoders, each masked by a single segment. Noise prediction for k-th branch is:EQUATIONand the output is a sum of region-masked predictions:ε = K-1 k=0 εk ⊙ m k (7)We train our model in an end-to-end manner, driven by the simple DDPM denoising objective. Model weights θ = {θ enc , θ mid , θ dec , θ mask } are optimized by minimizing the noise prediction loss:EQUATIONUnlike previous work, our method does not require a mask regularization loss term [54, 4, 9] , which predefines mask priors (e.g., object size). Algorithm 1 summarizes training.Once trained, our model can both segment novel input images and synthesize images from noise.Real Image Segmentation. Given clean input image x 0 , we first sample a noisy version x t through forward diffusion in Eq. 1. We then perform one-step denoising by passing x t to the model. We collect the predicted region masks as the segmentation for x 0 using Eq. 4.Image and Mask Generation. Using reverse diffusion, our model can generate realistic images and their corresponding segmentation masks, starting from a pure noise input x T ∼ N (0, 1). Reverse diffusion predicts x t-1 from x t :EQUATIONEQUATION)where σ t is empirically set according to the DDPM noise scheduler. We perform T steps of reverse diffusion to generate an image. We also collect its corresponding masks using Eq. 4 when t = 1. Algorithm 2 summarizes this process. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 We evaluate on: (1) real image segmentation, (2) image and region mask generation, using Flower [45] , CUB [62] , FFHQ [31] , CelebAMask-HQ [36] , and ImageNet [53] . In addition to the design of flat set of K regions, we also conduct a preliminary investigation into reorganizing our architectural design to support hierarchical segmentations; see Section A.1.Evaluation Metrics. For unsupervised segmentation on Flower and CUB, we follow the data splitting in IEM [54] and evaluate predicted mask quality under three commonly used metrics, denoted as Acc., IOU and DICE score [54, 9] . Acc. is the (per-pixel) mean accuracy of the foreground prediction. IOU is the predicted foreground region's intersection over union with the ground-truth foreground region. DICE score is defined as 2 F ∩F | F | [14] . On ImageNet, we evaluate our method on Pixel-ImageNet [68] , which provides human-labeled segmentation masks for 0.485M images covering 946 object classes. We report Acc., IOU and DICE score on a randomly sampled subset, each class containing at most 20 images. For face datasets, we train our model on FFHQ and only report per-pixel accuracy on the CelebAMask test set, using provided ground-truth.For image and mask generation, we use Fréchet Inception Distance (FID) [25] for generation quality assessment. Since we can not obtain the ground-truth for generated masks, we apply a supervised U-Net segmentation model, pre-trained on respective datasets, to the generated images and measure the consistency between masks in terms of per-pixel accuracy. In addition to quantitative comparisons, we show extensive qualitative results. Implementation Details. We train Flower, CUB and Face models at both 64 × 64 and 128 × 128 resolution. We also train class-conditioned ImageNet models with 64 × 64 resolution. For all Acc. IOU DICE GrabCut [52] 82.0 69.2 79.1 ReDO [9] 87.9 76.4 -IEM [54] 88. experiments, we use the U-Net [51] encoder-middle-decoder architecture similar to [27] . We use the decoder architecture as our mask generator and set the number of factorized masks K as 3. We note that K is the maximum number of regions the model may use. It could learn fewer components during training. 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. For segmentation evaluation, we simply select the mask channel that emerges as foreground and apply standard benchmarks. We use Adam to train all the models with a learning rate of 10 -4 and an exponential moving average (EMA) over model parameters with rate 0.9999. For all datasets except ImageNet, we train 64 × 64 and 128 × 128 models on 8 and 32 Nvidia V100 32GB GPUs, respectively. For Flower, CUB and FFHQ, we train the models for 50K, 50K, 500K iterations with batch size of 128, respectively. For ImageNet, we train 500K iterations on 32 Nvidia V100 GPUs with batch size 512. We adopt the linear noise scheduler as in Ho et al. [27] with T = 1000 timesteps.To evaluate our method on real image segmentation, we set t as 30 for the forward diffusion process. We also investigate the segmentation results with different noise levels in Figure 18 . For Flower and CUB, Figures 3 and 4 show test images and predicted segmentations. Tables 1 and 2 provide quantitative comparison with representative unsupervised image segmentation methods: GrabCut [52] , ReDO [9] and IEM [54] . As shown in Table 1 and Table 2 , our method outperforms all competitors.We also visualize the predicted face parsing results on FFHQ and CelebAMask datasets in Figure 1 (c)(d) and Figure 5 . Our model learns to accurately predict three segments corresponding to semantic components: skin, hair, and background. This particular semantic partitioning emerges from our unsupervised learning objective, without any additional prior. With ground-truth provided on CelebAMask-HQ, we also compare the pixel accuracy and mean IOU with a supervised U-Net and DatasetDDPM [2] . For the former, we train a supervised segmentation model with 3-class cross-entropy loss. For the unsupervised setting, we perform K-means (K=3) on the pre-trained DDPM, denoted as DatasetDDPM-unsup. Table 3 shows that we outperform DatasetDDPM by a large margin and achieve a relatively small performance gap with a supervised U-Net. Figure 6 shows the accurate segmentation results for ImageNet classes: ostrich, pekinese, papillon, and tabby. We compare with supervised U-Net and DatasetDDPM-unsup in Table 4 . We show more visualizations in the Appendix.We evaluate our method on image and mask generation. As shown in Figure 7 , 8, 1(c)(d) and 9, our method is able to generate realistic images. In the upper row of Table 5 , we see a consistent quality improvement over the original DDPM. This suggests that our method is a better generation architecture than the standard U-Net; separate computational paths for denoising individual image regions is a beneficial prior to impose when learning to model the image distribution. Additionally, our method produces accurate corresponding masks, closely aligned with the semantic partitions in the generated image.We also evaluate the quality of these segmentations. Since there is no ground-truth mask provided for generated images, we apply the U-Net segmentation models (pre-trained on respective labeled training sets) to the generated images to produce reference masks. We measure the consistency between the reference and the predicted parsing results in terms of pixel-wise accuracy. We compare our method with a pre-trained DDPM baseline, in which we first perform image generation, then pass them to DatasetDDPM-unsup to get masks. As shown in Table 5 (bottom), our method consistently achieves better segmentation on generated images than the DDPM baseline. Note that, different from the two-stage baseline, our method performs the computation in a single stage, generating image and mask simultaneously. The Appendix shows more visualizations. Table 6 : Ablations of decoding scheme on CUB.Multi-branch Decoders with Weight Sharing. Separating computation in multi-branch decoders with weight sharing is an essential design in our method. We show the effectiveness of this design by varying how to apply factorized masks in our decoding scheme: (1) concat: we use single branch to take concatenation of h and m. (2) masking h mid : we use m to mask h mid instead of h enc .(3) w/o weight sharing: we train decoders separately in our design. Table 6 shows separate design consistently yields better visual features than other designs for CUB. This suggests that our design benefits from fully utilizing mask information in the end-to-end denoising task and avoids a trivial solution where masks are simply ignored.Investigation on Mask Factorization. Our architecture is able to generate factorized representations, each representing a particular segment of the input image. We show this by visualizing the individual channels from softmax layer output in our mask generator. As shown in Figure 10 , skin, hair, and background are separated in different channels.Mask Refinement along Diffusion Process. In the DDPM Markov process, the model implicitly formulates a mapping between noise and data distributions. We validate that this occurs for both images and latent region masks by visualizing image and mask generation along the sequential reverse diffusion process in Figure 11 . We observe gradual refinement as denoising steps approach t = 0.Zero-shot Object Segmentation. We evaluate zero-shot object segmentation on both PASCAL VOC 2012 [17] and DAVIS-2017 videos [47] . Baseline DDPM generation is not solved for these datasets when training from scratch without external large-scale datasets (e.g., LAION [55] used in Stable Diffusion [50] ). We directly adopt zero-shot transfer from our pre-trained ImageNet model by applying the conditional label mapping to VOC. Table 7 details the mapping rule. Figure 13 shows the accurate segmentation results for images of classes: aeroplane, monitor, person, and sofa from VOC. Since our method does not require any pixel labels, we evaluate the performance of each object class individually. We report pixel accuracy and mIOU of each class in VOC in Table 7 , which demonstrates that our method can achieve reasonably high performance. Our method achieves an accuracy of 0.78 and mIOU of 0.54 when averaging over all 20 classes. We also show video segmentation on DAVIS-2017 in Figure 14 and the Appendix, without any labeled video pre-training.Face Interpolation. We also investigate face interpolation on FFHQ. Similar to standard DDPM [27] , we perform the interpolation in the denoising latent space with 250 timesteps of diffusion. Figure 12 shows good reconstruction in both pixels and region masks, yielding smoothly varying interpolations across face attributes such as pose, skin, hair, expression, and background. 
conclusion: 
 We propose a factorized architecture for diffusion models that is able to perform unsupervised image segmentation and generation simultaneously, while being trained once, from scratch, for image generation via denoising alone. Using model architecture as a constraint, via carefully designed component factorization and parallel decoding schemes, our method effectively and efficiently bridges these two challenging tasks in a unified framework, without the need for fine-tuning or altering the original DDPM training objective. Our work is the first example of engineering an architectural bottleneck so that learning a desired end task becomes a necessary byproduct of training to generate.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. Our initial investigation into hierarchical extensions suggests a promising future path towards handling complex scenes. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We conduct a further investigation to reorganize our architectural design to support hierarchical mask factorization in place of a flat set of K regions. We formulate a hierarchical factorized diffusion architecture to progressively refine segmentation results from a coarse initial prediction to a fine, detailed final segmentation. This approach helps in capturing both global context and fine details in the segmentation task. As shown in Figure 15 , the first level replicates the factorized diffusion architecture depicted in Figure 2 to generate initial region masks m 0 0 , m 0 1 , ..., each applied on the noisy input for the next level factorized diffusion process. Each branch of the second level architecture generates finer representations of region masks m 1 0 , m 1 1 , ..., constructing the final denoising output as i m 0 i(h 0 i + j h 1 j m 1 j ) 2. This nested architectural design can be instantiated as repeated levels of factorized diffusion, which is a promising way to handle multiscale scenes. As a proof of concept, we experiment on the shape 3D dataset [5] with a 2-level hierarchy. We first visualize each level's region mask in Figure 16 . We observe that the first level generates a coarse segmentation, based on which, second-level factorized diffusion generates fine segmentations of 3D shapes. Figure 17 We show more segmentation results for Flower, CUB, FFHQ, CelebA and ImageNet. As shown in Figures 19, 20, 21, 22 , and 23, our method consistently predicts accurate segmentations for real image inputs.We show more generation results for Flower, CUB, FFHQ, and ImageNet (classes: flamingo, water buffalo, garbage truck, and sports car). As shown in Figures 24, 25, 26 , and 27, our method consistently produces images and masks with high quality. We provide more segmentation results of 'bicycle', 'chair', 'potted plant' and 'train' in Figure 28 .We provide more DAVIS-2017 video segmentation results of 'classic-car', 'dance-jump' in Figure 29 . Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.(a) Real Images (b) Segmentation 


Decision: Authors present an unsupervised diffusion model that can not only
generate images but also provide segmentations, provided the number of
masks required as input. The approach and the architecture seems quite
novel as also stated by some reviewers.

The main concern raised during the review was that the extent of
experiments could be larger. Authors could have done experiments with
higher number of classes, full resolution images and more
datasets. During the rebuttal authors note that the main issue is the
computational cost. Given the number of experiments and the ablation
study, I tend to agree with the authors. Some reviewers also agree with
this.

The last reviewer also mentions theoretical insights are not
provided. Authors' response, while not answering it fully, provided
insights. This reviewer gave a BA for this paper.

There are two reviewers who gave BR to this article. Their main
concers were about experiments and comparisons. The second BR
reviewer also mentions clarity and motivation. Authors did their best
to answer these concerns during the rebuttal. Unfortunately, reviewer
did not engage in a discussion. However, reading the rebuttal, I think
authors response were valuable.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - Experiments only cover 2-3 class scenarios.
- ImageNet results only on downsampled 64x64 images.
questions: 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?
limitations: limitations: Limitations are discussed explicitly in the paper. No concerns about negative societal impact.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.)?
limitations: limitations: The ethical implications of using generative models to create harmful content should be addressed.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: see above the comments about the concatenation of ℎ𝑚𝑖d and ℎenc.
limitations: limitations: yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: 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.
limitations: 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.

Comments: title: To Reviewer FyWe
To Reviewer 9WST
Rebuttal is online - please respond
comment: **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.
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.
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?
> 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,
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,