Search papers, labs, and topics across Lattice.
This paper introduces WaveDiT, a conditional flow matching framework that synthesizes full-resolution 3D brain MRI images efficiently by operating in the coefficient space of a 3D Haar Discrete Wavelet Transform. By integrating factorized spatio-depth attention and modeling heteroscedastic uncertainty, WaveDiT achieves high-quality synthesis while maintaining computational feasibility on a single modern GPU. The results show that WaveDiT significantly improves the alignment of generated MRI distributions with real data and enhances performance in brain age prediction and anatomical accuracy compared to existing methods.
WaveDiT synthesizes high-fidelity 3D brain MRIs on a single GPU, achieving superior anatomical detail without the computational overhead of traditional methods.
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
