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

arXiv: https://arxiv.org/abs/2606.08670 · HTML: https://mlsanity.com/paper/2606.08670/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 21 of its 35 references that are in the map.
- First posted on arXiv 7 Jun 2026.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)

## Builds on (most cited first, 21 of 21)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Med3D: Transfer Learning for 3D Medical Image Analysis](https://mlsanity.com/paper/1904.00625/) (2019)
- [Brain Imaging Generation with Latent Diffusion Models](https://mlsanity.com/paper/2209.07162/) (2022)
- [Flow Matching on General Geometries](https://mlsanity.com/paper/2302.03660/) (2023)
- [Rotary Position Embedding for Vision Transformer](https://mlsanity.com/paper/2403.13298/) (2024)
- [On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks](https://mlsanity.com/paper/2203.09168/) (2022)
- [Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers](https://mlsanity.com/paper/2401.11605/) (2024)
- [FastSurfer -- A fast and accurate deep learning based neuroimaging pipeline](https://mlsanity.com/paper/1910.03866/) (2019)
- [Diffusion Probabilistic Models beat GANs on Medical Images](https://mlsanity.com/paper/2212.07501/) (2022)
- [Predicting brain age with deep learning from raw imaging data results in a reliable and heritable biomarker](https://mlsanity.com/paper/1612.02572/) (2016)
- [WDM: 3D Wavelet Diffusion Models for High-Resolution Medical Image Synthesis](https://mlsanity.com/paper/2402.19043/) (2024)
- [Medical Diffusion: Denoising Diffusion Probabilistic Models for 3D Medical Image Generation](https://mlsanity.com/paper/2211.03364/) (2022)
- [3D MedDiffusion: A 3D Medical Latent Diffusion Model for Controllable and High-quality Medical Image Generation](https://mlsanity.com/paper/2412.13059/) (2024)
- [Flow Matching for Medical Image Synthesis: Bridging the Gap Between Speed and Quality](https://mlsanity.com/paper/2503.00266/) (2025)
- [Heteroscedastic Uncertainty Estimation Framework for Unsupervised Registration](https://mlsanity.com/paper/2312.00836/) (2023)
- [FlowLet: Conditional 3D Brain MRI Synthesis using Wavelet Flow Matching](https://mlsanity.com/paper/2601.05212/) (2026)

## Built on by (most cited first, 1 of 1)

- [Now You Have My Healthy Attention: A U-DiT for Brain-MRI Inpainting](https://mlsanity.com/paper/2607.27974/) (2026)

## Abstract (arXiv)

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