Panoptic Diffusion Models: co-generation of images and segmentation maps

One of the initial works in this area, Denoising Diffusion Probabilistic Models (DDPM), use a Markov chain to gradually add scheduled noises to images in the forward process (Ho, Jain, and Abbeel 2020). The transition of the Markov chain is then parameterized by a neural network trained to predict the noise. During inference, a diffusion model starts from random noise and gradually reverses it to reconstruct the image.

A well-known drawback of diffusion models is that they require a large number of steps to generate samples iteratively. To address this issue and improve efficiency, researchers have proposed various modifications to diffusion models (Nichol and Dhariwal 2021). For instance, DDIM demonstrates that diffusion models can operate in a non-Markovian manner, resulting in shorter generative chains (Song, Meng, and Ermon 2021). Additionally, distillation algorithms have been introduced to further accelerate the multi-step inference process by progressively distilling a teacher model into a student model(Salimans and Ho 2022; Berthelot et al. 2023; Ren et al. 2024).

The backbone neural network for a diffusion model is typically a UNet, which is composed of convolutional layers and attention blocks, or a diffusion transformer that relies solely on attention mechanisms (Rombach et al. 2022; Peebles and Xie 2022). Another variant, UViT, is a type of diffusion transformer that retains skip connections, allowing later layers to access information from earlier layers, thereby enhancing alignment (Bao et al. 2023).

There are three main methods for applying conditions to a diffusion model. The first approach, used in stable diffusion, involves cross-attention between the image and the conditions (Rombach et al. 2022). The second method appends condition embeddings as tokens to the image patches (Bao et al. 2023). The third approach uses an adaptive norm layer to integrate conditions with the hidden states (Peebles and Xie 2022). In our panoptic diffusion models, we opt for the second method because the transformer can leverage self-attention to learn the relationships between images and maps, treating them as conditions for each other.

The solver for our panoptic diffusion model is a modified version of DPM Solver++ (Lu et al. 2023). Solving the reverse of the diffusion process is equivalent to solving an ordinary differential (ODE) equation, which can be decoupled as an exactly computed linear part and a non-linear part approximated by neural networks (Lu et al. 2022).

During inference, we apply classifier-free guidance similar to Nichol et al. (2021) and Ho and Salimans (2022). The diffusion model runs twice, once in an unconditioned setting and once in a conditioned setting, and the final output is obtained by taking a weighted sum of the unconditioned and conditioned outputs.

Object detection requires generating bounding boxes and fine-grained masks, tasks traditionally accomplished by convolutional neural networks such as Fast R-CNN (Girshick 2015) and Mask R-CNN (He et al. 2017). In Carion et al. (2020), researchers introduced the use of transformers to generate binary masks by inputting object queries. Building on this, Cheng et al. (2022) proposed a collaboration between an image encoder backbone and a masked transformer to generate masks, where masked attention replaces cross attention.

Recently, there has been growing interest in applying diffusion models to panoptic segmentation masks. For example, in Chen et al. (2023), a diffusion model comprising an image encoder and a mask decoder is used to extract image features and apply cross attention between these features and the masks. To address the challenge of handling discrete data with diffusion models, Chen, Zhang, and Hinton (2022) proposed converting panoptic masks into analog bits during preprocessing. On the other hand, Baranchuk et al. (2021) suggest that the intermediate features of diffusion models can capture semantic information useful for label-efficient segmentation. Similarly, DiffuMask (Wu et al. 2024) generates a synthetic image and a corresponding segmentation mask of an object using attention maps. However, directly extracting masks from attention maps lacks the ability to control the generated image in return. In contrast, our approach aims to co-generate pixel-level panoptic segmentation maps and images, allowing them to influence and control each other.

While previous studies use diffusion models for panoptic segmentation based on given images, our work leverages an additional dataset of panoptic maps to train a model capable of generating both maps and images. The generated maps are then used to condition the image generation, producing a photorealistic result.

