WFM: 3D Wavelet Flow Matching for Ultrafast Multi-Modal MRI Synthesis

Abstract: Diffusion models have achieved remarkable quality in multi-modal MRI synthesis, but their computational cost (hundreds of sampling steps and separate models per modality) limits clinical deployment. We observe that this inefficiency stems from an unnecessary starting point: diffusion begins from pure noise, discarding the structural information already present in available MRI sequences. We propose WFM (Wavelet Flow Matching), which instead learns a direct flow from an informed prior, the mean of conditioning modalities in wavelet space, to the target distribution. Because the source and target share underlying anatomy and differ primarily in contrast, this formulation enables accurate synthesis in just 1-2 integration steps. A single 82M-parameter model with class conditioning synthesizes all four BraTS modalities (T1, T1c, T2, FLAIR), replacing four separate diffusion models totaling 326M parameters. On BraTS 2024, WFM achieves 26.8 dB PSNR and 0.94 SSIM, within 1-2 dB of diffusion baselines, while running 250-1000x faster (0.16-0.64s vs. 160s per volume). This speed-quality trade-off makes real-time MRI synthesis practical for clinical workflows. Code is available at https://github.com/yalcintur/WFM.

• The paper introduces a novel 3D wavelet flow matching method that replaces iterative denoising with a learned velocity field for efficient multi-modal MRI synthesis.
• It consolidates synthesis of four BraTS modalities into a single 82M-parameter model, achieving competitive PSNR and SSIM scores with significantly reduced inference time.
• Experimental results demonstrate that WFM preserves key anatomical structures and provides a 250–1000× speedup over traditional diffusion models.

WFM: 3D Wavelet Flow Matching for Ultrafast Multi-Modal MRI Synthesis

Motivation and Background

Multi-modal MRI is a foundation of brain tumor analysis, with each protocol (T1, T1c, T2, FLAIR) providing complementary information. However, missing modalities are common in clinical practice due to acquisition time constraints or artifacts. Recent advances in conditional diffusion models have established a new quality ceiling for multi-modal MRI synthesis but suffer from two fundamental limitations: they require hundreds to thousands of iterative denoising steps per sample, and they often use separate, large models for each target modality. These inefficiencies preclude their adoption in routine clinical workflows, where sub-second synthesis latency is critical.

Existing diffusion models, e.g., cWDM, start from isotropic Gaussian noise and slowly reconstruct anatomical structure via stepwise denoising. This stochastic approach disregards structural information already present in available modalities. WFM ("Wavelet Flow Matching") proposes an orthogonal solution: replacing the forward process with a learned continuous flow from a highly informative, structurally aligned source, namely the mean of the conditioning modalities in 3D wavelet space.

WFM Formulation and Methodology

Informed Prior Construction

WFM constructs the initial source distribution by averaging the discrete wavelet transforms (DWT) of the three available modalities, yielding a latent representation rich in anatomical details and shared features between source and target. This approach rationalizes that contrast transformation, not structural synthesis, forms the core challenge in this setting. The explicit construction is

$$\mathbf{x}_{\text{source}} = \frac{1}{3}\sum_{i=1}^3 \text{DWT}(\mathbf{X}_{s_i})$$

where $\mathbf{X}_{s_i}$ denote the conditioning volumes.

Flow Matching Learning Objective

Rather than learning an iterative denoising process, WFM optimizes a neural network $f_\theta$ to predict the velocity field required to linearly interpolate in wavelet space between $\mathbf{x}_{\text{source}}$ and the target $\mathbf{x}_{\text{target}}$. The path is parameterized as

$$\mathbf{x}_t = (1 - t)\mathbf{x}_{\text{source}} + t\mathbf{x}_{\text{target}}$$

with $t \in [0,1]$ and a constant target velocity $$\mathbf{v} = \mathbf{x}_{\text{target}} - \mathbf{x}_{\text{source}}$$. Stochastic regularization noise is injected during training, maximizing generalization and robustness while maintaining clean samples at both endpoints.

Unified Multi-Modal Synthesis Architecture

WFM consolidates synthesis of all four BraTS modalities into a single 82M-parameter 3D U-Net with class conditioning, in contrast to the traditional approach of training four independent models (326M parameters total). Target modality is encoded as a class embedding and jointly learned with the continuous time embedding.

