JoDiffusion: Jointly Diffusing Image with Pixel-Level Annotations for Semantic Segmentation Promotion

AI summary

Overview

Research area: Computer Vision — synthetic training-data generation for semantic segmentation using latent diffusion models.

Technical level: Advanced. The paper assumes familiarity with latent diffusion models, variational auto-encoders (VAEs), text encoders, and standard segmentation benchmarks.

Scope: The paper proposes JoDiffusion, a diffusion framework that generates a synthetic semantic segmentation dataset — images paired with pixel-level annotation masks — from text prompts alone, and evaluates the resulting dataset by training standard segmenters on it.

What This Paper Is About

Training semantic segmentation models requires images paired with pixel-level masks, and producing those masks by hand is expensive and slow. Existing workaround pipelines are flawed: Image2Mask methods generate an image first and then predict a pseudo-mask (which often disagrees semantically with the image), while Mask2Image methods generate images from human-drawn masks (which caps diversity at whatever masks already exist). JoDiffusion instead learns the joint distribution of image and mask together, so a text prompt alone yields a semantically consistent image-and-mask pair at scale.

Key Contributions

1. A synthetic dataset generation framework for semantic segmentation that, per the authors, is the first attempt to simultaneously generate semantically consistent paired images and pixel-level annotation masks conditioned solely on text prompts.
2. An independent annotation VAE that maps annotation masks into the same latent space as images, using binary encoding of category indices and a lightweight architecture (roughly 50M parameters versus roughly 300M for the image VAE) trained without KL regularization and with cross-entropy loss only.
3. A joint diffusion scheme built on Unidiffuser that applies a shared noise perturbation to the image latent and the mask latent, so the denoising network predicts one joint noise term rather than two independent ones.
4. A boundary-mode mask optimization strategy that corrects small, noisy label regions by reassigning them to the most frequent category among their boundary pixels, using a threshold τ.

Main Findings

• Annotation VAE reconstruction is near-lossless: Reconstruction mIoU of the generated masks on validation sets is 99.50 on Pascal VOC, 98.85 on MS COCO, and 98.74 on ADE20K — the authors report this as exceeding 98% accuracy in compactly encoding masks.
• Synthetic-only training gains are large: With DeepLabV3 and a ResNet50 backbone, training on synthetic data alone gives 72.5 mIoU on Pascal VOC (40k generated samples) versus 61.6 for Dataset Diffusion (40k) and 60.4 for SDS (26k); on MS-COCO (80k generated samples) JoDiffusion reaches 42.6 versus 32.4 for Dataset Diffusion (80k) and 31.0 for SDS (50k).
• Gains persist with ResNet101: Synthetic-only Pascal VOC reaches 75.8 versus 64.8 (Dataset Diffusion, 40k) and 59.1 (SDS, 26k); MS-COCO reaches 44.9 versus 34.2 (Dataset Diffusion, 80k) and 31.8 (SDS, 50k).
• Mask2Former results follow the same pattern: With a ResNet50 Mask2Former, synthetic-only Pascal VOC is 74.5 versus 60.2 (Dataset Diffusion, 40k) and 59.8 (SDS, 26k); MS-COCO is 44.6 versus 31.0 (Dataset Diffusion, 80k) and 29.8 (SDS, 50k).
• Real-plus-synthetic training improves over real-only: On Pascal VOC with DeepLabV3-ResNet50, real-only training gives 77.4 mIoU at 11.5k real images; JoDiffusion's real-plus-synthetic gives 78.3. On MS-COCO, real-only is 48.9 at 118k real images and JoDiffusion's real-plus-synthetic is 56.4.
• Comparison against Mask2Image methods: On Pascal VOC with ResNet50, Raw Data gives 77.3, FreeMask (reproduced, 40k) gives 77.9, and JoDiffusion gives 79.4; on ADE20K with ResNet50, Raw Data gives 47.2, SegGen gives 49.9 with 1M data, FreeMask (reproduced, 40k) gives 48.2, and JoDiffusion (40k) gives 48.4.
• Swin-S backbone comparison: Pascal VOC Raw Data 83.8, FreeMask (reproduced, 40k) 84.2, JoDiffusion 85.1; ADE20K Raw Data 51.6, FreeMask (reproduced, 40k) 52.1, JoDiffusion 52.2.
• Mask optimization helps, and τ = 20 is best in the sweep: On Pascal VOC, τ = 0 gives 71.37 mIoU, τ = 20 gives 72.47, τ = 50 gives 72.38, and τ = 100 gives 72.38.
• More generated data consistently helps: On Pascal VOC, 5k samples give 68.54 mIoU, 10k give 70.02, 20k give 70.97, and 40k give 72.47.
• The paper does not report generative fidelity metrics such as FID for the synthetic images, nor does it report a separate ablation isolating the annotation VAE's contribution beyond its reconstruction mIoU.

Methodology in Plain English

The method runs in three stages.

