WaveDiT: 효율적인 3D 뇌 MRI 합성을 위한 분포 인식 웨이블릿 흐름 매칭

초록

크고 인구통계학적으로 균형 잡힌 데이터셋은 신뢰할 수 있는 신경영상 바이오마커에 필수적이다. 전체 해상도의 3D 뇌 MRI 합성은 이러한 환경에서 데이터 증강을 지원할 수 있지만, 기존 접근법은 체적 규모에서 엄청난 계산 비용을 초래하거나 해부학적 세부 구조를 손상시킬 수 있는 손실 잠재 압축에 의존한다. 결과적으로 실용적인 3D 생성적 증강은 종종 특수한 계산 인프라를 필요로 한다. 본 연구는 3D 하르 이산 웨이블릿 변환의 계수 공간에서 작동하는 조건부 흐름 매칭 프레임워크인 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