PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2601.05212v2cs.CV

TL;DR

Demographically imbalanced 3D brain MRI datasets limit the data available for reliable brain age prediction, while existing wavelet-based approaches remain computationally costly. FlowLet uses age-conditioned flow matching in an invertible wavelet domain to generate 3D MRIs with few sampling steps, preserving anatomical structure and improving brain age prediction performance.

  • Problem

    Demographically imbalanced 3D brain MRI datasets limit the availability of large, diverse, and age-balanced data needed for reliable brain age prediction.

  • Method

    FlowLet performs age-conditioned flow-matching generation directly in an invertible multi-scale wavelet representation for 3D brain MRI synthesis.

  • Results

    FlowLet achieves efficient generation with few sampling steps while preserving anatomically meaningful structure and improving performance on brain age prediction.

  • Takeaways & Limitations

    FlowLet provides a scalable, controllable approach to 3D MRI synthesis with anatomy-aware evaluation and clinically relevant downstream utility.

  • Takeaways & Limitations

    Region-based metrics and downstream brain age prediction do not replace expert-driven clinical assessment of anatomical realism.

Abstract

from arXiv · show

Brain Magnetic Resonance Imaging (MRI) plays a central role in studying neurological development, aging, and diseases. One key application is Brain Age Prediction (BAP), which estimates an individual's biological brain age from MRI data. Effective BAP models require large, diverse, and age-balanced datasets, whereas existing 3D MRI datasets are demographically skewed, limiting fairness and generalizability. Acquiring new data is costly and ethically constrained, motivating generative data augmentation. Current generative methods are often based on latent diffusion models, which operate in learned low dimensional latent spaces to address the memory demands of volumetric MRI data. However, these methods are typically slow at inference, may introduce artifacts due to latent compression, and are rarely conditioned on age, thereby affecting the BAP performance. In this work, we propose FlowLet, a conditional generative framework that synthesizes age-conditioned 3D MRIs by leveraging flow matching within an invertible 3D wavelet domain, helping to avoid reconstruction artifacts and reducing computational demands. Experiments show that FlowLet generates high-fidelity volumes with few sampling steps. Training BAP models with data generated by FlowLet improves performance for underrepresented age groups, and region-based analysis confirms preservation of anatomical structures.

1. Introduction

The introduction frames age-balanced, anatomically faithful 3D MRI synthesis as necessary for reliable Brain Age Prediction, which is limited by demographic imbalance, acquisition constraints, and computational challenges. It presents FlowLet as a conditional Flow Matching framework using invertible wavelets and multi-level age conditioning, evaluated for image fidelity, anatomy, and downstream BAP utility.

  • Brain Age Prediction estimates biological brain age from structural MRI, and its brain-age gap is clinically linked to cognitive decline, neurodegeneration, and altered aging.
  • BAP reliability depends on large, diverse, age-balanced datasets, but available 3D MRI data overrepresent young and middle-aged adults while undersampling pediatric and elderly populations.
  • Acquisition costs and ethical constraints motivate synthetic MRI generation for scalable augmentation, bias mitigation, and robustness improvement.
  • Full-resolution 3D MRI synthesis is challenged by high dimensionality and multi-scale anatomy, while existing autoencoding and diffusion approaches trade computational efficiency against anatomical fidelity.
  • FlowLet combines Flow Matching with an invertible wavelet representation and feature-wise and spatial age conditioning for efficient, anatomically faithful, localized 3D MRI synthesis.The framework is evaluated across three neuroimaging datasets spanning over 12 multi-site sources using global metrics, downstream BAP, and region-based anatomical analysis.

2. Related Work

