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

In plain words

Hospitals and researchers often need more 3D scans than they can share or collect, and a model that invents realistic CT scans can help. The earlier MAISI system could make whole-body scans of many kinds, but each one took a long time to generate, and the requested tumour did not always show up where it was asked for.

MAISI-v2 changes two things. First, it swaps the slow, many-step way the image model removes noise for a method that follows a nearly straight path from random noise to a scan, so far fewer steps are needed. Second, it adds a training penalty that checks the tumour specifically. The model is asked to draw the same scene twice, once with the tumour present and once with that tumour relabelled as ordinary tissue. The tumour area must change between the two results, and the rest of the body must stay the same.

The authors report that the faster model keeps image quality close to the older one and that the synthetic scans help a segmentation model find tumours a little better. These are the authors' own numbers on their own test setup.

TL;DR

MAISI-v2 is a 3D CT image generator built on MAISI. It replaces the DDPM sampler with rectified flow, so generation needs 30 steps instead of 1000, and adds a region-specific contrastive loss so that the tumour mask controls tumour appearance. Author-reported: the latent model is 33x faster in its sampling time, the FID is similar or better than MAISI's, and adding synthetic images improves Dice on all five tumour tasks, though only at p<0.1 for two of them. Code and weights are public; the weights carry a separate NVIDIA license.

Problem

Diffusion models make high-quality images, but in 3D they need hundreds of denoising steps, and each step runs a large network over a whole volume. Earlier 3D methods either cover only small volumes (256 cubed) or take over 10 minutes to generate a 512 cubed volume §1. Conditioning is also weak: a segmentation mask often does not produce the tumour it describes. The cycle-consistency fix from 2D work (ControlNet++) needs an extra reverse network trained on ground-truth conditions §1.

Method

Overview. MAISI-v2 keeps the MAISI pipeline and replaces the DDPM latent diffusion model with a rectified flow model. It reuses the pretrained MAISI autoencoder without fine-tuning §3.1. The conditioning for the latent model is the voxel spacing only. A ControlNet branch takes the segmentation mask and is trained with a weighted flow loss and the new contrastive term §3.1 §3.3 §3.4.

1. Autoencoder. The MAISI variational autoencoder compresses each volume by 4 in each spatial axis into 4 latent channels, a total compression of 16 §2.1. It is reused as is §3.1. At inference, the latent volume for a 512 cubed output is 4 x 128 x 128 x 128 §4.2.

2. Rectified flow. Let x0 be Gaussian noise and x1 the latent of a real image. The straight-line state at time t is x_t = (1 - t) x0 + t x1, and the training target is the constant velocity x1 - x0 §2.2 eq. 3. The network G predicts this velocity from x_t and the condition c eq. 4 §3.4. Sampling integrates the learned velocity field from noise to data. The paper uses 30 inference steps by default and reports results with 50 for the video comparison §4.2.

3. Training schedule of the latent model. The authors train in three stages with mixed precision §3.3:

• Pre-training on 128 cubed volumes, batch size 96, for one day.
• Main stage on full-resolution data grouped by shape (bucketed data parallelism), with a learning rate of 1e-3 and polynomial decay, for 16000 epochs.
• Fine-tuning on all data mixed with sampling weights to balance datasets, batch size 1, learning rate 1e-4, for about 2000 epochs. The stated optimiser is AdamW §3.3.

4. ControlNet with a weighted flow loss. The conditioning mask is encoded by a ControlNet, and the frozen latent model is driven by it §3.4. The flow loss is weighted by 100 inside tumour regions §3.4. The ControlNet is trained on CT images with pseudo segmentation maps of 127 anatomical classes and lesions from VISTA3D, replaced by manual labels where those exist §3.2.

5. Region-specific contrastive loss. This is the new term §3.4.1 alg. 1 fig. 1.

• Make two conditions from the same mask: the original mask c_orig, and a perturbed mask c_perturb in which tumour labels are replaced by the corresponding background label, such as a pancreatic tumour by pancreas §3.4.1.
• Feed the same noisy state x_t and the same time t to G twice, once with each mask §3.4.1 alg. 1.
• Form the ROI mask m where the two masks differ alg. 1, and the background mask m_minus as the complement of a dilated m eq. 7.
• ROI sensitivity: take the masked difference between the two outputs inside m, average it, and apply the bound D: the loss is minus the minimum of the average and delta eq. 5 eq. 6. The paper sets delta = 2 §3.4.1.
• Background consistency: a masked term on m_minus that should be small where the condition does not change eq. 8.
• Total objective: the flow loss plus lambda times (ROI loss plus background loss) eq. 9. The paper sets lambda by a quality check on CT intensities §3.4.4 §4.1.

6. Memory-aware variant. For large inputs, the paper computes the contrastive term on ControlNet encoder features instead of the final output, without the frozen latent model, to save memory. The loss terms are otherwise unchanged §3.4.2. See Paper vs code.

