VolDiT: Controllable Volumetric Medical Image Synthesis with Diffusion Transformers

2026-03-26

AI Summary

This work proposes VolDiT, the first pure Transformer-based 3D diffusion model for medical image synthesis, addressing the limitations of conventional convolutional U-Net architectures that struggle to capture global context due to their restricted receptive fields. VolDiT processes 3D volumes directly through voxel patch embeddings and leverages global self-attention to model long-range dependencies. To enable precise structural control, the model introduces a timestep-gated adapter that maps segmentation masks into learnable control tokens, facilitating token-level conditional modulation. Experiments demonstrate that VolDiT significantly outperforms U-Net baselines in high-resolution 3D medical image generation, achieving notable advances in global consistency, synthesis fidelity, and spatial controllability.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• 3D medical image synthesis
• diffusion models
• transformer architecture
• global context
• controllable generation

Innovation

• Diffusion Transformer
• 3D Medical Image Synthesis
• Volumetric Patch Embedding
• Global Self-Attention
• Token-level Conditioning