WaveDiT: Distribution-Aware Wavelet Flow Matching for Efficient 3D Brain MRI Synthesis
Summary
WaveDiT is a conditional flow matching framework for full-resolution 3D brain MRI synthesis that operates in wavelet coefficient space, enabling efficient generation on standard GPUs without lossy latent compression. It achieves improved alignment with real MRI distributions and downstream tasks.
View Cached Full Text
Cached at: 06/15/26, 09:02 AM
Paper page - WaveDiT: Distribution-Aware Wavelet Flow Matching for Efficient 3D Brain MRI Synthesis
Source: https://huggingface.co/papers/2606.08670
Abstract
WaveDiT generates full-resolution 3D brain MRIs using conditional flow matching in wavelet coefficient space with adaptive precision modeling, enabling efficient synthesis on standard hardware.
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, aconditional flow matchingframework operating in the coefficient space of a3D Haar Discrete Wavelet Transform. The model combinesfactorized spatio-depth attentionwithband-wise heteroscedastic uncertainty modelingderived from higher-order wavelet statistics. Predicted log-variance is integrated directly into both theflow objectiveand conditioning pathway, enabling adaptive precision consistent with the heavy-tailed and input-dependent variance structure of anatomical detail. This formulation supportsfull-resolution 3D synthesisunder 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 downstreambrain age predictionand region-levelanatomical agreementrelative to diffusion, latent, and wavelet-based baselines. Code is available at https://github.com/sisinflab/WaveDiT
View arXiv pageView PDFProject pageGitHub1Add to collection
Community
Paper author
Paper submitter
One full-size 3D brain. Every voxel, no lossy latent.
Size isn’t the only hard part: wavelet bands aren’t statistically equal, high-frequency anatomy is sparse, heavy-tailed, and grows sharper as noise becomes a brain. So we built Morpheus: a scheduler that focuses the model where the signal is predictable, not a uniform loss.
It builds directly on our previous work, FlowLet, and scales to native 182×218×182 volumes (and beyond), with training and inference on a single consumer GPU.
Upload images, audio, and videos by dragging in the text input, pasting, orclicking here.
Tap or paste here to upload images
Get this paper in your agent:
hf papers read 2606\.08670
Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash
Models citing this paper1
#### danesed/WaveDiT Unconditional Image Generation• Updated37 minutes ago • 1
Datasets citing this paper0
No dataset linking this paper
Cite arxiv.org/abs/2606.08670 in a dataset README.md to link it from this page.
Spaces citing this paper1
Collections including this paper0
No Collection including this paper
Add this paper to acollectionto link it from this page.
Similar Articles
FlowLet: Conditional 3D Brain MRI Synthesis using Wavelet Flow Matching
FlowLet is a conditional generative framework that synthesizes age-conditioned 3D brain MRIs using flow matching in an invertible wavelet domain, improving brain age prediction accuracy for underrepresented age groups with high efficiency.
Functional MRI Time Series Generation via Wavelet-Based Image Transform and Spectral Flow Matching for Brain Disorder Identification
This paper proposes DSFM, a novel generative framework that uses wavelet decomposition and spectral flow matching to synthesize realistic fMRI time series for brain disorder identification, addressing data scarcity and non-stationarity challenges.
Geodesic Flow Matching for Denoising High-Dimensional Structured Representations
This paper proposes Geodesic Flow Matching, a Riemannian transport method for denoising Spatial Semantic Pointers (SSPs) on toroidal manifolds, and demonstrates a 72% reduction in tracking error and 40% efficiency gain in a spiking neural SLAM system.
FM-fMRI: Event Conditioned Flow Matching for Rest-to-Task fMRI Time-Series Synthesis
Proposes FM-fMRI, an event-conditioned flow matching model that synthesizes task fMRI time series from resting-state fMRI, achieving superior spectral and connectivity agreement over baselines on the Human Connectome Project and an internal autism cohort, and improving downstream autism classification.
BrainG3N: A Dual-Purpose Tokenizer for Controllable 3D Brain MRI Generation
Introduces BrainG3N, a dual-purpose tokenizer for 3D brain MRI latent diffusion using a frozen masked autoencoder encoder for clinically informative embeddings and a CNN decoder for reconstruction, achieving state-of-the-art performance on a 23-task benchmark and enabling controllable generation and longitudinal forecasting.