WaveDiT : Appariement de flots par ondelettes conscient de la distribution pour la synthèse efficace d'IRM cérébrale 3D

Résumé

Des ensembles de données volumineux et équilibrés sur le plan démographique sont essentiels pour obtenir des biomarqueurs fiables en neuroimagerie. La synthèse d'IRM cérébrale 3D en pleine résolution peut soutenir l'augmentation de données dans ce contexte, mais les approches existantes entraînent soit un coût de calcul prohibitif à l'échelle volumétrique, soit reposent sur une compression latente avec pertes qui peut compromettre les détails anatomiques. En conséquence, l'augmentation générative 3D pratique nécessite souvent une infrastructure de calcul spécialisée. Nous proposons WaveDiT, un cadre de matching de flux conditionnel opérant dans l'espace des coefficients d'une transformée en ondelettes discrètes de Haar 3D. Le modèle combine une attention spatio-profondeur factorisée avec une modélisation de l'incertitude hétéroscédastique par bande, dérivée de statistiques d'ondelettes d'ordre supérieur. La log-variance prédite est intégrée à la fois dans l'objectif du flux et dans la voie de conditionnement, permettant une précision adaptative cohérente avec la structure de variance à queue lourde et dépendante de l'entrée des détails anatomiques. Cette formulation permet une synthèse 3D en pleine résolution sous des contraintes pratiques de mémoire et de temps sur une seule GPU moderne. L'évaluation sur une cohorte multi-site démontre un meilleur alignement entre les distributions d'IRM générées et réelles, ainsi qu'une amélioration de la prédiction en aval de l'âge cérébral et de l'accord anatomique au niveau régional par rapport aux références basées sur la diffusion, la latence et les ondelettes. Le code est disponible à l'adresse 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