Image guided image generation enables more precise control over the structure of the image and ensures faithfulness to users’ illustrative inputs. The input for guidance can have various forms, such as segmentation maps and layouts (Rombach et al. 2022; Zhang, Rao, and Agrawala 2023). Stochastic Differential Editing (SDEdit) perturbs user inputs with Gaussian noises and then synthesizes images by reversing SDE (Meng et al. 2022). They show that when the reverse SDE is not solved from the ending point but a particular timestep, the generated images can achieve a good balance between faithfulness and realism. Make-a-scene introduces scene-based conditioning for image generation by optionally providing tokens from segmentation maps (Gafni et al. 2022), but this method heavily relies on explicit strategies for tackling panoptic, human, and face semantics. SpaText (Avrahami et al. 2023) employs CLIP (Radford et al. 2021) to convert local text prompts that describe segments into image space and concatenate to the channel dimension of noises. ControlNet can accept user inputs such as canny edges and segmentation masks for conditional control of image generation (Zhang, Rao, and Agrawala 2023). Prompt-to-prompt image editing controls the generation by cross-attention to ensure similarity between images generated from similar prompts (Hertz et al. 2022). InstructPix2Pix combines Prompt-to-prompt method with stable diffusion to generate pairs of images from pairs of captions for training, then train the model to modify image pixels following the instructions (Brooks, Holynski, and Efros 2023).

These approaches demonstrate that providing various forms of guidance can more accurately control the structure of generated images. Building on this insight, our method assumes that such guidance is crucial for enhancing image quality. Additionally, panoptic diffusion models inherently generate segmentation maps alongside images, offering built-in guidance without the need for additional user input beyond the text prompt.

To reduce the number of trained parameters or adapt the model to a new domain, previous works have designed adaptive blocks to fine-tune convolutional neural networks or transformers (Houlsby et al. 2019; Long et al. 2021; Mou et al. 2023). In our two-stream panoptic diffusion model, the map stream functions similarly to an adapter. To prevent any negative impact on the pretrained weights, we employ zero-initialized convolutional blocks as proposed in Zhang, Rao, and Agrawala (2023).

As shown in Fig. 2,we process the panoptic segmentation maps through several steps before feeding them into the diffusion model. Instead of using a binary mask for each object, we load pixel-level panoptic annotations. In a segmentation map $M_{0}$, each pixel’s value is set to the corresponding category ID if it belongs to a segment; otherwise, its value is zero. We then convert these pixel values into analog bits (Chen, Zhang, and Hinton 2022). Analog bits are necessary because a standard diffusion model can only generate continuous data, while segmentation classes are discrete or categorical. Since the range of category ID is from 1 to 200, each pixel is represented by 8 binary bits. Prior to noise scheduling, these bits are scaled to the range $[-1,1]$, matching the range of the latent input to the diffusion model. To ensure that the noise can effectively flip the bits, its absolute value must exceed one. Therefore, we set the noise added to the maps as $\epsilon_{M}\sim\mathcal{N}(0,2*\mathbf{I})$.

Latent diffusion models use latent representations of images encoded by an autoencoder as inputs. However, no autoencoder exists for encoding and decoding high-resolution segmentation maps into latents. We address this issue by pooling and using a larger patch size for the maps. To achieve high-resolution maps and enable more precise control, we first pool the maps to match one, two, or four times the height and width of the image latents. We use min pooling to prioritize smaller category numbers, as the COCO dataset annotations categorize 1-91 as thing categories and 92-200 as stuff categories. Next, we set the patch size of the maps to be one, two, or four times that of the images. This approach ensures that, after patchifying, the sizes of the image and map features align. Given that images have three RGB channels while maps have only one channel for the category ID before preprocessing, using a larger patch size is effective for extracting hidden features from segmentation maps. Consequently, this method allows us to generate higher-resolution maps without the need for an additional autoencoder or a larger latent size.

For postprocessing, the output values predicted by the diffusion model are thresholded at zero. Negative values are treated as zero bits, while positive values are considered one bits. Subsequently, these output bits are converted back into category numbers.

