Med3D-JADE: 3D Joint Attentive Diffusion Engine for Volumetric Medical CT and Mask Co-generation

ICLR 2026 Conference Submission17483 Authors

Keywords: Medical imaging, Data Augmentation, joint generation, diffusion model, 3D

Abstract: Data scarcity is a critical bottleneck for training robust 3D medical image segmentation models. Current generative approaches for paired data synthesis are often limited to conditional generation (e.g., mask-to-image), which cannot produce novel anatomical structures and thus fail to address the lack of structural diversity in training data. To overcome this, we introduce Med3D-JADE, the first diffusion-based framework, to our knowledge, that learns the true joint distribution p(image, mask) to simultaneously generate entirely new 3D CT volumes and their corresponding segmentation masks. Our method adapts a pre-trained 3D medical generation foundation model (MAISI) into a dual-branch latent diffusion architecture. We preserve the foundation model's high-fidelity image synthesis by freezing its original branch and training a new, parallel branch for segmentation. Our proposed Volumetric Joint Attention (VJA) modules enforce coherence between the modalities, while the reuse of the model's powerful Volumetric Compression Network (VCN) facilitates efficient, high-resolution generation for both domains without needing to train a new encoder from scratch. We rigorously validate our approach by using the generated pairs for data augmentation across four datasets, including public benchmarks like SegTHOR and challenging MSD tumor datasets. Augmenting with our synthetic data leads to significant performance gains for diverse segmentation models like nnU-Net, SwinUNETR, and SegResNet, establishing Med3D-JADE as a generalizable and practical solution for overcoming 3D data scarcity in medical imaging.

Primary Area: applications to computer vision, audio, language, and other modalities

Submission Number: 17483