VolDiT: Controllable Volumetric Medical Image Synthesis with Diffusion Transformers

TL;DR

VolDiT introduces a fully transformer-based 3D diffusion model for medical image synthesis, offering improved global coherence, fidelity, and controllability over traditional U-Net based methods.

Contribution

It is the first to develop a purely transformer-based 3D diffusion model for volumetric medical image synthesis with structured control capabilities.

Findings

Outperforms state-of-the-art U-Net based models in global coherence and fidelity.

Enables precise spatial control via token-level conditioning.

Demonstrates superior controllability and scalability in 3D medical image generation.

Abstract

Diffusion models have become a leading approach for high-fidelity medical image synthesis. However, most existing methods for 3D medical image generation rely on convolutional U-Net backbones within latent diffusion frameworks. While effective, these architectures impose strong locality biases and limited receptive fields, which may constrain scalability, global context integration, and flexible conditioning. In this work, we introduce VolDiT, the first purely transformer-based 3D Diffusion Transformer for volumetric medical image synthesis. Our approach extends diffusion transformers to native 3D data through volumetric patch embeddings and global self-attention operating directly over 3D tokens. To enable structured control, we propose a timestep-gated control adapter that maps segmentation masks into learnable control tokens that modulate transformer layers during denoising. This token-level conditioning mechanism allows precise spatial guidance while preserving the modeling advantages of transformer architectures. We evaluate our model on high-resolution 3D medical image synthesis tasks and compare it to state-of-the-art 3D latent diffusion models based on U-Nets. Results demonstrate improved global coherence, superior generative fidelity, and enhanced controllability. Our findings suggest that fully transformerbased diffusion models provide a flexible foundation for volumetric medical image synthesis. The code and models trained on public data are available at https://github.com/Cardio-AI/voldit.