In the forward pass of the diffusion process (Ho, Jain, and Abbeel 2020), random noise $\epsilon\sim\mathcal{N}(0,\mathbf{I})$ is added to the image latent $x_{0}$ according to the noise scheduler. With a total of $n$ steps, each step updates the noisy image $x_{t}$ from the previous step $x_{t-1}$, using scaling factors $\alpha$ and $\beta$ provided by the noise scheduler. This process forms a Markov chain. Consequently, the noisy image $x_{t}$ can be simplified and calculated directly from $x_{0}$.

where $\alpha_{t}$ are close to 1 and $\beta_{t}=1-\alpha_{t}$. The cumulative factor $\bar{\alpha}=\prod_{i=1}^{t}\alpha_{i}$, and the noise is scaled by $\sigma_{t}=\sqrt{1-\bar{\alpha}}$.

To learn to denoise panoptic segmentation maps, we create another random Gaussian noise $\epsilon_{M}\sim\mathcal{N}(0,2*\mathbf{I})$ and add it to the ground-truth maps $M_{0}$. The same noise scheduler is used to add noises to maps.

where $M_{t}$ is the noised map at timestep $t$.

The panoptic diffusion model outputs $\epsilon_{\theta}$, which estimates the noise $\epsilon$. Using this estimated noise, we compute the predicted image $\tilde{x_{0}}$. When incorporating the map as an additional input to the diffusion model, the equation for predicting the image is given by Eq. 4. To accelerate inference, we utilize a fast DPM solver to compute $x_{t_{i-1}}$ from $x_{t_{i}}$ (Lu et al. 2022, 2023). By using discontinuous time steps $t_{i}$ and $t_{i-1}$, this method can skip intermediate steps, reducing the total number of sampling steps required. The first-order solver is described in Equation 5, where $h_{i}$ represents the difference in the log signal-to-noise ratio between different steps ($h_{i}=\log(\alpha_{t_{i}}/\sigma_{t_{i}})-\log(\alpha_{t_{i-1}}/\sigma_{t_{i-1}})$). Details on a third-order solver can be found in Appendix A.

The other output of a panoptic diffusion model is $M_{\theta}$, which is a prediction of $M_{0}$. Drawing inspiration from DPM-solver++, we use the following equation to estimate $M_{t_{i-1}}$ from the previous step. It is important to note that the model directly estimates $M_{0}$ rather than the noise added to the segmentation map, as predicting $\epsilon_{M}$ does not provide effective guidance for the images. By training the diffusion model with panoptic segmentation maps, it incorporates intrinsic self-control into the image generation process.

In a special case where ground truth maps are provided as conditions, the diffusion model will focus solely on predicting the images. This allows users to have customized control for generating desired images, similar to existing methods (Zhang, Rao, and Agrawala 2023). However, this approach limits the diversity of the generated images.

Since the generation of $x_{t-1}$ and $M_{t-1}$ relies on $x_{t}$ and $M_{t}$, they form a dual problem. Improvements in the quality of the generated masks and images influence each other. Consequently, according to the scaling law, a larger diffusion model can produce more accurate masks, which in turn provides better control and further enhances image quality.

Let the inputs to a panoptic diffusion model at each timestep be image latent $x_{t}$, mask $M_{t}$, text condition encoded by a text encoder $C$, and timestep $t$. The conditional probability of $x_{t-1}$ and $M_{0}$ is given by

Equation 6 show that it is feasible to predict the segmantation map $M_{0}$ first, then use it as a condition to predict $x_{t-1}$. However, when using a unified model to predict both $x_{t-1}$ and $M_{0}$, the intermediate features already contain the segmentation information used to predict $M_{0}$. Through self-attention, the map features can inherently condition $x_{t-1}$. Therefore, it is reasonable to predict $x_{t-1}$ and $M_{0}$ simultaneously. By taking the logarithm of the probability, we can optimize the model by combining the losses associated with image denoising and segmentation map generation.

