WaveDiT用小波变换实现高效3D脑部MRI生成,单卡即可完成高分辨率合成。
WaveDiT: Distribution-Aware Wavelet Flow Matching for Efficient 3D Brain MRI Synthesis

- 在小波系数空间中构建条件流匹配模型,避免传统压缩损失细节。
- 支持全分辨率3D生成,在单张现代GPU上实现,内存和时间开销可控。
- 生成图像与真实数据分布更接近,提升脑龄预测和解剖结构一致性。
大规模且人口统计学均衡的数据集对可靠的神经影像生物标志物至关重要。全分辨率3D脑部MRI合成可支持此类场景下的数据增强,但现有方法或在体素尺度带来难以承受的计算成本,或依赖有损潜在压缩,可能损害解剖细节。因此,实际的3D生成增强通常需要专用计算基础设施。我们提出WaveDiT,一种在三维哈尓离散小波变换系数空间中运行的条件流匹配框架。该模型结合了分层时空注意力与基于高阶小波统计的带级异方差不确定性建模。预测的对数方差直接融入流目标函数和条件路径,实现与解剖细节重尾、输入相关的方差结构相一致的自适应精度。该设计使模型在单张现代GPU上,于合理内存和时间约束下,完成全分辨率3D合成。在多中心队列上的评估显示,生成图像与真实MRI分布对齐更好,且相比扩散、潜在和小波基基线,在下游脑龄预测和区域解剖一致性方面均有提升。代码已开源:https://github.com/sisinflab/WaveDiT
原文摘要 · Abstract (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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。