WaveDiT: 分布感知的小波流匹配用于高效3D脑部MRI合成

摘要

大规模且人口平衡的数据集对于可靠的神经影像生物标志物至关重要。在此背景下，全分辨率三维脑部MRI合成可支持数据增强，但现有方法要么在体素尺度上产生高昂的计算成本，要么依赖有损潜在压缩而可能损害解剖细节。因此，实用的三维生成增强通常需要专门的计算基础设施。我们提出WaveDiT，一种在三维Haar离散小波变换系数空间中运行的条件流匹配框架。该模型将分解的深度-空间注意力与基于高阶小波统计的带状异方差不确定性建模相结合。预测的对数方差直接集成到流目标函数和条件路径中，从而能够实现与解剖细节的重尾和输入依赖方差结构相一致的适应性精度。该公式支持在单块现代GPU上在实用的内存和时间约束下进行全分辨率三维合成。在多中心队列上的评估表明，与扩散、潜在和小波基线方法相比，生成图像与真实MRI分布之间的对齐程度有所改善，同时下游脑年龄预测和区域级解剖一致性也得到增强。代码已公开于https://github.com/sisinflab/WaveDiT。

English

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