WaveDiT: Distribution-Aware Wavelet Flow Matching for Efficient 3D Brain MRI Synthesis

Abstract: Large and demographically balanced datasets are essential for reliable neuroimaging biomarkers. Full-resolution 3D brain MRI synthesis can support data augmentation in this setting, but existing approaches either incur prohibitive computational cost at volumetric scale or rely on lossy latent compression that may compromise anatomical detail. As a result, practical 3D generative augmentation often requires specialized compute infrastructure. We propose WaveDiT, a conditional flow matching framework operating in the coefficient space of a 3D Haar Discrete Wavelet Transform. The model combines factorized spatio-depth attention with band-wise heteroscedastic uncertainty modeling derived from higher-order wavelet statistics. Predicted log-variance is integrated directly into both the flow objective and conditioning pathway, enabling adaptive precision consistent with the heavy-tailed and input-dependent variance structure of anatomical detail. This formulation supports full-resolution 3D synthesis under practical memory and time constraints on a single modern GPU. Evaluation on a multi-site cohort demonstrates improved alignment between generated and real MRI distributions, together with enhanced downstream brain age prediction and region-level anatomical agreement relative to diffusion, latent, and wavelet-based baselines. Code is available at https://github.com/sisinflab/WaveDiT

• The paper presents a novel conditional flow matching framework that uses wavelet decomposition to generate anatomically faithful 3D brain MRIs.
• It leverages the Morpheus module to model heteroscedastic uncertainty, ensuring precise weighting of high-frequency details and tissue boundaries.
• Experiments demonstrate state-of-the-art performance in FID, brain age prediction MAE, and ROI fidelity, synthesizing full-resolution MRIs in approximately 1 second on a single GPU.

WaveDiT: Distribution-Aware Wavelet Flow Matching for Efficient 3D Brain MRI Synthesis

Motivation and Context

Synthesizing high-fidelity, anatomically plausible 3D brain MRIs is critical for neuroimaging research, especially in the context of augmenting imbalanced cohorts and mitigating acquisition and privacy constraints in clinical datasets. Conventional diffusion models operating in voxel or latent spaces are limited by either prohibitive compute requirements or lossy anatomical reconstructions that degrade downstream utility. Recent wavelet-domain generative models alleviate some of these limitations, but typically neglect the heterogeneity of wavelet band statistics, significantly constraining anatomical preservation in high-frequency subbands. The paper introduces WaveDiT, a distribution-aware, conditional flow matching framework that unifies a wavelet coefficient representation with heteroscedastic uncertainty modeling and an Hourglass Diffusion Transformer (HDiT) backbone, enabling single-GPU, fast, and precise volumetric MRI synthesis.

Wavelet-Domain Modeling and Statistical Justification

WaveDiT leverages the 3D Haar DWT to decompose volumetric MRI data into eight subbands: a single low-frequency approximation (LLL) and seven high-frequency (HF) detail bands. The energetic and statistical analysis in the paper demonstrates that nearly all signal energy is concentrated in LLL (98.11\%), while HF subbands exhibit heavy-tailed, non-Gaussian distributions with strong, input-dependent heteroscedasticity. The kurtosis of HF bands rapidly escalates during the generative trajectory—from Gaussianity ($\kappa \approx 3$) at $t=0$ to extreme values between $27$ (single-axis) and $270$ ($HHH$) at $t=1$.

These findings directly motivate the necessity of modeling wavelet-band-specific predictive uncertainty and adaptive weighting during training. Uniform MSE losses over- or under-penalize errors depending on location and frequency, biasing spatial fidelity especially at tissue boundaries.

Morpheus: State-Aware Heteroscedastic Uncertainty Scheduler

To address this, the Morpheus module is introduced—a lightweight MLP that extracts wavelet-domain statistics (mean, std, L2 norm, skewness, kurtosis, max value) from the input at each time-step $t$, concatenating them with sinusoidal time embeddings. Morpheus regresses the per-band log-variance, forming the basis of a fully Bayesian, heteroscedastic training objective, which dynamically down-weights unpredictable HF content and up-weights anatomically consistent LLL regions. This uncertainty estimate enters both the loss and the conditioning pathway of the generative backbone, propagating frequency-aware reliability along the flow trajectory.

