WaveDiT: 分布認識型ウェーブレットフローマッチングによる効率的な3D脳MRI合成

要旨

大規模で人口統計的にバランスの取れたデータセットは、信頼性の高い神経画像バイオマーカーにとって不可欠である。完全解像度の3D脳MRI合成は、この設定におけるデータ拡張を支援できるが、既存のアプローチは、体積スケールで法外な計算コストを伴うか、解剖学的詳細を損なう可能性のある損失のある潜在圧縮に依存している。その結果、実用的な3D生成的拡張には、しばしば特殊な計算インフラストラクチャが必要となる。我々は、3D Haar離散ウェーブレット変換の係数空間で動作する条件付きフローマッチングフレームワークであるWaveDiTを提案する。本モデルは、因子分解された空間-深度注意と、高次ウェーブレット統計から導出されたバンド別の不均一分散不確実性モデリングを組み合わせる。予測された対数分散は、フロー目的関数と条件付け経路の両方に直接統合され、解剖学的詳細の重い裾を持ち入力に依存する分散構造と一致した適応精度を可能にする。この定式化により、単一の最新GPU上で実用的なメモリと時間の制約の下で、完全解像度の3D合成が可能となる。マルチサイトコホートでの評価は、生成されたMRI分布と実際の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