WaveDiT: Emparejamiento de flujo wavelet consciente de la distribución para la síntesis eficiente de MRI cerebral 3D

Resumen

Los conjuntos de datos grandes y demográficamente equilibrados son esenciales para obtener biomarcadores fiables de neuroimagen. La síntesis 3D de resonancia magnética cerebral de resolución completa puede respaldar el aumento de datos en este contexto, pero los enfoques existentes incurren en costos computacionales prohibitivos a escala volumétrica o dependen de una compresión latente con pérdida que puede comprometer el detalle anatómico. Como resultado, el aumento generativo 3D práctico a menudo requiere infraestructura de cómputo especializada. Proponemos WaveDiT, un marco de coincidencia de flujo condicional que opera en el espacio de coeficientes de una Transformada Wavelet Discreta 3D de Haar. El modelo combina atención factorizada espacio-profundidad con modelado de incertidumbre heterocedástica por bandas derivado de estadísticas wavelet de orden superior. La log-varianza predicha se integra directamente tanto en el objetivo de flujo como en la vía de acondicionamiento, lo que permite una precisión adaptativa coherente con la estructura de varianza de cola pesada y dependiente de la entrada del detalle anatómico. Esta formulación admite la síntesis 3D de resolución completa bajo restricciones prácticas de memoria y tiempo en una sola GPU moderna. La evaluación en una cohorte multisitio demuestra una mejor alineación entre las distribuciones de MRI generadas y reales, junto con una mejora en la predicción descendente de la edad cerebral y en la concordancia anatómica a nivel regional en comparación con las líneas base basadas en difusión, latente y wavelet. El código está disponible en 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