Figure 1: Training pipeline: wavelet decomposition, HDiT backbone with Morpheus scheduling, and Bayesian heteroscedastic loss.

Network Architecture and Flow Matching

The model extends the Hourglass Diffusion Transformer to 3D, exploiting factorized spatial–depth attention. Instead of costly cubic attention, WaveDiT reshapes wavelet tensors to a collection of 2D slices, where each slice is tokenized and processed by a transformer stack integrating time, position (via random Fourier features), age metadata, and the Morpheus frequency hint.

Two hierarchical attention types are used:

• Level 1 (Neighborhood Attention): Shallow transformer layers use local 2D attention, consistent with dominant edge-encoding properties in HF coefficients.
• Level 2 (Factorized Spatio-Depth Attention): Deeper layers combine in-plane (slice) global attention with cross-slice (depth) attention, preserving volumetric context with drastically reduced computational complexity.

The generative process relies on conditional flow matching (CFM)—in particular, rectified conditional paths—resulting in efficient synthesis via a small, fixed number of Heun integration steps.

Experimental Evaluation and Results

Benchmarking on a multi-site cohort (OpenBHB, ADNI, OASIS-3) with strict subject-level data separation, WaveDiT is evaluated along three axes: (1) global distributional metrics, (2) downstream augmentation utility for brain age prediction (BAP), and (3) region-of-interest (ROI) anatomical fidelity.

Strong quantitative performance is demonstrated. At 10 sampling steps, WaveDiT achieves:

• Best FID: 0.0039, surpassing FlowLet and WDM (even at 1000 steps).
• State-of-the-art BAP MAE: 2.44 years, lower than models trained with real or other synthetic data.
• Highest ROI fidelity: lowest region-wise intensity MAE and KL divergence, and highest Dice coefficient (0.46, compared to 0.44 for FlowLet and lower for all diffusion/latent baselines).

Ablation studies show removal of Morpheus substantially degrades both generative and downstream metrics, confirming that heteroscedastic, distribution-aware weighting is necessary to preserve anatomical fidelity under severe compositional diversity and limited compute. Furthermore, the CFM objective yields a more optimal trade-off between fidelity and sample efficiency compared to rectified or optimal transport flow variants.

Figure 2: Visual comparison of models. Axial, coronal, and sagittal views of a real 72 years old subject and age-conditioned generations at the same target age.

The speed–quality trade-off is significant: WaveDiT synthesizes full-resolution 3D MRIs in ~1 second (10 steps, single H100), compared to ~6 seconds for FlowLet and ~150 seconds for voxel-domain WDM at 1000 steps.

Implications and Future Directions

WaveDiT demonstrates that conditional flow matching in the wavelet domain, when paired with subband-wise uncertainty modeling and computationally efficient transformers, constitutes an effective paradigm for 3D neuroimaging data augmentation. The heteroscedastic approach is critical for robust learning over non-stationary anatomical content, and the tractability of single-GPU synthesis opens avenues for broader deployment in resource-constrained research and clinical settings.

The theoretical framework suggests applicability beyond MRI, notably to CT and other 3D modalities exhibiting similar signal–noise and local–global statistical patterns. Moreover, the Morpheus scheme could be extended to model richer forms of metadata, potentially conditioning on diagnosis, phenotype, or other task-relevant variables. Extensions incorporating multi-scale wavelet decompositions, adaptive integration of more expressive Fourier-based frequency representations, or integration with domain-specific anatomical priors are natural next steps.

Conclusion

WaveDiT advances the state-of-the-art in 3D MRI generative modeling by synthesizing anatomically faithful, demographically controllable volumes with highly efficient sampling and strong performance on clinical utility proxies for data augmentation. The results validate heteroscedastic, distribution-aware flow matching as essential for high-fidelity generative modeling in the medical imaging domain, highlighting the future potential for anatomy-preserving augmentations under hard computational and statistical constraints.