Related work highlights the need for generative augmentation in 3D neuroimaging and the trade-offs among GANs, VAEs, diffusion, latent, wavelet, and flow-matching approaches. FlowLet targets the remaining gap between efficient generation, artifact avoidance, and scalable high-resolution volumetric synthesis.

  • Data augmentation: 3D neuroimaging datasets remain limited by small cohorts and high-dimensional voxel data, restricting population diversity and statistical reliability.Deep learning models are particularly affected because they require large, diverse datasets to generalize well.
  • Generative modeling: Generative models enable anatomically plausible medical-image synthesis beyond geometric or intensity-based augmentation, with GANs and VAEs establishing early foundations.GANs can produce sharp samples but often suffer from training instability and mode collapse.
  • Diffusion models: Diffusion models achieve high-quality, diverse synthesis but typically require hundreds or thousands of iterative sampling steps.Their generation reverses a gradual noising process through repeated evaluations.
  • Compressed and wavelet domains: Latent and wavelet approaches reduce computational or memory demands, but learned latent compression can introduce artifacts and diffusion-based wavelet models still require costly iterative inference.Wavelet methods use fixed multi-resolution frequency decompositions that preserve spatial structure and avoid learned-compression artifacts.
  • Flow Matching: Flow Matching addresses diffusion inefficiency by learning a continuous-time velocity field and transporting samples through an ordinary differential equation, motivating its integration with wavelet-domain generation.The approach encourages straighter transport trajectories and reduces the number of inference steps required.

3. Proposed method

FlowLet synthesizes age-conditioned 3D brain MRIs by combining an invertible 3D Haar wavelet representation with conditional flow matching. Its single-stage pipeline decomposes volumes, learns velocity fields in the multi-scale wavelet domain, and reconstructs full-resolution MRIs through the inverse transform.

  • Framework overview: FlowLet’s three-stage pipeline performs wavelet decomposition, flow-matching generation in the multi-scale frequency domain, and inverse-wavelet reconstruction.The framework operates entirely in the wavelet domain.
  • Wavelet representation: The invertible 3D Haar DWT produces one low-frequency subband for coarse anatomy and seven high-frequency subbands for localized fine details.The eight subbands are concatenated into an 8-channel tensor while preserving structural information.
  • Conditional flow matching: A conditional neural network learns continuous-time velocity fields that transport Gaussian-noise samples to the target wavelet-domain distribution, with variables such as age used for conditioning.The model is trained with an MSE objective on parameterized velocity predictions.
  • Sampling and reconstruction: At inference, FlowLet integrates the learned velocity-field ODE from Gaussian noise and applies the inverse DWT to produce a full-resolution 3D brain MRI.The ODE is solved using Euler integration.
  • Flow-matching variants: FlowLet supports RFM, CFM, VP, and Trigonometric Flow, spanning straight, nonlinear, and circular interpolation paths with different trajectory curvature.CFM introduces non-zero curvature and greater sensitivity as t approaches 1, whereas RFM uses a straight path with constant velocity.

4. Datasets and preprocessing

The training cohort integrates cognitively normal T1-weighted MRI scans from OpenBHB, ADNI, and OASIS-3 to broaden age coverage for age-conditioned synthesis and BAP evaluation. All scans undergo a unified preprocessing pipeline, including registration, skull stripping, resampling, and intensity normalization.

  • Datasets: The training cohort combined cognitively normal T1-weighted MRI scans from OpenBHB, ADNI, and OASIS-3 to address OpenBHB’s younger-age imbalance and improve older-age coverage.This integration supports age-conditioned synthesis and downstream BAP evaluation.
  • Datasets: ADNI contributed 769 scans from cognitively normal individuals aged 60–91 years, retaining one scan per subject and excluding cognitive impairment.The cohort mean age was 76.97 ± 4.99 years.
  • Datasets: OASIS-3 contributed 1,041 cognitively normal T1-weighted MRIs covering ages 42–95 years after filtering an initial pool of 1,314 scans.The retained scans had a mean age of 71.10 ± 8.93 years.
  • Age distribution: Integrating the datasets produced younger-age concentration from OpenBHB alongside richer representation of ages 60–95 years from ADNI and OASIS-3.The combined distribution provides broader coverage across adulthood for learning age-dependent anatomical variability.
  • Evaluation splits: Generative evaluation used a 20% split of the integrated training distribution, whereas BAP evaluation used the independent OpenBHB validation set.The generative split preserved the integrated cohort’s age profile, while the BAP benchmark followed OpenBHB’s original age distribution.

