[논문 리뷰] Generation of 3D Brain MRI Using Auto-Encoding Generative Adversarial Networks

Gihyun Kwon, Chihye Han외 1명|arXiv (Cornell University)|2019. 08. 07.

Generative Adversarial Networks and Image Synthesis참고 문헌 18인용 수 11

한 줄 요약

이 논문은 랜덤 노이즈에서 고해상도이고 다양한 3D 뇌 MRI 볼륨을 생성하기 위해 VAE 스타일의 잠재공간 정규화와 WGAN-GP 학습을 결합한 3D 오토인코딩 GAN(3D-α-WGAN-GP)을 제안한다. 이 모델은 분포 적합도와 샘플 다양성 측면에서 베이스라인 GAN 및 VAE-GAN을 능가하며, 최소한의 훈련 데이터로도 다양한 모odalities에서 정상 및 병변이 있는 뇌 MRI를 현실적으로 합성할 수 있다.

ABSTRACT

As deep learning is showing unprecedented success in medical image analysis tasks, the lack of sufficient medical data is emerging as a critical problem. While recent attempts to solve the limited data problem using Generative Adversarial Networks (GAN) have been successful in generating realistic images with diversity, most of them are based on image-to-image translation and thus require extensive datasets from different domains. Here, we propose a novel model that can successfully generate 3D brain MRI data from random vectors by learning the data distribution. Our 3D GAN model solves both image blurriness and mode collapse problems by leveraging alpha-GAN that combines the advantages of Variational Auto-Encoder (VAE) and GAN with an additional code discriminator network. We also use the Wasserstein GAN with Gradient Penalty (WGAN-GP) loss to lower the training instability. To demonstrate the effectiveness of our model, we generate new images of normal brain MRI and show that our model outperforms baseline models in both quantitative and qualitative measurements. We also train the model to synthesize brain disorder MRI data to demonstrate the wide applicability of our model. Our results suggest that the proposed model can successfully generate various types and modalities of 3D whole brain volumes from a small set of training data.

연구 동기 및 목표

• 딥 러닝 기반 의료 영상 분석에서 3D 의료 영상 데이터의 부족이라는 핵심 과제를 해결하기 위해.
• 하이브리드 VAE-GAN 프레임워크를 통해 3D GAN 기반의 MRI 생성에서 모드 붕괴와 이미지 흐림 현상을 극복하기 위해.
• 쌍체 또는 교차 도메인 데이터가 필요 없이 랜덤 잠재 벡터에서 다양하고 현실적인 3D 뇌 MRI 볼륨을 생성하기 위해.
• 다양한 영상 시퀀스(T2, FLAIR, T1w)에서 정상, 종양 유발, 뇌졸중 영향을 받은 뇌 MRI를 합성함으로써 모델의 유연성을 입증하기 위해.

제안 방법

• VAE의 KL 발산 대신 GAN 스타일의 JS 발산을 사용하여 인코더의 사후분포와 사전분포 간의 적대적 매칭을 수행하는 코드 디스크림이너 네트워크를 통합함으로써, 3D로 확장된 α-GAN 아키텍처를 구현한다.
• 4×4×4 커널을 사용하고 배치 정규화(BatchNorm)와 LeakyReLU 활성화 함수를 적용한 5개의 3D 컨볼루션 계층을 갖춘 3D 인코더와 3D 생성기로 구성되어 있으며, 부피 영상 데이터의 공간 계층을 모델링한다.
• 학습 안정성을 향상시키고 모드 커버리지를 개선하기 위해 기울기 페널티를 포함한 WGAN-GP 손실을 사용하여 3D GAN에서 흔한 학습 불안정성 문제를 완화한다.
• 생성기가 인코딩된 잠재 벡터에서 재구성된 이미지를 생성함으로써, 분리 가능하고 의미 있는 잠재 표현을 가능하게 하는 변동형 오토인코더 구조를 활용한다.
• 재구성 손실, 이미지 디스크림이너에서의 적대적 손실, 코드 디스크림이너에서의 적대적 손실을 동시에 최적화하는 통합 목표 함수를 사용하여 모델을 엔드 투 엔드로 훈련시킨다.
• 다양성과 분포 정확도의 균형을 맞추기 위해 잠재 벡터 크기 조정을 시행하였으며, 중간 크기의 벡터(예: 512)가 최적의 성능을 낼 수 있음을 발견했다.

실험 결과

연구 질문

• RQ1 쌍체 또는 교차 도메인 데이터가 필요 없이 랜덤 노이즈에서 현실적이고 다양한 3D 뇌 MRI 볼륨을 생성할 수 있는가?
• RQ2 VAE 스타일의 잠재 정규화를 GAN 학습과 융합함으로써 3D MRI 생성에서 모드 붕괴와 이미지 흐림 현상을 효과적으로 완화할 수 있는가?
• RQ3 기존 GAN과 비교해 WGAN-GP 손실이 3D MRI 생성에서 학습 안정성과 분포 적합도를 향상시키는가?
• RQ4 제안된 모델은 종양, 뇌졸중 병변 등의 다양한 뇌 병변과 T2, FLAIR, T1w 등의 다양한 영상 모odalities에 일반화 가능한가?
• RQ5 3D MRI 생성에서 샘플 다양성과 분포 정확도를 균형 잡는 데 최적의 잠재 벡터 크기는 얼마인가?

주요 결과

• 제안된 3D-α-WGAN-GP 모델은 실제 데이터에서 가장 높은 Fréchet MRI 거리(FID) 점수 12.3을 기록하여, 3D-WGAN-GP(FID: 21.5)와 3D-VAE-GAN(FID: 18.7)을 크게 능가하며, 분포 매칭 능력이 뛰어나다는 것을 보여준다.
• 모델은 1000개의 쌍에 대해 MS-SSIM 점수 0.92를 기록하여 실제 데이터(MS-SSIM: 0.93)와 밀도 높은 유사성을 보였으며, 베이스라인 모델들은 생성된 샘플 간 유사성이 높아 모드 붕괴를 시사한다.
• PCA 시각화 결과, 제안된 모델에서 생성된 샘플은 실제 데이터와 가장 넓고 현실적인 분포 범위를 보이며, 3D-WGAN-GP와 3D-VAE-GAN이 제한된 또는 흐린 출력을 생성하는 것과 대비된다.
• 잠재 벡터 크기가 512일 때, 다양성과 분포 정확도 사이의 최적 균형을 달성한다. 더 작은 크기(예: 100)는 심각한 모드 붕괴를 유도하고, 더 큰 크기(예: 1024)는 실제 데이터 분포를 벗어난 샘플을 생성한다.
• 모델은 FLAIR, T2, T1w 등 다양한 모달리티에서 종양과 뇌졸중 병변이 있는 현실적인 3D 뇌 MRI를 성공적으로 생성하여, 다양한 병변과 영상 대비에 대한 일반화 능력을 입증한다.
• 보조 결과에서는 다른 모델(예: 3D-WGAN-GP, 3D-VAE-GAN)이 심각하게 열화되거나 비현실적인 이미지를 생성하는 것으로 확인되어, 제안된 아키텍처의 우수성을 확인한다.