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

• URL: http://arxiv.org/abs/2606.08670v1
• Date: Sun, 07 Jun 2026 15:21:02 GMT
• Title: WaveDiT: Distribution-Aware Wavelet Flow Matching for Efficient 3D Brain MRI Synthesis
• Abstract summary: Large demographically balanced datasets are essential for reliable neuroimaging biomarkers.<n>WaveDiT is a conditional flow matching framework operating in the coefficient space of a 3D Haar Discrete Wavelet Transform.<n>It supports full-resolution 3D synthesis under practical and memory time constraints on a single modern GPU.
• Score: 7.371811584771131
• License: http://creativecommons.org/licenses/by/4.0/
• 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

Related papers

• ProgFormer: Hierarchical Voxel Diffusion Transformer for Longitudinal Brain MRI Prediction [81.07048249760953]
  longitudinal changes are often subtle and confined to specific anatomical regions.<n>Most subject-specific brain structure remains stable over time.<n>ProgFormer is a hierarchical voxel-space Diffusion Transformer for longitudinal brain MRI prediction.
  arXiv  Detail & Related papers  (2026-07-30T00:13:42Z)
• Anatomically-conditioned Latent Diffusion Model for Data-Efficient Few-Shot Cross-Domain 3D Glioma MRI Synthesis [1.4312099612976772]
  Anatomically-conditioned Latent Diffusion Model (ALDM) is a novel framework for data-efficient, few-shot 3D MRI synthesis.<n> evaluated in an extreme few-shot setting with only 16 target images.
  arXiv  Detail & Related papers  (2026-06-24T04:42:42Z)
• Wavelet-Fusion Diffusion Model for Multimodal Brain MRI Synthesis with Modality and Metadata Conditioning [0.0]
  Multimodal MRI provides complementary information for neuroimaging analysis, where different imaging modalities capture distinct anatomical, tissue, and pathological features.<n>Large-scale structural MRI resources are increasingly available, but their modality coverage is often uneven across public and pooled neuroimaging datasets.<n>Our approach combines a Wavelet-Fusion variational autoencoder (WF-VAE) latent compressor with a conditional 3D U-Net diffusion model trained in the learned latent space.
  arXiv  Detail & Related papers  (2026-05-30T11:57:27Z)
• EPC-3D-Diff: Equivariant Physics Consistent Conditional 3D Latent Diffusion for CBCT to CT Synthesis [36.790666885042356]
  We propose EPC-3D-Diff, a novel conditional 3D latent diffusion framework for CBCT to CT synthesis.<n>We exploit the fact that an in plane rotation of the volume corresponds to an angular shift in its projections.<n>We validate on a paired head CBCT/CT phantom dataset, including repeat scans, and paired clinical data using patient wise splits.<n> EPC-3D-Diff generalizes well and achieved substantial improvements, +7.4 dB (phantom) and +1.8 dB (clinical data) in PSNR compared to state of the art methods.
  arXiv  Detail & Related papers  (2026-05-19T20:34:09Z)
• MedPruner: Training-Free Hierarchical Token Pruning for Efficient 3D Medical Image Understanding in Vision-Language Models [59.180043227905294]
  We propose MedPruner, a training-free and model-agnostic hierarchical token pruning framework for efficient 3D medical image understanding.<n>We show that MedPruner enables models such as MedGemma to maintain or even exceed their original performance while retaining fewer than 5% of visual tokens.
  arXiv  Detail & Related papers  (2026-03-12T07:37:00Z)
• Exploiting Completeness Perception with Diffusion Transformer for Unified 3D MRI Synthesis [9.857855424798732]
  We propose CoPeDiT, a latent diffusion model equipped with completeness perception for unified synthesis of 3D MRIs.<n>CoPeDiT significantly outperforms state-of-the-art methods, achieving superior robustness, generalizability, and flexibility.
  arXiv  Detail & Related papers  (2026-02-20T18:05:39Z)
• Improving 2D Diffusion Models for 3D Medical Imaging with Inter-Slice Consistent Stochasticity [47.52394044948656]
  Inter-Slice Consistentity (ISCS) is a strategy that encourages interslice consistency during diffusion sampling.<n>ISCS can be dropped into any 2D trained diffusion based 3D reconstruction pipeline without additional computational cost.
  arXiv  Detail & Related papers  (2026-02-04T02:54:08Z)
• FlowLet: Conditional 3D Brain MRI Synthesis using Wavelet Flow Matching [7.371811584771131]
  We propose FlowLet, a conditional generative framework that synthesizes age-conditioned 3D MRIs.<n> Experiments show that FlowLet generates high-fidelity volumes with few sampling steps.
  arXiv  Detail & Related papers  (2026-01-08T18:36:29Z)
• Adapting HFMCA to Graph Data: Self-Supervised Learning for Generalizable fMRI Representations [57.054499278843856]
  Functional magnetic resonance imaging (fMRI) analysis faces significant challenges due to limited dataset sizes and domain variability between studies.<n>Traditional self-supervised learning methods inspired by computer vision often rely on positive and negative sample pairs.<n>We propose adapting a recently developed Hierarchical Functional Maximal Correlation Algorithm (HFMCA) to graph-structured fMRI data.
  arXiv  Detail & Related papers  (2025-10-05T12:35:01Z)
• Accelerating 3D Photoacoustic Computed Tomography with End-to-End Physics-Aware Neural Operators [74.65171736966131]
  Photoacoustic computed tomography (PACT) combines optical contrast with ultrasonic resolution, achieving deep-tissue imaging beyond the optical diffusion limit.<n>Current implementations require dense transducer arrays and prolonged acquisition times, limiting clinical translation.<n>We introduce Pano, an end-to-end physics-aware model that directly learns the inverse acoustic mapping from sensor measurements to volumetric reconstructions.
  arXiv  Detail & Related papers  (2025-09-11T23:12:55Z)
• A Skull-Adaptive Framework for AI-Based 3D Transcranial Focused Ultrasound Simulation [1.662610796043078]
  Transcranial focused ultrasound (tFUS) is an emerging modality for non-invasive brain stimulation and therapeutic intervention.<n>TFUScapes is the first large-scale, high-resolution dataset of tFUS simulations through anatomically realistic human skulls.<n>DeepTFUS is a deep learning model that estimates normalized pressure fields directly from input 3D CT volumes and transducer position.
  arXiv  Detail & Related papers  (2025-05-19T11:37:51Z)
• Leveraging Frequency Domain Learning in 3D Vessel Segmentation [50.54833091336862]
  In this study, we leverage Fourier domain learning as a substitute for multi-scale convolutional kernels in 3D hierarchical segmentation models. We show that our novel network achieves remarkable dice performance (84.37% on ASACA500 and 80.32% on ImageCAS) in tubular vessel segmentation tasks.
  arXiv  Detail & Related papers  (2024-01-11T19:07:58Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.