5. Experimental setup

The experimental setup compares FlowLet with seven 3D brain MRI synthesis baselines, including age-conditioned variants for fair age-conditioned synthesis and BAP evaluation. Experiments use standardized training and controlled ablations covering FlowLet’s major design choices.

  • Baseline models: FlowLet is compared with WDM, MD, MLDM, BS, MOTFM, and age-conditioned variants WDMa and MOTFMa.The baselines include wavelet diffusion, medical diffusion, latent diffusion, VQ-VAE/Transformer, and flow-matching models.
  • Baseline models: Age conditioning for WDMa and MOTFMa normalizes scalar age, projects it into a 512-dimensional embedding, and injects it into the U-Net through cross-attention.These variants adapt originally unconditional public implementations for fair age-conditioned synthesis and downstream BAP comparison.
  • Baseline models: Subject-specific brain aging models are excluded because they transform an existing subject MRI and preserve identity rather than synthesize from noise.The paper treats this as a fundamentally different task from the evaluated synthesis setting.
  • Implementation: FlowLet uses PyTorch, AdamW, cosine annealing, mixed precision, and optional xformers attention, with experiments run on an NVIDIA A6000 GPU with 48 GB VRAM.FlowLet can be trained in 24 GB, whereas baseline models required the larger memory capacity.
  • Ablation studies: Controlled ablations vary the wavelet basis, inference steps, conditioning mechanisms, and numerical ODE solver while assessing each component in isolation.FiLM and spatial cross-attention are among the conditioning mechanisms examined.

6. Evaluation metrics

The evaluation combines global fidelity, distributional alignment, diversity, and region-based anatomical metrics with downstream brain age prediction (BAP) performance. It also accounts for conditioning-age effects and uses age-matched regional comparisons to assess clinical and anatomical fidelity.

  • Global metrics: FID and MMD measure synthetic–real distributional alignment using medical-image ResNet-50 features, with lower values indicating closer alignment.Both metrics evaluate image fidelity and distributional similarity between generated and real samples.
  • Global metrics: Intra-set MS-SSIM assesses sample diversity, where lower values indicate greater inter-sample variability and higher values may suggest mode collapse.Pairwise comparisons use two-sided Wilcoxon rank-sum tests with Bonferroni correction at α= 0.05.
  • Age-stratified evaluation: Fidelity and diversity metrics are additionally computed within non-overlapping age bins of 15–30, 40–55, and 65–80 years to account for conditioning-range effects.Intra-set MS-SSIM depends on the age interval over which samples are synthesized.
  • Region-based evaluation: Region-based evaluation uses one-to-one age-matched real references and 95 FastSurfer5 cortical and subcortical ROIs to prevent mode collapse from appearing artificially strong.Metrics are computed independently per ROI and then averaged across the full brain.
  • Region-based evaluation: Regional anatomical plausibility is summarized with iMAE, KLD, and DICE over unified real–synthetic voxel supports, with lower iMAE/KLD and higher DICE indicating better fidelity.Dice ranges from 0 for no overlap to 1 for perfect overlap.
  • Downstream clinical evaluation: Clinical usefulness is evaluated through BAP, which predicts chronological age from structural T1-weighted MRI using a 3D DenseNet-121 regression model.The protocol focuses on cognitively normal subjects because older adults are underrepresented and are most susceptible to cognitive decline.

7. Results

