FlowLet: Age-Conditioned 3D MRI Synthesis

• FlowLet is a conditional generative framework for age-conditioned 3D MRI synthesis that leverages an invertible 3D wavelet domain and continuous-time flow matching to produce synthetic MRIs.
• The method employs a 3D Haar DWT for exact invertibility, achieving minimal round-trip error (~6e-8 MAE) and rapid synthesis speeds (~1.6 seconds per volume).
• FlowLet improves downstream brain age prediction by reducing MAE from 4.91 to 4.01 years, promoting fairness and enhanced anatomical fidelity.

FlowLet refers to a conditional generative framework for age-conditioned 3D MRI synthesis that leverages continuous-time flow matching within an invertible 3D wavelet domain. Its primary application is the augmentation of brain MRI datasets for improving fairness and generalizability in brain age prediction (BAP) models. FlowLet directly addresses limitations of conventional latent diffusion models for 3D volumetric medical imaging, including slow inference speeds, compression-induced artifacts, and a lack of explicit conditioning. The framework is characterized by exact invertibility, rapid sampling, and demonstrated preservation of fine-grained anatomical structure in synthetic images (Danese et al., 8 Jan 2026).

1. Invertible 3D Wavelet Domain Representation

FlowLet operates in the orthonormal 3D Haar Discrete Wavelet Transform (DWT) domain. A 3D MRI volume $x \in \mathbb{R}^{D \times H \times W}$ is mapped via Haar DWT $W$ into eight spatial-frequency subbands (LLL, LLH, ..., HHH). The forward transform,

$$x_w = W(x) \in \mathbb{R}^{8 \times (D/2) \times (H/2) \times (W/2)},$$

is exactly invertible via $W^{-1}$. Parseval’s theorem ensures energy preservation, supporting stable learning dynamics. The invertibility property is maintained throughout the generative process, eliminating reconstruction artifacts and negating the need for a learned decoder. Haar wavelets yielded the lowest round-trip MAE ($\sim 6 \times 10^{-8}$) among common alternatives, providing the most stable and faithful anatomical synthesis (Danese et al., 8 Jan 2026).

2. Continuous-Time Flow Matching in Latent Space

FlowLet adopts continuous-time Flow Matching (FM) to facilitate rapid and accurate generation of synthetic 3D MRIs. FM defines a time-dependent velocity field $v_\theta(x, t, c)$ that transports samples from a Gaussian prior $x(0) \sim \mathcal{N}(0, I)$ in wavelet space to the target data distribution $x(1) = W(x_{\mathrm{real}})$, conditioned on metadata $c$ (such as age). The generative process solves the ODE:

$$\frac{dx(t)}{dt} = v_\theta(x(t), t, c), \qquad t \in [0, 1].$$

Sampled paths and velocity fields are generated via several path families, including:

• Rectified Flow Matching (RFM): $x_t = (1-t)x_0 + t x_1$, $v_{\mathrm{target}}(x_t, t) = x_1 - x_0$
• Conditional Flow Matching (CFM): $$v_{\mathrm{target}}(x_t, t) = \frac{x_1 - x_t}{1-t + \varepsilon}$$
• Variance-Preserving Diffusion and Trigonometric Flow routes, with the analytic conditional score provided via Tweedie’s formula.

Training minimizes the mean-squared error between $v_\theta(x_t, t, c)$ and $v_{\mathrm{target}}(x_t, t)$ over random interpolations in $t$ (Danese et al., 8 Jan 2026).

3. Explicit Age-Conditioning Architecture

Age-conditioning in FlowLet is achieved by normalizing scalar age to $[0, 1]$ and projecting it to a 512-dimensional embedding $e_{\mathrm{cond}}$ using a two-layer MLP with SiLU activations. In multi-condition settings, individual embeddings are summed. This conditioning signal is injected into each residual block of a conditional 3D U-Net backbone via FiLM modulation:

