MAISI-v2: 3D Medical Image Synthesis

• The paper introduces MAISI-v2, a framework combining rectified flow acceleration and region-specific contrastive loss to generate anatomically plausible 3D synthetic images.
• It employs a deterministic ODE in a latent diffusion architecture, reducing inference from 198 seconds to 6 seconds while achieving state-of-the-art FID scores.
• The method enhances downstream augmentation utility, improving segmentation Dice scores by up to 7.5% across multiple tumor types.

MAISI-v2 is a 3D high-resolution medical image synthesis framework that combines rectified flow–based acceleration with a region-specific contrastive loss to produce anatomically plausible and condition-consistent synthetic data at significantly reduced inference times. The design addresses three principal challenges in medical image synthesis with diffusion models: restricted generalizability to specific body regions or voxel spacings, prohibitive inference latency, and weak conditioning alignment. By leveraging rectified flow within a latent diffusion architecture and introducing a contrastive objective focused on regions of interest (ROI), MAISI-v2 establishes state-of-the-art fidelity with a substantial throughput advantage, and improves synthetic utility in downstream data augmentation (Zhao et al., 7 Aug 2025).

1. Rectified Flow Integration for Fast 3D Synthesis

MAISI-v2 incorporates rectified flow to accelerate latent diffusion model (LDM) inference by replacing the multi-step, stochastic denoising procedure of DDPMs with a deterministic ordinary differential equation (ODE) driven by a learned time-dependent velocity field $v_{t}(x)$. Latent variables are transported from isotropic Gaussian noise ($x_0 \sim \pi_0$) to encoded image latents ($x_1 \sim \pi_1$) with straightness regularization. The transport map $\varphi_t(x)$ is governed by:

$$\frac{d\varphi_t(x)}{dt} = v_t(\varphi_t(x)), \quad \varphi_0(x) = x_0, \quad \varphi_1(x) = x_1$$

In practice, $v_t$ is parameterized by a neural network $G_\theta$ and trained to minimize:

$$\mathcal{L}_{\mathrm{flow}} = \int_{0}^{1} \mathbb{E}_{x_0 \sim \pi_0,\, x_1 \sim \pi_1,\, t}\, \lVert v_t(x_t, c) - (x_1 - x_0) \rVert^2 \, dt$$

where $x_t = (1-t)x_0 + t x_1$ and $c$ represents conditioning variables such as voxel spacing (Eq. 5 in (Zhao et al., 7 Aug 2025)). The deterministic ODE is efficiently integrated (typically 30 steps), outperforming the canonical 1,000-step DDPM LDM with a 33-fold LDM acceleration (198 s $x_0 \sim \pi_0$0 6 s, Table 1). The process is: sample $x_0 \sim \pi_0$1, step through $x_0 \sim \pi_0$2 in $x_0 \sim \pi_0$3 steps updating $x_0 \sim \pi_0$4, then decode to image space via a VAE.

2. Region-specific Contrastive Loss for Condition Sensitivity

Condition fidelity, particularly for small but clinically salient structures, is improved by a region-specific contrastive loss. Given a segmentation mask $x_0 \sim \pi_0$5 (e.g., including a tumor) and a perturbed mask $x_0 \sim \pi_0$6 (e.g., same anatomy, but tumor region relabeled as organ), a binary ROI mask $x_0 \sim \pi_0$7 is derived and its complement $x_0 \sim \pi_0$8 (after one-voxel dilation) defines background voxels.

For a latent sample $x_0 \sim \pi_0$9, two outputs are generated: $x_1 \sim \pi_1$0 ROI sensitivity is quantified by the $x_1 \sim \pi_1$1 difference over ROI voxels,

$x_1 \sim \pi_1$2

capped by $x_1 \sim \pi_1$3 to prevent gradient explosion. The region-specific contrastive loss combines: $x_1 \sim \pi_1$4 and the training objective is

$x_1 \sim \pi_1$5

(Eq. 12). By maximizing differences in the ROI and minimizing them in the background, the model aligns synthetic variation tightly with the conditioning mask.

3. Architecture and Conditioning Mechanisms

The architecture consists of three principal modules:

• VAE (frozen): A 3D encoder compresses 1-channel CT volumes into 4-channel latent representations (downsampled by $x_1 \sim \pi_1$6, 16$x_1 \sim \pi_1$7 compression), with a sliding-window decoder for high-resolution 3D synthesis.
• Rectified-flow LDM ($x_1 \sim \pi_1$8): A 3D U-Net backbone (four downsampling/upsampling stages, 64–512 channels, 3×3×3 residual blocks with GroupNorm and SiLU). FiLM or cross-attention layers condition on time $x_1 \sim \pi_1$9 and voxel spacing. The noise prediction head of traditional LDM is replaced by a velocity-prediction head for rectified flow.
• ControlNet: A separate 3D U-Net encoder ($\varphi_t(x)$0) ingests segmentation masks, yielding multi-scale feature maps that modulate $\varphi_t(x)$1 through cross-attention and convolutional fusion.

All convolution, attention, and resampling operates natively in 3D. Bucketed data parallelism and mixed-precision training enable handling of volumes for which all dimensions are divisible by 128, with efficient sliding window decoding at $\varphi_t(x)$2 resolution patches.

4. Datasets, Training Procedure, and Model Generalizability

The training corpus consists of 11,000–12,000 CT scans covering head-and-neck, chest, and abdomen, with voxel spacings ranging $\varphi_t(x)$3 to $\varphi_t(x)$4, standardized so each dimension is divisible by 128. High-resolution scans are augmented by downsampling to lower resolution, yielding 107,000 training volumes. ControlNet masks comprise 127 pseudo-labels (VISTA3D) and additional manual annotation.

Training uses:

• Rectified-flow LDM: Pretrained on $\varphi_t(x)$5 crops, then mixed shapes with bucketed Distributed Data Parallelism (DDP), followed by fine-tuning on all shapes (16,000 + 2,000 epochs across 64$\varphi_t(x)$6A100 GPUs).
• ControlNet: Trained over 60 epochs using AdamW; region weighting and $\varphi_t(x)$7 scheduling are used to manage ROI emphasis and stability (weight=100 on tumor regions, $\varphi_t(x)$8).
• Regularization includes HU-intensity quality checks for loss scaling and $\varphi_t(x)$9 gradient clipping via $$\frac{d\varphi_t(x)}{dt} = v_t(\varphi_t(x)), \quad \varphi_0(x) = x_0, \quad \varphi_1(x) = x_1$$0.

A single LDM and ControlNet are trained jointly for all body regions and spacings, leveraging explicit spacing conditioning (see Fig. 8, Appendix).

5. Performance Metrics and Empirical Outcomes

Quantitative performance is evaluated on the out-of-distribution AutoPET2023 benchmark (512$$\frac{d\varphi_t(x)}{dt} = v_t(\varphi_t(x)), \quad \varphi_0(x) = x_0, \quad \varphi_1(x) = x_1$$1 at 1 mm$$\frac{d\varphi_t(x)}{dt} = v_t(\varphi_t(x)), \quad \varphi_0(x) = x_0, \quad \varphi_1(x) = x_1$$2 spacing). The mean 2D Fréchet Inception Distance (FID) on central slices across three views (lower is better) is as follows (Table 1):

Method	FID₍avg₎
HA-GAN	13.595
GenerateCT	10.757
MedSyn	24.709
MAISI (DDPM, 1000 steps)	2.441
MAISI (DDIM, 30 steps)	4.776
MAISI-v2 (Rectified Flow, 30 steps)	2.322

MAISI-v2 achieves both SOTA FID and pronounced speedup, with LDM inference reduced from 198s (MAISI, DDPM) to 6s and end-to-end (including VAE decode) to 21s. For 50 steps, end-to-end conditioned inference is 34s (Table 2).

Downstream segmentation with Auto3DSeg on five tumor types (5-fold cross-validation) quantifies augmentation utility. Dice scores (increment over real-only training) are:

Tumor Type	Dice Δ (MAISI-v2)	Final Dice
Liver	+3.3%	0.695
Lung	+7.5%	0.655
Pancreas	+6.4%	0.497
Colon	+4.2%	0.491
Bone-Lesion	+3.3%	0.537

Significant improvements are observed in 4/5 tumor categories (Table 4).

6. Practical Implications and Future Directions

MAISI-v2 demonstrates that rectified flow can be effectively adapted to 3D latent diffusion, drastically reducing inference steps without compromising generative fidelity, and that targeted contrastive objectives meaningfully enhance conditioning accuracy—particularly for fine-scale and clinically relevant ROIs. The framework's released codebase, weights, and GUI facilitate reproducibility and extension. A plausible implication is further adoption of rectified flow methods for high-throughput, controllable synthesis in clinical training and data augmentation, with potential generalization to broader classes of 3D conditional generation tasks (Zhao et al., 7 Aug 2025).