FlowLet achieves strong image fidelity, diversity, sampling efficiency, and anatomical consistency, while improving downstream brain age prediction through explicit age conditioning. Results also identify Haar wavelets and 10 sampling steps as effective choices, whereas Trigonometric paths and incomplete conditioning degrade performance.

  • Generative performance: All FlowLet variants achieve competitive FID and MMD using only 10 ODE sampling steps, outperforming diffusion-based baselines such as WDM, MD, and MLDM.Real-to-real baselines provide additional context for the reported global similarity values.
  • Generative performance: FlowLet preserves realistic sample diversity while improving global distributional alignment, with intra-set MS-SSIM comparable to or lower than baseline models.The results indicate that improved similarity does not arise from mode collapse or reduced inter-sample variability.
  • Sampling efficiency: 10 sampling steps provide an effective quality-efficiency trade-off for RFM, CFM, and VP, while the Trigonometric path becomes unstable and degrades at 200 steps.FID improves with additional steps before saturating at approximately 10 steps for the stable formulations; BAP similarly improves rapidly and then saturates.
  • Downstream brain age prediction: FlowLet-generated data lowers BAP test MAE relative to real-only training and baseline-generator augmentation, with RFM and CFM achieving the best FlowLet performance.The gains are attributed to explicit age-conditioned synthesis, which provides complementary age-relevant information and mitigates dataset imbalance.
  • Region-based evaluation: FlowLet variants, particularly RFM and VP, rank among the strongest methods on regional anatomical metrics, whereas global fidelity does not guarantee fine-grained anatomical preservation.The MD baseline illustrates this discrepancy with a DICE score of 0.294 despite relatively low global FID.
  • Ablation studies: Haar provides the best wavelet reconstruction fidelity and stable generative behavior, with a mean MAE of 6.08 × 10^-8, while removing either conditioning component substantially harms evaluation outcomes.Spatial-only and FiLM-only variants have higher BAP error and reduced anatomical consistency, and the unconditional variant performs worst overall.

8. Discussion

FlowLet combines invertible wavelet-domain flow matching with explicit age conditioning to deliver efficient, anatomically meaningful 3D MRI synthesis. The discussion emphasizes that transport-path geometry, conditioning design, and regional/downstream evaluation are more informative than global fidelity or solver order alone.

  • Efficiency and representation: FlowLet achieves competitive or superior global fidelity while requiring an order-of-magnitude fewer sampling steps through flow matching in an invertible wavelet domain.The wavelet representation preserves exact invertibility and fine-grained spatial detail without learned compression.
  • Efficiency and representation: Performance saturates at approximately ten ODE steps across FlowLet variants, indicating that smooth trajectories in a multi-scale representation make long stochastic sampling paths unnecessary.Unlike diffusion models, FlowLet remains stable at low step counts rather than relying on additional steps to reduce noise-induced artifacts.
  • Efficiency and representation: Approximately 8× lower memory consumption enables batch size 4 using approximately 22 GB of VRAM, whereas WDM and MLDM require over 40 GB under comparable settings.These resource savings support scalability to large neuroimaging cohorts and research environments with limited computational resources.
  • Age conditioning: Age conditioning is essential because unconditional models can achieve competitive global fidelity yet show weaker regional anatomy and poor brain age prediction performance.The age-conditioned design targets coherent age trajectories and biologically meaningful control for downstream age-related applications.
  • Age conditioning: Combining FiLM feature-wise conditioning with spatially adaptive cross-attention consistently preserves age-relevant anatomical variations, unlike either conditioning mechanism alone.The MOTFM comparison further indicates that effective conditioning, rather than architectural differences alone, drives downstream gains.
  • Flow geometry and evaluation: Straight or near-straight transport paths provide more stable global, regional, and downstream performance than higher-order solver accuracy or curved interpolation paths.Replacing Euler with RK4 does not improve the Trigonometric formulation, while curved paths become less anatomically consistent as steps increase.

9. Limitations and future directions