$$h_{\mathrm{mod}} = \mathrm{Norm}(h)(1 + \gamma_t) + \beta_t,$$

$h_{\mathrm{mod}} = h'(1 + \gamma_c) + \beta_c,$

where parameters $(\gamma_t, \beta_t)$, $(\gamma_c, \beta_c)$ are derived from temporal and condition embeddings, respectively. Additionally, spatial cross-attention layers at coarser scales enable age information to modulate global and regional features, improving structure preservation and control (Danese et al., 8 Jan 2026).

4. Sampling Procedure and Computational Efficiency

Synthetic volumes are generated by integrating the ODE in wavelet space using a fixed-step Euler solver (typically $K=10$ steps), starting at $x(0) \sim \mathcal{N}(0, I)$, then applying the inverse DWT to reconstruct the MRI volume $x_{\mathrm{MRI}} = W^{-1}(x(1))$. FlowLet achieves high-fidelity synthetic generation with as few as 5–10 steps (FID~0.30), distinctly faster than diffusion-based methods, which require several hundred to a thousand steps for similar fidelity (FlowLet: $\approx 1.6$ seconds/volume on an A6000 GPU vs. $\approx 70$ seconds/volume for 1,000-step diffusion) (Danese et al., 8 Jan 2026).

5. Quantitative Results and Anatomical Fidelity

FlowLet was evaluated on 5,794 T1-weighted MRIs from OpenBHB, OASIS-3, and ADNI, processed to a common template. Key metrics:

• Global fidelity and diversity: FID = 0.298, MMD = 0.0119, MS-SSIM = 0.9508 (all for RFM, 10 steps)—competitive or superior to conditional latent diffusion (MLDM) and wavelet diffusion (WDM) baselines.
• Downstream BAP impact: Using a 3D DenseNet-121, BAP models trained on FlowLet-augmented data achieved lower MAE for held-out real subjects over age 44 (MAE = $4.01 \pm 3.38$ years) compared to models trained on real data alone (MAE = $4.91 \pm 3.92$ years), with the greatest improvement in underrepresented older age groups.
• Fine-grained anatomical preservation: On region-based anatomical metrics (mean voxel error, histogram KL divergence, Dice overlap of 95 ROIs), FlowLet outperforms baselines, e.g., $i$MAE = 37.7, KL = 0.855, Dice = 0.420, indicating accurate retention of neuroanatomical detail.

Ablation studies show that the Haar wavelet, step count $\geq$10, and explicit conditional modulation (both FiLM and spatial attention) are necessary for optimal fidelity and downstream utility (Danese et al., 8 Jan 2026).

6. Network and Implementation Details

FlowLet’s generative backbone is a conditional 3D U-Net with four down/up levels. It uses a base channel count of 32 with multipliers (1,2,4,8), GroupNorm and SiLU activations, and skip connections between encoder and decoder features. Mixed-precision training (batch size 4) is supported via PyWavelets and xformers for efficient memory use (total $\sim$85M parameters, $\sim$24 GB VRAM per GPU). The process is strictly invertible at all stages—$\hat{x} = W^{-1}(x_w)$ recovers the MRI exactly. No latent-space compression is used (Danese et al., 8 Jan 2026).

7. Implications, Limitations, and Significance

FlowLet establishes that exact invertible wavelet regimes combined with continuous-time flow matching significantly accelerate and improve the fidelity of large-scale volumetric MRI synthesis, yielding samples that enhance performance and equity in downstream clinical prediction tasks. The substantially reduced computational cost and explicit conditional control distinguish FlowLet from diffusion-based competitors, particularly in settings where synthetic augmentation supports demographic balancing.

Limitations include current restriction to T1-weighted MRI, dependency on preprocessing pipelines, and potential instability of certain ODE solvers (e.g., trigonometric flow at high step counts). A plausible implication is that extension to multimodal or pathology-aware conditioning could further advance fairness and application breadth in neuroimaging.

References: (Danese et al., 8 Jan 2026)