Unified Multi-plane Autoregressive Diffusion for 3D Multi-contrast MRI Synthesis

AI summary

Overview

Research area: Medical image analysis / computer vision — specifically 3D multi-contrast MRI synthesis using latent diffusion and autoregressive generative modeling.

Technical level: Advanced. The paper assumes familiarity with latent diffusion models, autoencoders, spatially adaptive normalization, and autoregressive generation.

Scope: This paper introduces MPAD (Multi-Plane Autoregressive Diffusion), a latent diffusion framework that synthesizes full 3D MRI volumes of a missing contrast using only 2D slice-wise operations, and evaluates it against five baselines on the ADNI and IXI brain MRI datasets.

What This Paper Is About

MRI exams often acquire several contrasts (such as T1w, T2w, and PDw), each highlighting different tissue properties, but collecting all of them is slow and uncomfortable for patients. The paper's goal is to synthesize the missing contrasts from the ones already acquired, which is inherently a 3D problem because anatomy must remain coherent across axial, sagittal, and coronal planes. Fully 3D generative models handle this coherence but are impractical because their computational cost scales cubically with volume size, so the authors instead aim to get 3D-quality results from efficient 2D operations.

Key Contributions

1. Reformulating 3D MRI synthesis as 2D masked-slice prediction. The authors train a 2D diffusion model to reconstruct masked latent slices of a target contrast, conditioned on source-contrast slices, unmasked target-contrast slices, and a text prompt, thereby avoiding the cubic scaling of volumetric diffusion.

2. A multi-plane autoregressive inference scheme with inter-plane priors. Slices are generated autoregressively in random order within one plane orientation (Intra-plane Autoregressive Synthesis, IAS), and each completed plane serves as a conditioning prior for orthogonal orientations (Inter-plane Prior Inference, IPI), followed by a refinement pass and voxel-wise averaging of four candidate volumes.

3. A unified one-to-many model. A single model handles multiple target contrasts, rather than requiring a separately trained model for each contrast pair. Among the compared baselines, only ALDM shares this one-to-many capability.

4. Demonstrated efficiency and quality gains. MPAD reduces training FLOPs by 7× and inference FLOPs by 3× relative to 3D latent diffusion baselines, while also lowering inference time and peak memory consumption, and it improves synthesis fidelity across multiple datasets.

Main Findings

• Computational savings: MPAD reduces training FLOPs per iteration by 7× and inference FLOPs per volume by 3× compared to 3D latent diffusion baselines, with lower inference time and peak memory.

• ADNI performance: MPAD achieves the highest SSIM in five out of six synthesis tasks and consistently the lowest NMSE. For the challenging PD→T1 and PD→T2 tasks, it improves PSNR by over 2 dB relative to the second-best method. Its reported PSNRs are 23.10 (T1→T2), 23.78 (T1→PD), 23.99 (T2→T1), 25.24 (T2→PD), 22.56 (PD→T1), and 22.59 (PD→T2).

• IXI performance: MPAD ranks first across all metrics on the IXI dataset, with particularly strong results on T2→T1 and PD→T2. Its reported PSNRs are 30.09, 28.11, 28.98, 29.21, 28.88, and 30.94 across the six tasks.

• Qualitative results: On the ADNI T2→T1 task, CycleGAN-3D and EaGAN produce blurring and loss of fine detail; LDM-3D and cWDM preserve structure but introduce intensity distortions; ALDM shows misalignments in complex anatomical regions. MPAD produces sharper boundaries and more accurate tissue contrast across axial, coronal, and sagittal views. Similar trends appear on IXI for T1→T2.

• Autoregression helps modestly: Adding IAS improves PSNR and NMSE on both datasets relative to a non-AR baseline (ADNI PSNR 22.79 to 22.85 at g=4; IXI 28.89 to 29.08). Smaller group sizes give small incremental gains, with g=1 reaching 22.87 (ADNI) and 29.14 (IXI), and g=4 chosen for all experiments.

• 3D conditioning encoder beats 2D: Replacing the 2D MCE with a 3D version adds +0.72 PSNR and +0.019 SSIM on ADNI and +1.19 PSNR and +0.007 SSIM on IXI.

• Gaussian noise masking beats learnable tokens: Gaussian noise masking adds +0.32 PSNR / +0.029 SSIM on ADNI and +2.00 PSNR / +0.017 SSIM on IXI over learnable mask tokens.

• Inter-plane priors matter: Enabling priors adds +1.12 PSNR and +0.037 SSIM on ADNI, and +1.88 PSNR and +0.013 SSIM on IXI.

• More planes help: Performance improves as planes are added, from a single plane (23.54 is reached only with all four passes on ADNI: 23.54 PSNR / 0.832 SSIM; IXI 29.37 / 0.886), with the refinement pass on the first plane providing a small but consistent extra gain.

• Plane ordering is robust: Different generation orders produce nearly identical results (23.48 to 23.54 PSNR, 0.830 to 0.832 SSIM on ADNI), indicating the process is largely invariant to plane sequence.

• Remaining inference headroom: Reducing inter-plane refinement iterations lowers inference time to approximately 1.8 s per volume, indicating a controllable trade-off between speed and multi-plane consistency.

Methodology in Plain English

The approach has two training stages and a multi-pass inference scheme.

Stage 1 — compressed 3D representation. A KL-regularized 3D autoencoder compresses each MR volume into a compact latent tensor. The encoder produces a latent space with 3 channels at 32×32×32 resolution, using a 3D CNN with channel multipliers of 1, 2, and 4 and two residual blocks per resolution level, trained for 1000 epochs. The key property is that the latent is isotropic, so it can be sliced along axial, coronal, or sagittal planes without geometric bias.

