Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MAISI-v2: Accelerated 3D High-Resolution Medical Image Synthesis with Rectified Flow and Region-specific Contrastive Loss

Can Zhao, Pengfei Guo, Dong Yang, Yucheng Tang, Yufan He, Benjamin Simon, Mason Belue, Stephanie Harmon, Baris Turkbey, Daguang Xu

TL;DR

MAISI-v2 tackles the bottlenecks of 3D medical image synthesis by integrating rectified flow with a latent diffusion framework to dramatically accelerate generation while enhancing condition fidelity via a region-specific loss. It expands training data, employs a three-stage mixed-precision training on large GPUs, and introduces memory-aware controls for large volumes. It achieves state-of-the-art image quality with about $33\times$ acceleration on a 3D CT dataset and demonstrates that synthetic images can improve downstream segmentation, with open-source release.

Abstract

Medical image synthesis is an important topic for both clinical and research applications. Recently, diffusion models have become a leading approach in this area. Despite their strengths, many existing methods struggle with (1) limited generalizability that only work for specific body regions or voxel spacings, (2) slow inference, which is a common issue for diffusion models, and (3) weak alignment with input conditions, which is a critical issue for medical imaging. MAISI, a previously proposed framework, addresses generalizability issues but still suffers from slow inference and limited condition consistency. In this work, we present MAISI-v2, the first accelerated 3D medical image synthesis framework that integrates rectified flow to enable fast and high quality generation. To further enhance condition fidelity, we introduce a novel region-specific contrastive loss to enhance the sensitivity to region of interest. Our experiments show that MAISI-v2 can achieve SOTA image quality with $33 \times$ acceleration for latent diffusion model. We also conducted a downstream segmentation experiment to show that the synthetic images can be used for data augmentation. We release our code, training details, model weights, and a GUI demo to facilitate reproducibility and promote further development within the community.

MAISI-v2: Accelerated 3D High-Resolution Medical Image Synthesis with Rectified Flow and Region-specific Contrastive Loss

TL;DR

MAISI-v2 tackles the bottlenecks of 3D medical image synthesis by integrating rectified flow with a latent diffusion framework to dramatically accelerate generation while enhancing condition fidelity via a region-specific loss. It expands training data, employs a three-stage mixed-precision training on large GPUs, and introduces memory-aware controls for large volumes. It achieves state-of-the-art image quality with about acceleration on a 3D CT dataset and demonstrates that synthetic images can improve downstream segmentation, with open-source release.

Abstract

Medical image synthesis is an important topic for both clinical and research applications. Recently, diffusion models have become a leading approach in this area. Despite their strengths, many existing methods struggle with (1) limited generalizability that only work for specific body regions or voxel spacings, (2) slow inference, which is a common issue for diffusion models, and (3) weak alignment with input conditions, which is a critical issue for medical imaging. MAISI, a previously proposed framework, addresses generalizability issues but still suffers from slow inference and limited condition consistency. In this work, we present MAISI-v2, the first accelerated 3D medical image synthesis framework that integrates rectified flow to enable fast and high quality generation. To further enhance condition fidelity, we introduce a novel region-specific contrastive loss to enhance the sensitivity to region of interest. Our experiments show that MAISI-v2 can achieve SOTA image quality with acceleration for latent diffusion model. We also conducted a downstream segmentation experiment to show that the synthetic images can be used for data augmentation. We release our code, training details, model weights, and a GUI demo to facilitate reproducibility and promote further development within the community.

Paper Structure

This paper contains 29 sections, 5 equations, 4 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. MAISI: Medical AI for Synthetic Imaging
  4. Rectified Flow
  5. Methodology
  6. Overview of MAISI-v2
  7. Training Data
  8. Rectified Flow-based LDM Training
  9. ControlNet Training
  10. Region-specific Contrastive Loss
  11. Memory-Aware Region-specific Contrastive Loss
  12. Final Objective
  13. Quality Check with CT HU Intensities
  14. Experiments
  15. Training Data and Implementation Details
  16. ...and 14 more sections

Figures (4)

  • Figure 1: Region-specific Contrastive Loss
  • Figure 2: Qualitative comparison across axial (top row), sagittal (middle row), and coronal (bottom row) views. Columns correspond to different methods. MAISI-DDPM and MAISI-v2 in this figure are unconditional synthesis which do not use ControlNet or segmentation maps.
  • Figure 3: MAISI-v2 segmentation-guided results for small to large volume size and three different regions.
  • Figure 4: MAISI-v2 segmentation-guided results for five types of tumors. We show results for different voxel spacing and volume size to demonstrate the flexibility of MAISI-v2. Different Hounsfield Unit window is used to better show the contrast between tumor and normal tissues.