[Paper] FlowLet: Conditional 3D Brain MRI Synthesis using Wavelet Flow Matching

Source: arXiv - 2601.05212v1

Overview

The paper introduces FlowLet, a new conditional generative model that can synthesize realistic 3‑D brain MRI volumes tailored to a specific age. By operating in a wavelet‑based invertible space, FlowLet sidesteps the heavy memory and artifact problems of existing diffusion‑based generators, enabling fast, high‑quality MRI creation that can be used to balance age‑biased datasets for brain‑age prediction (BAP) tasks.

Key Contributions

• Age‑conditioned 3‑D MRI synthesis: Generates whole‑brain volumes that reflect a user‑specified chronological age.
• Wavelet flow‑matching architecture: Uses an invertible 3‑D wavelet transform combined with flow‑matching, eliminating the need for latent compression and reducing memory footprints.
• Fast inference: Produces high‑fidelity scans in just a few sampling steps, dramatically faster than latent diffusion pipelines.
• Improved BAP performance: Augmenting training data with FlowLet‑generated scans boosts brain‑age prediction accuracy, especially for under‑represented age groups.
• Anatomical fidelity verification: Region‑wise analyses demonstrate that synthesized volumes preserve key brain structures and tissue contrasts.

Methodology

1. Wavelet Decomposition – Each input MRI is decomposed into multi‑scale wavelet coefficients, yielding a set of high‑frequency detail maps and a low‑frequency approximation. This representation is invertible and far more compact than the raw voxel grid.
2. Conditional Flow Matching – A neural network learns a continuous-time vector field that transports a simple Gaussian distribution to the distribution of wavelet coefficients, conditioned on the target age. The flow‑matching objective directly matches the instantaneous velocity of the model to the optimal transport field, avoiding the costly iterative denoising steps of diffusion models.
3. Training – The model is trained on publicly available 3‑D brain MRI datasets (e.g., OASIS, ADNI) with associated age labels. The loss combines a flow‑matching term with a reconstruction penalty that ensures the inverse wavelet transform yields realistic intensity values.
4. Sampling – To generate a new scan, the desired age is fed as a conditioning vector, a Gaussian sample is drawn, and the learned flow is integrated for a small number of steps (often ≤ 4). The resulting wavelet coefficients are then reconstructed into a full‑resolution MRI via the inverse wavelet transform.

Results & Findings

• Visual quality – Qualitative inspection shows that FlowLet‑generated volumes retain fine cortical folds and subcortical structures without the blurring or checkerboard artifacts common in latent diffusion outputs.
• Quantitative fidelity – Measured by Fréchet Inception Distance (FID) adapted for medical images and by structural similarity (SSIM), FlowLet outperforms state‑of‑the‑art diffusion baselines by ~15 % on FID and ~0.04 on SSIM.
• Sampling speed – Generation takes ≈ 0.8 seconds per volume on a single NVIDIA A100 GPU, compared to 5–10 seconds for comparable diffusion models.
• BAP boost – Training a ResNet‑based brain‑age predictor on the original dataset plus 30 % synthetic scans improves mean absolute error (MAE) from 4.2 years to 3.6 years for the 60‑80 year age band, a 14 % relative gain.
• Region‑level consistency – Volumetric analyses of hippocampal and ventricular regions show < 2 % deviation from real scans, confirming anatomical realism.

Practical Implications

• Data balancing for clinical AI – Researchers can quickly generate age‑specific MRIs to fill gaps in longitudinal or cross‑sectional studies, leading to fairer models that generalize across the lifespan.
• Reduced acquisition costs – Hospitals and research consortia can augment limited imaging cohorts without the expense and ethical hurdles of additional scans.
• Fast prototyping – Because synthesis is near‑real‑time, developers can experiment with synthetic data pipelines (e.g., on‑the‑fly augmentation during model training) without incurring major compute overhead.
• Potential for other modalities – The wavelet‑flow framework is modality‑agnostic, opening doors to conditional synthesis of CT, PET, or even multimodal neuroimaging stacks.
• Regulatory‑friendly synthetic data – Since the generated volumes are derived from a learned distribution rather than direct patient images, they may be easier to share across institutions under privacy regulations.

Limitations & Future Work

• Dataset diversity – FlowLet was trained on a handful of publicly available cohorts; performance on scans from different scanners, protocols, or pathologies (e.g., tumors) remains untested.
• Conditioning granularity – Age is the sole conditioning variable; extending to disease labels, genetics, or cognitive scores would increase utility.
• Evaluation scope – While FID/SSIM and region volumes are informative, downstream clinical validation (e.g., diagnostic accuracy when using synthetic data) is still pending.
• Scalability to ultra‑high resolution – The current wavelet depth balances memory and detail; future work could explore adaptive wavelet schemes to handle sub‑millimeter isotropic voxels.

FlowLet demonstrates that clever use of invertible transforms and flow‑matching can make high‑quality, conditionally controlled 3‑D medical image synthesis both fast and practical—an advance that could reshape how developers build and train neuroimaging AI systems.

Authors

• Danilo Danese
• Angela Lombardi
• Matteo Attimonelli
• Giuseppe Fasano
• Tommaso Di Noia

Paper Information

• arXiv ID: 2601.05212v1
• Categories: cs.CV
• Published: January 8, 2026
• PDF: Download PDF