Stage one trains an annotation VAE. Because annotation masks store a single category index per pixel, and adjacent index values would look too similar to a network, each category is converted into a binary representation before encoding. The encoder stacks a few convolutional and downsampling layers (channels growing from 128 to 256 to 512 to 512), and the decoder uses three transposed-convolution upscaling blocks to reconstruct the mask, which is then read out with an argmax over the softmax output. Since this VAE is only a compression tool and not a generative model, no standard normal prior and no KL divergence term are used — only cross-entropy loss.

Stage two trains the diffusion model jointly. An image is captioned automatically (BLIP-2 on Pascal VOC and ADE20K; COCO Caption on MS-COCO). Text, image and mask are each encoded into a latent space, and the image latent and mask latent are diffused together using one shared Gaussian noise perturbation. During denoising, the network takes the noisy image latent, the noisy mask latent, the text condition and the timestep, and predicts that single joint noise term, trained with mean squared error loss. This bidirectional interaction is why the image and mask stay aligned instead of one being derived from the other afterward.

Stage three is post-processing. Because diffusion can leave speckled or inconsistent labels, especially in tiny regions and along boundaries, any region smaller than a threshold τ has all of its pixels relabeled to the most frequent category found among its boundary pixels. The authors frame this as a maximum-likelihood-style estimate, on the assumption that boundary pixels are more likely to be correct than interior pixels of a small noisy blob.

The resulting image-and-mask pairs are then used to train standard segmentation models. Images and masks are resized to 512×512; AdamW is used throughout with random horizontal flipping as augmentation. At generation time, DPMSolverMultistep is used with 50 steps. On Pascal VOC 40,000 samples are generated, on MS-COCO 80,000, and on ADE20K 40,000.

Why This Matters

Impact on research: The paper reframes segmentation dataset synthesis as a joint-distribution problem rather than a two-step prediction problem, and shows that conditioning on text alone removes the Mask2Image bottleneck of depending on existing hand-drawn masks. The synthetic-only numbers — for example 72.5 versus 61.6 mIoU on Pascal VOC with DeepLabV3-ResNet50 — suggest the quality of generated supervision matters more than the sheer count of generated images.

Real-world applications:

• Autonomous driving, where pixel-level road-scene annotation is costly and rare conditions (unusual weather, rare obstacles) can be synthesized on demand.
• Medical image analysis, where expert pixel-level annotation is expensive and privacy rules restrict sharing real patient data.
• Robot navigation, where a robot needs to parse scenes and where training environments can be described in text.
• Any dense-prediction domain with many fine-grained categories and sparse annotation budgets, such as indoor scene parsing.

Industry relevance: The pipeline produces both the image and its label in one pass, which means a data-hunger problem can be attacked with text prompts instead of annotation labor. The lightweight annotation VAE (roughly 50M versus roughly 300M parameters for the image VAE) and the released code at https://github.com/00why00/JoDiffusion lower the practical barrier to reproducing or adapting the method.

Future Directions

• Extend the joint generation framework beyond the three benchmarks studied (Pascal VOC, MS-COCO, ADE20K) to other dense-prediction tasks and label types, such as instance or panoptic annotations.
• Reduce reliance on a real annotated dataset for supervision during training of the generative model, since the current setup still depends on real image-mask pairs.
• Improve caption generation quality, as the pipeline currently uses BLIP-2 for Pascal VOC and ADE20K and COCO Caption for MS-COCO — differences in caption source and any effect on the generated data are not analyzed.
• Replace the dataset-dependent mask optimization threshold τ with an adaptive or learned criterion; the best value found in the sweep (τ = 20) was tuned per dataset (τ = 20 was also used for generation on all three datasets).

Target Audience

Researchers and graduate students working on diffusion-based data generation, semantic segmentation, or dataset distillation; practitioners who need to scale dense-label training data without proportional annotation cost; and engineers evaluating whether synthetic supervision can substitute for or augment real annotation. Readers without prior exposure to latent diffusion models and VAEs will find the methodological sections dense.

Authors’ abstract

Given the inherently costly and time-intensive nature of pixel-level annotation, the generation of synthetic datasets comprising sufficiently diverse synthetic images paired with ground-truth pixel-level annotations has garnered increasing attention recently for training high-performance semantic segmentation models. However, existing methods necessitate to either predict pseudo annotations after image generation or generate images conditioned on manual annotation masks, which incurs image-annotation semantic inconsistency or scalability problem. To migrate both problems with one stone, we present a novel dataset generative diffusion framework for semantic segmentation, termed JoDiffusion. Firstly, given a standard latent diffusion model, JoDiffusion incorporates an independent annotation variational auto-encoder (VAE) network to map annotation masks into the latent space shared by images. Then, the diffusion model is tailored to capture the joint distribution of each image and its annotation mask conditioned on a text prompt. By doing these, JoDiffusion enables simultaneously generating paired images and semantically consistent annotation masks solely conditioned on text prompts, thereby demonstrating superior scalability. Additionally, a mask optimization strategy is developed to mitigate the annotation noise produced during generation. Experiments on Pascal VOC, COCO, and ADE20K datasets show that the annotated dataset generated by JoDiffusion yields substantial performance improvements in semantic segmentation compared to existing methods.