Efficient Few-step ODE Integration

At inference, the learned velocity field is integrated either with a single-step Euler or two-step Heun ODE solver, typically requiring just 1 or 2 function evaluations (NFE), to transform the source into the predicted target. The final image is reconstructed via the inverse DWT. This results in an extreme reduction in compute time compared to standard diffusion pipelines.

Figure 1: (A) cWDM requires 1000 denoising steps (165.4s); WFM achieves comparable synthesis quality in 1-2 steps (0.34-0.64s) starting from an informed prior. (B) Pipeline overview: conditioning modalities are wavelet-averaged to form the source, which is transported to the target via a learned velocity field.

Experimental Analysis

Dataset and Metrics

Experiments were conducted on the BraTS 2024 benchmark (1,032 train / 219 validation cases). Performance was measured using PSNR and SSIM within the brain mask. Notably, validation set standard deviations are high (±2–3 dB), reflecting significant anatomical heterogeneity in glioma cohorts.

Quantitative Results

WFM delivers an average of 26.8 dB PSNR and 0.94 SSIM on the validation set, using 1-2 integration steps (NFE=1-2). The best diffusion baseline, cWDM, achieves 28.4 dB PSNR and 0.95 SSIM but requires 1000 steps and a 4× larger parameter footprint. Thus, WFM trades approximately 1.6 dB PSNR for a 250–1000× speedup (0.16–0.64 seconds per volume vs. 160s for cWDM).

Unified vs. Separate Models

The unified WFM model shows negligible performance loss (≤0.1 dB) compared to hypothetical separate models, confirming class conditioning's effectiveness and the benefits of parameter sharing.

Qualitative Evaluation

Detailed visualizations indicate that WFM reliably preserves anatomical boundaries and pathological features (e.g., tumors, ventricle boundaries, edema), particularly in regions where source modalities provide relevant cues. However, some loss of detail is evident in cases where target contrast (e.g., T1c enhancement) is unpredictable from available inputs.

Figure 2: Axial slices: left—successful case with precise structural details; right—tumor boundaries and enhancement patterns well preserved. Error maps highlight areas of synthesis difficulty (e.g., tumor cores).

Figure 3: Sagittal views: structural fidelity is consistently maintained across orientations, with errors localized primarily to challenging regions.

Figure 4: Additional examples demonstrating strong generalization and structural preservation across diverse validation cases.

Efficiency and Theoretical Implications

WFM's primary advance is reframing generation as a flow-matching problem with a domain-informed source. This design enables parameter sharing, drastically reduces computational demands, and leverages structural redundancy across modalities. The nearly constant velocity field—justified by structural alignment of source and target—renders higher-order integration steps unnecessary, achieving optimal quality in as few as 1–2 steps.

Compared to bridge-based diffusion (e.g., Schrödinger bridges, Brownian bridges) or latent diffusion, WFM leverages the medical imaging context (shared anatomy, contrast-dominated translation) to maximize tractability and speed. The wavelet space operation, as opposed to deep learned latent codes, ensures invertibility and avoids lossy compression that could erase subtle pathologies.

Limitations and Future Research

WFM is evaluated primarily on BraTS glioma data with fixed triplet-to-single modality synthesis. Its effectiveness for other pathologies, variable conditioning set sizes, and less structurally aligned source-target pairs (e.g., CT-to-MRI) remains unproven. Additionally, downstream task efficacy (e.g., segmentation using synthetic modalities), uncertainty quantification, and rigorous clinical validation must be pursued.

Given the growing importance of rapid, interactive imaging tools in clinical workflows, extensions to more generalized, robust informed-prior frameworks and automated error flagging will be critical. Methods such as HACA3 for harmonization or I$^2$SB for unpaired translation could serve as comparators for future benchmarks.

Conclusion

WFM establishes a new paradigm for MRI synthesis in the clinical regime, demonstrating that computational and memory burdens of diffusion models are not inescapable. By leveraging structured, domain-informed priors and unified architectures, WFM achieves near-diffusion fidelity in a fraction of the time and compute, with a single model applicable to all standard modalities. This work highlights informed source construction and flow matching as practical alternatives to pure-noise initialization, with immediate implications for real-time medical image analysis and generative modeling in data-rich, structurally aligned domains (2604.21146).