WaveDiT：分布感知的小波流匹配用於高效3D腦部MRI合成

摘要

大型且人口統計學上均衡的數據集對於可靠的神經影像生物標記至關重要。全解析度3D腦部MRI合成可在這一背景下支援資料擴增，但現有方法在體積尺度上要么面臨高昂的計算成本，要么依賴可能損害解剖細節的有損潛在壓縮。因此，實用的3D生成式擴增通常需要專門的計算基礎設施。我們提出WaveDiT，這是一個在3D哈爾離散小波變換係數空間中運作的條件流匹配框架。該模型將分解式空間-深度注意力機制與源自高階小波特性的頻帶異方差不確定性建模相結合。預測的對數變異數直接整合到流目標與條件路徑中，從而實現與解剖細節的重尾及輸入依賴變異數結構一致的自適應精度。此公式支援在單一現代GPU上，於實際記憶體與時間限制下進行全解析度3D合成。在多站點隊列上的評估顯示，與擴散、潛在及小波基準方法相比，生成與真實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