The study’s evaluation does not replace expert clinical assessment and remains limited in conditioning scope, fairness and privacy auditing, and imaging modality. Future work will pursue neuroradiologist validation, multi-attribute conditioning, broader demographic and privacy analyses, additional modalities, and uncertainty estimation.

  • Limitations and future directions: Region-based anatomical metrics and brain age prediction are stronger proxies than global similarity scores but do not replace expert-driven clinical assessment.Future work should incorporate structured evaluations by neuroradiologists to validate anatomical realism.
  • Limitations and future directions: FlowLet’s explicit conditioning is currently limited to age, while multi-attribute conditioning and covariate-based matching require systematic investigation.Potential attributes include sex, pathology, and cognitive scores; synthetic and real subjects are currently matched solely by age because other covariates are not consistently available.
  • Limitations and future directions: Comprehensive fairness audits, demographic-subgroup robustness analyses, and privacy-leakage assessments remain open challenges for deployment in sensitive clinical settings.The preprocessing pipeline mitigates spurious correlations, but these broader evaluations remain necessary.
  • Limitations and future directions: The study is limited to T1-weighted structural MRI, with future work targeting other 3D imaging modalities and uncertainty estimation.Uncertainty estimation is intended to better characterize the reliability of synthetic cohorts.

10. Conclusion

FlowLet combines flow matching with an invertible wavelet representation for conditional 3D brain MRI synthesis, enabling efficient generation while preserving anatomically meaningful structure. Its global fidelity improvements support regional plausibility and improved brain age prediction, positioning it as a scalable, controllable approach to 3D MRI synthesis.

  • 10. Conclusion: FlowLet combines flow matching with an invertible wavelet representation for conditional 3D brain MRI synthesis, avoiding learned latent compression.This design supports efficient generation with few sampling steps.
  • 10. Conclusion: Improvements in global fidelity translate into stronger regional anatomical plausibility and improved performance on brain age prediction.The evaluation connects image quality with preservation of meaningful anatomy and a clinically relevant downstream task.
  • 10. Conclusion: FlowLet provides a scalable, open-source, and controllable approach to 3D MRI synthesis, while limitations remain.The work emphasizes efficient generative modeling, explicit conditioning, and anatomy-aware evaluation for volumetric medical images.

A. Appendix

The appendix provides implementation details and extended experiments supporting FlowLet, covering efficiency, qualitative generation, age-stratified performance, wavelet analyses, and significance testing.

  • The appendix documents Flow Matching implementations, hyperparameters, and efficiency benchmarks.
  • It reports step-wise qualitative results and age-stratified metrics for the generated MRIs and downstream evaluation.
  • Additional analyses examine wavelet representations and statistical significance.

A.1. Flow Matching implementations

The implementations train flow-matching models by minimizing MSE between predicted and target velocity fields, with uniformly sampled time and wavelet-domain data/noise variables. RFM, CFM, and Trigonometric Flow use distinct interpolation paths and corresponding target velocities.

  • A.1. Flow Matching implementations: All implementations minimize MSE between predicted and target velocity fields, sample t uniformly from [0, 1], and represent data and noise as x1_wavelet and x0_wavelet.The noise sample is x0 ∼ N(0, I).
  • A.1. Flow Matching implementations: RFM implements linear interpolation, xt = (1 - t_broadcast) * x0_wavelet + t_broadcast * x1_wavelet, with v_target = x1_wavelet - x0_wavelet.This corresponds to the path xt = (1 − t)x0 + tx1 and constant velocity v_target = x1 − x0.
  • A.1. Flow Matching implementations: CFM computes the state-dependent target velocity as v_target = (x1_wavelet - xt) / (1 - t_broadcast + 1e-8) using the same linear interpolation.The added ϵ = 10^-8 provides numerical stability.
  • A.1. Flow Matching implementations: Trigonometric Flow uses circular interpolation, xt = cos(πt/2)x0 + sin(πt/2)x1, and computes velocity as its time derivative.The implementation uses angle for πt/2 and v_target = -sin(angle)(π/2)x0_wavelet + cos(angle)(π/2)x1_wavelet.

