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

2026-01-08

arXiv.org

AI Summary

Existing 3D brain MRI datasets exhibit demographic biases that hinder the fairness and generalization of brain age prediction models, while conventional generative approaches suffer from limitations in inference speed, reconstruction fidelity, and precise age conditioning. To address these challenges, this work proposes FlowLet—the first generative framework that integrates conditional flow matching within an invertible 3D wavelet domain. Trained end-to-end, FlowLet enables high-fidelity synthesis with fast, few-step sampling while accurately controlling the age attribute of generated images. By operating directly in the wavelet domain, the method avoids artifacts caused by latent space compression, preserving fine anatomical details. This approach significantly improves brain age prediction performance on underrepresented age groups, thereby enhancing the diversity and balance of training data.

Technology Category

• Computer Vision: Generative Adversarial Networks (GANs) for Vision
• Natural Language Processing: Generation
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Web Mining and Content Analysis: Web data generation and simulation

Abstract

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.

Problem

• Brain Age Prediction
• 3D MRI Synthesis
• Data Augmentation
• Generative Modeling
• Age-Conditioned Generation

Innovation

• Flow Matching
• Wavelet Domain
• Conditional Generation
• 3D Brain MRI Synthesis
• Data Augmentation