Stage 2 — slice-wise diffusion training. A 2D diffusion model is trained to reconstruct masked slices of the target contrast. At each iteration, a plane orientation is sampled randomly; target latent slices are stacked and a random subset is masked with a ratio uniformly sampled between 0.7 and 1.0. A Multi-modal Conditioning Encoder fuses source-contrast latents and the partially masked target latents, together with a text embedding from BioMedCLIP naming the desired contrast (for example, "T2-weighted MR image"). These conditioning features are injected into the diffusion model through SPADE modulation layers, which apply spatially varying scale and shift at each location. The model is trained with the standard noise-prediction objective, using a 2D CNN UNet backbone with 2 residual blocks per level and 32 channels per attention head, a linear noise schedule from β1 = 0.0015 to βT = 0.0195 over T = 1000 diffusion steps, and AdamW at a learning rate of 4×10⁻⁶ for 1000 epochs. Slices from different planes are batched together so the network becomes plane-agnostic.

Inference — autoregressive multi-plane generation. Within a plane, slices are generated one group at a time in random order, each conditioned on previously generated slices. When moving to an orthogonal plane, the already-completed volume is used as a prior: rather than starting from pure noise, the model extracts the corresponding slice from the completed volume, applies forward diffusion to an intermediate timestep τ, and reverses from there. The first plane is generated from pure noise using 10 DDIM steps, while prior-conditioned planes use only 2 timesteps. After all three planes, a refinement pass on the first plane incorporates information from all orientations. The four resulting candidates are combined by voxel-wise averaging, and the final volume is obtained by decoding the aggregated latent.

Why This Matters

Impact on research: The paper shows that global 3D anatomical consistency does not strictly require a 3D denoising network. By factorizing a volume into sequences of 2D slices and propagating context across orthogonal planes, it offers an alternative to the cubic cost of volumetric diffusion and opens a path toward unified models that handle many target contrasts rather than one model per contrast pair.

Real-world applications:

• Shortening MRI exam protocols by acquiring only a subset of contrasts and synthesizing the rest, reducing scan time and patient discomfort.
• Reducing susceptibility to motion artifacts that arise during prolonged scans, since shorter acquisitions mean less opportunity for patient movement.
• Standardizing imaging workflows across sites and scanners by filling in contrasts that were not acquired under a given protocol.
• Enabling practical deployment on standard GPUs, where the paper reports lower peak memory and inference time than 3D baselines.

Industry relevance: Hospitals and imaging centers face throughput pressure on scanners, and scanner manufacturers and clinical software vendors have a direct interest in reducing acquisition burden without losing diagnostic information. The reported 7× training and 3× inference FLOP reductions, plus support for one-to-many translation in a single model, are relevant to cost-conscious deployment. The work was supported by the MSIT Global Research Support Program in the Digital Field (RS-2024-00436680) through IITP and by Microsoft Research Asia.

Future Directions

• Further inference acceleration: The paper notes that reducing inter-plane refinement iterations brings inference time to approximately 1.8 s per volume, suggesting additional efficiency headroom without architectural changes.
• Efficient aggregation or parallel generation: The authors propose exploring more efficient aggregation mechanisms or parallel generation strategies to speed up inference while preserving multi-view consistency.
• Distillation and fewer-step sampling: Distillation techniques or fewer-step sampling methods are suggested as ways to reduce the sampling burden without compromising quality.
• Broader validation: Details of the datasets are deferred to supplementary material and experiments are limited to ADNI and IXI brain MRI, leaving open how well the framework generalizes to other anatomies, contrasts, and acquisition settings.

Target Audience

This paper is most useful to machine learning researchers working on medical image synthesis and generative models, particularly those interested in diffusion architectures, autoregressive decoding, and latent-space methods. It also serves medical imaging scientists and clinical researchers evaluating whether synthesized contrasts are trustworthy for diagnostic workflows, and engineers building deployable imaging pipelines who care about FLOPs, memory, and inference latency. Readers need a solid background in diffusion models and 3D convolutional networks to follow the technical details.

Authors’ abstract

Acquiring a complete set of magnetic resonance imaging (MRI) contrasts is time-intensive and uncomfortable for patients, despite the diagnostic value of multi-contrast imaging. This motivates synthesizing missing contrasts from those already acquired, which is an inherently 3D problem requiring anatomical coherence across axial, sagittal, and coronal planes. However, fully 3D generative models are often impracti- cal under computational resources that scale cubically with volume size. We propose a unified Multi-Plane Autoregressive Diffusion (MPAD), a latent diffusion framework that achieves full-volume 3D synthesis using efficient plane-wise 2D operations while preserving volumetric coherence. A 3D autoencoder first compresses MRI scans into an isotropic 3D la- tent representation. A 2D diffusion model is then trained to reconstruct masked latent slices of the target contrast, conditioned on both source- contrast slices and unmasked target-contrast slices. During inference, we introduce plane-wise autoregressive synthesis with inter-plane priors. Slices are generated autoregressively in random order within one plane orientation to maintain intra-plane continuity, then propagated as con- ditioning priors to orthogonal plane orientations to enforce inter-plane consistency. Compared to 3D latent diffusion baselines, MPAD reduces training and inference FLOPs by 7x and 3x, respectively, while also lowering inference time and peak memory consumption. Experiments on multiple datasets demonstrate that MPAD achieves superior perfor- mance, generating high-fidelity 3D volumes and supporting one-to-many translation within a single unified model.