A.2. Variance-Preserving (VP) Diffusion Matching

This section connects VP diffusion matching’s SDE formulation to the deterministic velocity field used by FlowLet. It derives a computable conditional training target by replacing the unknown marginal score with an analytical conditional score.

  • Core definitions and time conventions: VP diffusion evolves continuously from clean data at t = 0 to pure noise at t = 1 under a linear variance schedule.The implementation uses βmin = 0.1 and βmax = 20.0.
  • Core definitions and time conventions: FlowLet samples generative-path time t_flow ∈ [0, 1] and maps it to theoretical data-to-noise time through t = 1 − t_flow.All subsequent VP formulas use the theoretical time t.
  • From stochastic SDE to deterministic ODE: The Probability Flow ODE preserves the diffusion process’s marginals and yields the deterministic VP flow that FlowLet trains to approximate.The ODE construction explains the factor 1/2 multiplying the score term.
  • Deriving the computable target velocity: Because the true marginal score is unknown, conditional flow matching substitutes the analytical conditional score from Tweedie’s formula to obtain a computable target velocity.Substituting the conditional score into the ODE velocity field produces the target learned by the network.

A.3. Training hyperparameters … A.14. External DLBS validation

The supplementary analyses specify FlowLet’s reproducible training and preprocessing, characterize efficiency and sampling behavior, and evaluate age conditioning, wavelet representations, and external generalization. Across these analyses, FlowLet preserves anatomical and high-frequency structure, supports efficient age-conditioned synthesis, and improves brain-age prediction on an external cohort.

  • A.3. Training hyperparameters: FlowLet training used deterministic PyTorch, NumPy, and CuDNN behavior, a fixed 80%–20% split of 5,794 cognitively normal subjects, wavelet preprocessing, and standardized optimization settings.The split produced 4,635 training and 1,159 validation subjects; volumes were converted to eight wavelet subbands, and final selection used the lowest validation MSE.
  • A.4. Efficiency benchmarks: FlowLet reduces computational demands, using approximately 22 GB of training VRAM and achieving roughly a 45× inference speedup over the conditional diffusion baseline WDM.It scales to 256^3 volumes with 42 GB of VRAM at batch size 1, while sampling increases from roughly 1.6s at 112^3 to 6.8s at 256^3 using 10 ODE steps.
  • A.5. Additional qualitative assessment of Flow: RFM, CFM, and VP generate stable, coherent volumes with 2–5 ODE steps, whereas Trigonometric flow exhibits instability and structural artifacts in low-step regimes.At 200 steps, all formulations converge toward anatomically plausible structures under the fixed seed and age condition.
  • A.6. Additional age-stratified quantitative: Across age groups, increasing sampling steps monotonically improves FID before a plateau between 10 and 200 steps, supporting 10 steps as an efficiency–fidelity balance.The 65–80 group achieves the best FID and intra-set MS-SSIM performance, while the Trigonometric variant can become unstable at 200 steps.
  • A.7. Significance testing and P-Value analysis.: Statistical comparisons use FlowLet-RFM at 10 steps as the reference for FID, MMD, and MS-SSIM, with Bonferroni-corrected pairwise significance testing.The reported heatmaps identify non-significant differences using p>0.05, including results in the 200-step rows.
  • A.8. Additional wavelet evaluation: Haar provides the lowest reconstruction error, with mean MAE 6.08 × 10^-8, while longer-support or non-orthogonal wavelets show larger errors or localized edge artifacts.Coiflet-2 achieves the second-lowest reconstruction error at MAE ∼10^-8, whereas db4 and bior3.3 are around ∼10^-7.
  • A.9. Low-frequency ablations: Low-frequency ablations remain competitive, but full-wavelet training with high-frequency suppression at inference is consistently stronger than training exclusively on the approximation band.The results indicate that coarse anatomy carries a substantial fraction of the measurable evaluation signal, while full-wavelet training retains additional strength.
Loading 2601.05212v2…