7. Quality gate. For each organ, the median HU of generated images must lie in a range set from training images: the more extreme of the 5th percentile and the 6-sigma bound on the low side, and the same on the high side with the 95th percentile §3.4.4. The lambda is chosen so that nearly 100% of generated images pass §3.4.4 §4.1.

# Region-specific contrastive step  [alg. 1] [§3.4.1]
# Inputs: latent x1 (B, C, X, Y, Z); mask c_orig (B, 1, X, Y, Z)
# Output: total training loss (scalar)

def contrastive_training_loss(x1, c_orig):
    # Noise and time on the straight path  [§2.2] [eq. 3]
    x0 = sample_gaussian(x1.shape)
    t = sample_time()
    x_t = (1 - t) * x0 + t * x1
    # Same condition with tumour labels replaced by organ labels  [§3.4.1]
    c_perturb = replace_tumour_labels(c_orig)
    # ROI mask: voxels where the two conditions differ  [alg. 1]
    roi = (c_perturb != c_orig)
    # Background: complement of a dilated ROI  [eq. 7]
    background = 1 - dilate(roi)
    # Two predictions from the same noisy input  [eq. 4] [alg. 1]
    v_orig = G(x_t, c_orig)
    v_perturb = G(x_t, c_perturb)
    # Flow loss against the constant velocity  [eq. 3]
    flow = squared_norm(v_orig - (x1 - x0))
    # ROI sensitivity: capped at delta = 2  [eq. 5] [eq. 6]
    diff = v_orig - v_perturb
    d_roi = mean_over_mask(abs(diff), roi)
    loss_roi = -minimum(d_roi, delta)
    # Background consistency  [eq. 8]
    loss_bg = mean_over_mask(abs(diff), background)
    # Weighted sum  [eq. 9]
    return flow + lambda_contrast * (loss_roi + loss_bg)

Hyperparameters

Results

Author-reported. Setup: AutoPET2023 as an out-of-distribution test set at 512 x 512 x 512 and 1 mm spacing, H100 GPUs, MONAI and PyTorch §4.1 §4.2. Baselines are HA-GAN, MedSyn, GenerateCT and MAISI (DDPM and DDIM) §4.2 tab. 1.

Image quality and speed tab. 1

Caption of tab. 1: FID is 2D in three planes, computed on the central 50% of slices §4.2. Time includes the LDM and VAE decoding (15 s) tab. 1.

Step ablation tab. 3

Inference time across modalities tab. 2, 50 steps, single 80 GB GPU per the caption:

The caption notes that the other methods' times come from their papers on other devices tab. 2.

Downstream segmentation (Dice, five-fold, Auto3DSeg) tab. 4. Real-only is the baseline. Augmented sets add the same number of synthetic images.

Caption of tab. 4: significance marks are p<0.1 and p<0.001 from a paired t-test §4.3. The text says the improvement is significant for four of five tumour types §4.3.

Qualitative results. Axial, sagittal and coronal slices compare methods fig. 2. The paper notes sagittal and coronal artefacts for the 2D GenerateCT and blur with mosaic artefacts for MedSyn and HA-GAN §4.2. Multi-region ControlNet samples are shown in fig. 3 and five tumour types in fig. 4.

Paper vs code

The official repository is linked in the paper's links block. Its owner, NVIDIA-Medtech, belongs to NVIDIA, the authors' employer. It is Apache-2.0 at commit f6362e4d0bc572e387f35b522369f0f15b27dade. Model weights are under a separate NVIDIA license (LICENSE.weights) and are not used on this card.

