[논문 리뷰] Multiscale Metamorphic VAE for 3D Brain MRI Synthesis

Jaivardhan Kapoor, Jakob H. Macke외 1명|arXiv (Cornell University)|2023. 01. 09.

Generative Adversarial Networks and Image Synthesis인용 수 4

한 줄 요약

이 논문은 고정 해부학적 템플릿에 다중 척도의 미분 가능하고 강도 변환을 적용하여 고해상도 MRI 볼륨을 생성하는 다중 척도 변형 변분 오토인코더(M³AE)를 제안한다. 이 방법은 상태 기반의 FID 점수를 달성하면서도 분리된 해석 가능한 잠재 공간을 유지하여 기존 VAE 및 GAN 기반 접근법보다 분포 커버리지와 해부학적 정확도에서 뚜렷한 성능 향상을 보였다.

ABSTRACT

Generative modeling of 3D brain MRIs presents difficulties in achieving high visual fidelity while ensuring sufficient coverage of the data distribution. In this work, we propose to address this challenge with composable, multiscale morphological transformations in a variational autoencoder (VAE) framework. These transformations are applied to a chosen reference brain image to generate MRI volumes, equipping the model with strong anatomical inductive biases. We structure the VAE latent space in a way such that the model covers the data distribution sufficiently well. We show substantial performance improvements in FID while retaining comparable, or superior, reconstruction quality compared to prior work based on VAEs and generative adversarial networks (GANs).

연구 동기 및 목표

• 고해상도 3D 뇌 MRI를 충분한 데이터 분포 커버리지로 생성하는 데 도전하는 것.
• 더 나은 생성 모델링을 위해 변분 오토인코더 프레임워크에 강력한 해부학적 인덕티브 바이어스를 통합하는 것.
• 재구성 품질을 유지하면서도 후속 분석을 위한 분리된 가로지르는 잠재 공간을 가능하게 하는 것.
• 2D에서 3D MRI 볼륨으로 형태학적 변환 기반 합성을 다중 척도 조합을 통해 확장하는 것.
• GAN의 한계(예: 모드 붕괴, 잠재 공간 부재)와 VAE의 한계(예: 흐릿한 출력)를 극복하는 것

제안 방법

• 모델은 생성의 기초로 고정된 기준 뇌 템플릿을 사용하는 변분 오토인코더를 사용한다.
• 생성 과정을 두 개의 디코더로 분해한다: 하나는 미분 가능 변형 필드(ϕ)를 위한 것이고, 다른 하나는 덧셈 강도 변환(A)을 위한 것이다.
• 변환은 굵기에서 세밀함으로 향하는 다중 척도 캐스케이드 방식으로 적용되며, 각 수준에서 해상도가 두 배로 증가하여 계층적 정밀도 향상이 가능하다.
• 최종 출력은 M(T|ϕ,A) = (T + A) ∘ ϕ로 변환을 조합하여 형성되며, 마지막 단계는 순수하게 강도 기반이다.
• 손실 함수는 ELBO에 변형 필드와 강도 맵에 대한 척도별 정규화 항을 추가로 포함하며, 스무스니스 및 수렴 제약 조건을 포함한다.
• 모델은 AdamW를 사용하여 학습하며 학습률은 0.0003이며, 추론 시 표준 정규 분포에서 샘플링된 잠재 벡터를 사용한다.

Figure 1 : Description of the proposed M 3 AE model. The decoder consists of two separate backbones Decoder ϕ and Decoder A that output diffeomorphic deformation fields and additive intensity transformations, respectively. These transformations are composed from coarse-to-fine scales that increase b

실험 결과

연구 질문

• RQ1 표준 VAE 및 GAN과 비교해 복수 척도의 조합 가능한 형태학적 변환은 3D 뇌 MRI 합성의 정확도와 분포 커버리지 향상에 기여하는가?
• RQ2 변형과 강도 변환을 적용한 고정 해부학적 템플릿을 사용함으로써 생성된 MRI 볼륨의 해부학적 일관성이 향상되는가?
• RQ3 고해상도 3D MRI 생성기에서 VAE 기반 모델이 고해상도 출력을 달성하면서도 분리된 잠재 공간을 유지할 수 있는가?
• RQ4 FID, SSIM, PSNR, MSE 지표에서 제안된 방법은 최고 수준의 기준 모델과 비교해 어떻게 성능을 내는가?
• RQ5 다중 척도 변형 필드로 인해 발생하는 위상 다양성과 시각적 품질의 한계는 무엇인가?

주요 결과

• M³AE 모델은 2D 슬라이스에서 FID 점수 11.2를 기록하여 α WGAN 기준 모델과 β VAE를 크게 앞서며, 더 나은 분포 커버리지와 샘플 품질을 입증했다.
• 해부학적 제약 조건에도 불구하고 재구성 품질은 β VAE와 유사하며, SSIM, PSNR, MSE 값이 좁은 범위 내에 유지되었다.
• 정성적 분석 결과 M³AE 샘플은 캄버와 구조적 특징이 잘 보존되어 있어 β VAE 및 α WGAN에서 관찰되는 아티팩트와는 대조적으로 가장 해부학적으로 정확했다.
• 잘못 생성된 세밀 척도 변형 필드로 인해 일부 경우에 파도 모양의 시각적 아티팩트가 발생하여 고해상도에서의 정규화 향상 필요성을 시사했다.
• 피질 굴곡의 위상 다양성은 제한적이었으며, 향후 연구에서 주로 굴곡과 피질 굴곡의 변동성을 더 잘 모델링할 필요가 있음을 시사했다.
• 분리된 잠재 공간 덕분에 GAN 기반 방법이 그러한 구조화된 잠재 표현을 갖지 못하는 것과는 달리 의미 있는 후속 분석이 가능했다.

Figure 2 : Preprocessed volumes from the ADNI dataset