The training algorithm is outlined in Algorithm 1. We use Mean Squared Error (MSE) loss to optimize the predicted noises for both image and segmentation map denoising. Specifically, the target for image denoising is the noise $\epsilon$, while the target for mask generation is the ground-truth $M_{0}$. The losses for images and maps are summed to perform gradient backpropagation. During inference, the diffusion model iteratively denoises both images and maps, as detailed in Algorithm 2.

We propose a new metric to evaluate the quality of generated segmentation maps by measuring the difference in the number of pixels labeled as each category. While Panoptic Quality (Kirillov et al. 2018) uses Intersection over Union (IoU) to assess segmentation maps based on the weighted sum of true positives, false positives, and false negatives, this approach is not suitable for maps generated by diffusion models. These models produce maps probabilistically based on text prompts and co-generated images, making it impractical to compute IoU with ground-truth maps due to inherent differences in the generated images. Instead, we introduce the Mean Count Difference (MCD) metric. MCD evaluates the quality of generated maps by counting the frequency $f$ of each category in both the ground-truth and generated maps, then summing their absolute differences. This sum is divided by the total number of pixels, calculated as the product of the height and width. Given that object locations on the generated map are not fixed, comparing category frequencies rather than direct pixel values provides a more meaningful assessment. The metric ranges from $[0,2]$, where zero indicates identical segmentation maps and larger values indicate greater differences.

We evaluate the quality of generated images using FID (Heusel et al. 2017) and CLIP scores (Hessel et al. 2022). FID assesses the quality and fidelity of the generated images by employing an Inception model, while CLIP scores gauge how well the generated images correspond to the text prompts. For CLIP scores, we use the ViT-B/32 model (Radford et al. 2021). We generate 30,000 images and segmentation maps from 5,000 text files in the COCO dataset’s validation set, with each text file containing five captions describing the same scene. We compute the average CLIP scores by comparing these five captions with the generated images.

In Table.1, we compare the FID and CLIP scores of our models with those of state-of-the-art methods. The results indicate that while our panoptic diffusion models (PDMs) are trained with a combined loss of images and segmentation maps, they achieve comparable fidelity (FID scores) and improved relevance between image and text (higher CLIP scores). This improvement is due to the enhanced connectivity between the image, text, and segmentation map. The two-stream PDM performs better due to its pretrained stream and larger number of parameters. When ground-truth maps are provided, the model performs optimally because it focuses solely on optimizing image generation.

Increasing the patch factor results in a higher MCD because generating higher-resolution maps with a fixed number of latents becomes more challenging. This creates a trade-off between map resolution and quality. We find that a patch factor of 2 offers the best balance, yielding the highest FID and CLIP scores. However, increasing the patch factor to 4 results in worse performance, suggesting that unbalanced patch sizes for maps and images are detrimental.

In Fig. 1, we compare the images and masks generated by PDM with images generated by U-ViT. By training with segmentation masks, PDM learns that the shape of a stop sign should be octagon, while U-ViT cannot guarantee to generate an octagon stop sign. Similarly, PDM ensures to generate correct shapes for a fire hydrant and a human. In the last row of Fig. 1, PDM generates masks for not only elephants but also for the river, while a regular diffusion model misses the required component of the text prompt.
Figure 4 displays images generated with either ground-truth segmentation maps or co-generated maps. The generated maps in the bottom left show objects of the same categories and similar shapes as the ground-truth maps. The images on the right are conditioned on these segmentation maps, demonstrating the PDM’s ability to simultaneously generate correlated images and maps. While images generated with ground-truth maps exhibit slightly better quality, co-generation removes the need for a segmentation input and produces diverse maps and images. Note that the pixel values in the generated segmentation maps correspond to category IDs (1-200), which are mapped to random RGB colors for visualization. The color map used is detailed in Appendix C.

Additional examples generated by PDMs are provided in Appendix B. Zero-shot results on the CIFAR-10 dataset demonstrate that our model can generate segmentation maps for various categories across different image datasets.