WaveDiT: Verteilungsbewusstes Wavelet-Flow-Matching für effiziente 3D-Gehirn-MRT-Synthese

Zusammenfassung

Große und demografisch ausgewogene Datensätze sind für zuverlässige Neuroimaging-Biomarker unerlässlich. Die Synthese von 3D-Gehirn-MRT in voller Auflösung kann in diesem Kontext die Datenaugmentation unterstützen, aber bestehende Ansätze verursachen entweder prohibitive Rechenkosten im volumetrischen Maßstab oder verlassen sich auf verlustbehaftete latente Kompression, die anatomische Details beeinträchtigen kann. Infolgedessen erfordert die praktische 3D-generative Augmentation oft spezialisierte Recheninfrastruktur. Wir schlagen WaveDiT vor, ein Framework für bedingtes Flow Matching, das im Koeffizientenraum einer 3D-Haar-Diskreten-Wavelet-Transformation arbeitet. Das Modell kombiniert faktorisierte räumlich-Tiefen-Aufmerksamkeit mit bandweiser heteroskedastischer Unsicherheitsmodellierung, die aus Wavelet-Statistiken höherer Ordnung abgeleitet wird. Die vorhergesagte Log-Varianz wird direkt in das Flow-Ziel und den Konditionierungspfad integriert, was eine adaptive Präzision ermöglicht, die mit der schwer-tailed und eingabeabhängigen Varianzstruktur anatomischer Details konsistent ist. Diese Formulierung unterstützt die Synthese von 3D-Bildern in voller Auflösung unter praktischen Speicher- und Zeitbeschränkungen auf einer einzelnen modernen GPU. Die Evaluierung an einer multizentrischen Kohorte zeigt eine verbesserte Übereinstimmung zwischen generierten und echten MRT-Verteilungen sowie eine verbesserte nachgelagerte Vorhersage des Gehirnalters und eine verbesserte regionale anatomische Übereinstimmung im Vergleich zu Diffusions-, latenten und waveletbasierten Baselines. Der Code ist verfügbar unter 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