WaveDiT: Correspondência de Fluxo Wavelet Consciente da Distribuição para Síntese Eficiente de RM Cerebral 3D

Resumo

Conjuntos de dados grandes e demograficamente equilibrados são essenciais para biomarcadores confiáveis em neuroimagem. A síntese de ressonância magnética cerebral 3D em resolução total pode apoiar a aumentação de dados nesse contexto, mas as abordagens existentes ou incorrem em custo computacional proibitivo em escala volumétrica ou dependem de compressão latente com perdas que pode comprometer detalhes anatômicos. Como resultado, a aumentação generativa 3D prática frequentemente requer infraestrutura computacional especializada. Propomos o WaveDiT, uma estrutura de correspondência de fluxo condicional que opera no espaço de coeficientes de uma Transformada Wavelet Discreta de Haar 3D. O modelo combina atenção fatorada espaço-profundidade com modelagem de incerteza heterocedástica por banda derivada de estatísticas wavelet de ordem superior. A log-variância prevista é integrada diretamente tanto no objetivo de fluxo quanto no caminho de condicionamento, permitindo precisão adaptativa consistente com a estrutura de variância de cauda pesada e dependente da entrada dos detalhes anatômicos. Essa formulação suporta síntese 3D em resolução total sob restrições práticas de memória e tempo em uma única GPU moderna. A avaliação em uma coorte multissítio demonstra melhor alinhamento entre as distribuições de ressonância magnética geradas e reais, juntamente com previsão aprimorada de idade cerebral e concordância anatômica em nível regional em comparação com abordagens baseadas em difusão, latente e wavelet. O código está disponível em 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