• Flow loss form. Paper eq. 3 is a squared norm eq. 3. The ControlNet and latent trainers use an L1 loss train_controlnet.py diff_model_train.py diff_model_train.py.
• Velocity target. The code's target is images minus noise train_controlnet.py, which matches x1 - x0 with x0 the noise eq. 3. This is consistent.
• ROI term. Paper eq. 6 is minus the minimum of D and delta eq. 6. The code writes it as a ReLU of (-D + delta) minus delta train_controlnet.py. These are equal. The code's D is a sum divided by the count of ROI voxels times channels, which the paper's norm notation does not state train_controlnet.py.
• Background term. Paper eq. 8 measures the difference between the two outputs on the background eq. 8. The code instead compares the ROI-free output with the ground-truth velocity over the background mask train_controlnet.py. The code has no difference term there.
• Dilation. The background mask uses MONAI dilation with filter size 3 train_controlnet.py. The paper states the mask as 1 minus dilate(m) with no size eq. 7.
• Memory-aware variant. Not in the release. The ControlNet block outputs are requested off and unused in the loss train_controlnet.py train_controlnet.py train_controlnet.py.
• Lambda schedule. The paper uses 0.01 for 30 epochs and then 0.001 §4.1. The released rflow-ct config sets a constant 0.01 config_maisi_controlnet_train_rflow-ct.json. No second value appears in the scripts.
• ControlNet hyperparameters. The paper gives LR 5e-5 for 60 epochs §3.4. The released rflow-ct config gives LR 1e-5 for 100 epochs and batch size 1 config_maisi_controlnet_train_rflow-ct.json. The scheduler is PolynomialLR with power 2.0 train_controlnet.py.
• Tumour weight label. The paper weights "tumours" by 100 §3.4. The released config weights only label 23 config_maisi_controlnet_train_rflow-ct.json. Label 23 is the lung tumour in the tumour remapping augmentation.py.
• Removing tumours. The paper's example replaces a pancreatic tumour with pancreas §3.4.1. The code maps 26 to liver 1, 24 to pancreas 4, 27 to colon 62, 116 to 14, 117 to 5 and brain tumours to 22 augmentation.py. Lung tumour 23 is replaced by a majority vote of neighbouring organ labels augmentation.py.
• Freezing. The latent model is frozen during ControlNet training train_controlnet.py. The ControlNet starts from the diffusion weights or a checkpoint train_controlnet.py.
• Latent model optimiser and schedule. The paper says AdamW with three stages §3.3. The trainer uses Adam diff_model_train.py. The released diffusion config has LR 1e-5, batch 2, 1000 epochs config_maisi_diff_model_rflow-ct.json. This is not the paper's three-stage schedule.
• Timestep sampling. The released network config uses an RFlow scheduler with timestep transform and scale 1.4, which the paper does not mention config_network_rflow.json.
• Inference steps. The rflow diffusion config sets 30 steps config_maisi_diff_model_rflow-ct.json. The 80 GB release inference config sets 1000 steps config_infer_80g_512x512x512.json.
• Quality check thresholds. The paper uses the 5th and 95th percentiles §3.4.4. The code uses the 0.5th and 99.5th percentiles together with the 6-sigma bounds quality_check.py image_median_statistics_ct.json.
• FID setup. The paper uses the central 50% of slices §4.2. The example command in the FID script uses a centre ratio of 0.4 compute_fid_2-5d_ct.py. The paper says "Inception" §4.2. The script defaults to RadImageNet ResNet50 features compute_fid_2-5d_ct.py. The feature network behind the reported FID is not stated.

Pitfalls

• The 33x figure is for the latent model's sampling time only. It compares 6 s for 30 steps with 198 s for 1000 steps tab. 1. Adding the 15 s VAE decode gives 21 s against 213 s, about 10x overall. This is derived from tab. 1.
• Flow loss form and the paper's loss differ. Implement the squared norm of eq. 3 if you follow the paper, or L1 if you follow the code. The two behave differently at small residuals.
• The contrastive ROI term has no gradient once it exceeds delta. The clamp at delta = 2 eq. 6 stops the gradient for voxels once the normalised difference exceeds it. Without the clamp, training becomes unstable and produces NaNs §3.4.1. The paper says delta = 2 corresponds to a plus-or-minus one standard deviation range, which assumes the latents are near Gaussian §3.4.1.
• Lambda is hand-tuned. The authors chose it by a visual and HU criterion §3.4.4 §4.1. Large lambda aligns tumours but breaks other organs. The 0.01 to 0.001 schedule is empirical §4.1.
• Significance in the segmentation table is weak for two tumours. The paper's text claims significance for four of five tumours §4.3. Two of those are at p<0.1, which is weaker than the usual 0.05 threshold tab. 4.
• Rectified flow tumour alignment can be weaker than DDPM. The authors report weaker tumour alignment for rectified flow than for DDPM, especially for colon tumours §5. They propose two causes and leave the mechanism open §5.
• Memory limits. Inference needs about 40 GB for 512 x 512 x 768 and training needs 80 GB §5. The LDM trained on 64 A100 80 GB GPUs for three weeks §3.3.
• Data is not fully public. The training set includes 237 in-house bone lesion CTs and 1386 in-house chest and abdomen CTs §4.1. The released code cannot reproduce that set.
• Data counts differ between sections. Section 3.2 gives 1.4k in-house chest and abdomen scans and 11k usable scans §3.2. Section 4.1 adds 237 bone lesion scans and 1386 chest and abdomen scans §4.1.
• The sampler and timestep distribution are not stated. The paper names rectified flow and 30 steps but not the solver. The code uses MONAI's RFlowScheduler with a timestep transform config_network_rflow.json. Re-implementers must choose a solver and time distribution.
• Weights carry a separate license. The released weights are under the NVIDIA License in LICENSE.weights, which this card has not assessed. The card does not depend on the weights.
• Synthetic CT is not for diagnosis. The authors say the synthetic images are for training other models, not clinical use (Acknowledgments).

Lineage

• Builds on: MAISI, a 3D latent diffusion framework for CT synthesis (Guo et al., WACV 2025) §2.1. The paper gives no arXiv id, and there is no card for it. Rectified flow is cited as Liu et al. 2022 §2.2. No arXiv ids are given in the source.
• Compared against: HA-GAN, GenerateCT, MedSyn, MAISI (DDPM and DDIM), DiffTumor and 3D MedDiffusion §4.2 §4.3 tab. 1. None of these has a card on this site.
• Superseded by: not stated.