Source-linked AI summary
WaveDiT: Distribution-Aware Wavelet Flow Matching for Efficient 3D Brain MRI Synthesis
Danilo Danese, Angela Lombardi, Giuseppe Fasano, Matteo Attimonelli, Tommaso Di Noia
TL;DR
Full-resolution 3D brain MRI synthesis remains computationally demanding, while lossy compression may compromise anatomical detail. WaveDiT uses wavelet-domain conditional flow matching with state-aware uncertainty modeling and improves distributional, brain-age, and regional anatomical metrics under low-step sampling.
Problem
Full-resolution 3D brain MRI synthesis remains computationally demanding, while lossy latent compression may compromise fine-grained anatomical detail.
Method
WaveDiT performs conditional flow matching in 3D Haar wavelet-coefficient space with factorized spatio-depth attention and band-wise state-aware uncertainty modeling.
Results
WaveDiT-CFM improves global distributional metrics, brain age prediction, and ROI-level anatomical metrics over diffusion, flow-based, and wavelet-based baselines under low-step sampling.
Takeaways & Limitations
WaveDiT provides a quantitative route to full-resolution 3D brain MRI synthesis with improved regional anatomical plausibility under practical sampling constraints.
Takeaways & Limitations
The study is limited to T1-weighted MRI and lacks expert reader studies, so conclusions rely on quantitative proxies of anatomical fidelity and clinical utility.
Abstract
from arXiv · showhide
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
1 Introduction
High-resolution brain MRI supports clinically relevant neuroimaging tasks but requires large, demographically balanced cohorts that are difficult to acquire. Full-resolution 3D generative modeling remains computationally demanding, limiting accessibility outside specialized infrastructure.
- Motivation: High-resolution brain MRI supports brain age prediction, disease-risk stratification, and longitudinal monitoring of neurodegeneration, motivating robust biomarkers from large, demographically balanced cohorts.Such cohorts are difficult to acquire because of high acquisition costs, privacy constraints, and site-specific heterogeneity.
- Computational challenge: Full-resolution 3D MRI generation is computationally demanding because pixel-space diffusion requires hundreds to thousands of denoising steps over millions of voxels.Training at full volumetric resolution often requires high-memory GPUs and prolonged compute schedules.
2 Methods
WaveDiT synthesizes full-resolution 3D brain MRI in single-level 3D Haar wavelet coefficient space, modeling band- and state-dependent uncertainty with Morpheus. Its factorized spatio-depth attention reduces volumetric attention cost while preserving coherence across slices.
- Wavelet representation: A single-level 3D Haar DWT maps each volume into one LLL approximation band and seven directional high-frequency detail bands at half resolution per spatial dimension.The transform maps R1×D×H×W to R8×D′×H′×W′, where D′ = D/2, H′ = H/2, and W′ = W/2.
- Wavelet statistics: 98.11% of total energy lies in LLL, while HF-band kurtosis rises from approximately 3 at t = 0 to 27–31, 84–93, and approximately 270 across increasingly multidirectional bands at t = 1.The HHH band reaches kurtosis approximately 270, corresponding to an 89.8× increase relative to t = 0.
- Uncertainty modeling: Morpheus predicts band-wise log-variance from wavelet statistics and time, using it in a Bayesian heteroscedastic flow objective and as frequency conditioning for the backbone.The objective uses adaptive precision, while the conditioning pathway is active during both training and sampling.
- Efficient attention: For a 224^3 input, wavelet decomposition produces approximately 1.4 × 10^6 spatial tokens, and factorized depth attention reduces complexity from approximately 2 × 10^12 to approximately 1.8 × 10^10.The architecture treats the volume as a batch of 2D slices and restores volumetric coherence through factorized depth attention.
- Sampling and reconstruction: Sampling integrates Gaussian noise with a second-order Heun solver and reconstructs the synthesized volume in voxel space using the inverse 3D Haar DWT.Morpheus supplies frequency conditioning at each integration step.
3 Experiments and Results
Experiments use a multi-site, subject-disjoint cohort with standardized preprocessing to evaluate generative quality and downstream utility. WaveDiT-CFM shows stronger global distributional alignment, brain-age performance, and regional anatomical agreement than the reported baselines, while global metrics alone may miss regional discrepancies.
- Generative quality: At 10 steps, WaveDiT-CFM achieves the lowest reported generative-quality metrics among all methods and improves over the conditioned wavelet-based FlowLet baseline.Compared with WDM, WaveDiT-CFM also requires fewer steps and achieves lower FID.
- Downstream evaluation: WaveDiT-CFM achieves the lowest brain-age prediction MAE, outperforming the reference model trained solely on real data and other conditional baselines for synthetic augmentation.BAP evaluates chronological-age prediction from MRI using separate 3D DenseNet models trained with each method’s synthetic samples.
- Downstream evaluation: WaveDiT-CFM obtains the best ROI-level scores, with lower iMAE and KL divergence and higher Dice than the alternatives.ROI evaluation averages results across 95 regions.
- Downstream evaluation: WDM and MD can achieve competitive FID while showing reduced Dice and higher KLD, indicating that global metrics may underestimate regional brain-structure discrepancies.The results motivate combining generative-quality metrics with brain-age prediction and ROI-level analysis.
4 Conclusion
WaveDiT is a wavelet-domain conditional flow matching model for 3D brain MRI synthesis that combines an HDiT backbone with the Morpheus state-aware uncertainty scheduler. Under low-step sampling, it improves distributional, brain age prediction, and ROI-level anatomical results versus diffusion and flow-based baselines, with ablations supporting both wavelet representation and uncertainty weighting.
- Contributions: WaveDiT combines an HDiT backbone with the Morpheus state-aware uncertainty scheduler for wavelet-domain conditional flow matching in 3D brain MRI synthesis.The model operates in the wavelet domain and uses state-aware uncertainty scheduling.
- Results: Under low-step sampling, WaveDiT-CFM achieves strong global distributional scores while improving brain age prediction and ROI-level anatomical metrics over diffusion and flow-based baselines.The reported improvements span global distributional evaluation, downstream brain age prediction, and region-of-interest anatomical assessment.
- Ablations: Ablations indicate that wavelet representation and state-aware uncertainty weighting stabilize training and preserve region-level anatomical plausibility.Both components contribute to the reported training stability and anatomical plausibility.