[논문 리뷰] Neuromorphologicaly-preserving Volumetric data encoding using VQ-VAE

Petru-Daniel Tudosiu, Thomas Varsavsky외 6명|arXiv (Cornell University)|2020. 02. 13.

Medical Image Segmentation Techniques참고 문헌 33인용 수 12

한 줄 요약

이 논문은 전체 해상도의 뇌 MRI 볼륨을 원본 크기의 0.825%로 압축하면서도 영상 품질과 신경형태학적 구조를 유지하는 3D VQ-VAE 아키텍처를 제안한다. SubPixel Deconvolutions, FixUp 블록, 그리고 3D 적응형 손실을 통합함으로써, 모델은 최신 기술 수준의 재구성 품질을 달성하고 건강한 뇌와 알츠하이머병 영향을 받은 뇌 모두에서 뛰어난 형태학적 유지 능력을 보여준다.

ABSTRACT

The increasing efficiency and compactness of deep learning architectures, together with hardware improvements, have enabled the complex and high-dimensional modelling of medical volumetric data at higher resolutions. Recently, Vector-Quantised Variational Autoencoders (VQ-VAE) have been proposed as an efficient generative unsupervised learning approach that can encode images to a small percentage of their initial size, while preserving their decoded fidelity. Here, we show a VQ-VAE inspired network can efficiently encode a full-resolution 3D brain volume, compressing the data to $0.825\%$ of the original size while maintaining image fidelity, and significantly outperforming the previous state-of-the-art. We then demonstrate that VQ-VAE decoded images preserve the morphological characteristics of the original data through voxel-based morphology and segmentation experiments. Lastly, we show that such models can be pre-trained and then fine-tuned on different datasets without the introduction of bias.

연구 동기 및 목표

• 고차원적이고 데이터 볼륨이 적은 3D 의료 영상 문제를 해결하기 위해 효율적이고 비지도 기반의 압축을 가능하게 한다.
• GAN(모드 붕괴, 학습 불안정성)과 VAE(저품질 재구성)의 한계를 극복하여 3D 의료 데이터에서의 성능을 향상시킨다.
• 압축된 재구성 영상이 임상 및 연구 응용에 필수적인 신경해부학적 형태학적 특성을 유지하도록 보장한다.
• 한 데이터셋에서 사전학습하고 다른 데이터셋에서 미세조정함으로써 이식 가능성을 확보한다. 이 과정에서 편향이 유입되지 않도록 한다.
• 생성 모델이 부피형 신경영상에 대해 작고 생물학적으로 타당한 형태로 구현될 수 있음을 입증한다.

제안 방법

• VAE의 가우시안 사후분포를 크기 K×D의 학습된 벡터 양자화 임베딩 공간으로 대체하여 3D MRI에 VQ-VAE 프레임워크를 적응시킨다.
• 3D 부피 재구성에서 격자 무늬 아티팩트를 완화하기 위해 SubPixel Deconvolutions를 사용한다.
• 소규모 배치 조건에서도 학습 안정성을 확보하기 위해 배치 정규화 없이 FixUp 잔여 블록을 활용한다.
• Barron(2019)의 적응형 손실을 3D로 확장하고, Baur 등(2019)의 영감을 받은 손실과 비교하여 시각적 품질을 향상시킨다.
• 재구성 손실과 인지 정규화를 조합하여 엔드 투 엔드로 학습함으로써 구조적 정밀도를 향상시킨다.
• 이산 임베딩 코드에서의 생성 품질 향상을 위해 PixelCNN을 사용하여 순차적 복원을 수행한다.

실험 결과

연구 질문

• RQ1 3D VQ-VAE는 전체 해상도의 뇌 MRI를 고품질, 저압축 비율로 재구성하면서도 신경해부학적 구조를 유지할 수 있는가?
• RQ2 제안된 VQ-VAE 모델은 최신 기술 수준의 GAN 기반 방법에 비해 영상 품질과 형태학적 유지 능력에서 어떻게 비교되는가?
• RQ3 VQ-VAE 재구성 영상이 회색질, 흰질, 뇌척수액의 세분화 정확도를 어느 정도 유지하는가?
• RQ4 압축 후 알츠하이머병 환자군과 건강한 대조군 간의 인구 수준의 형태학적 차이가 유지되는가?
• RQ5 모델은 새로운 데이터셋에서 효과적으로 미세조정될 수 있는가? 이 과정에서 편향이나 성능 저하가 발생하지는 않는가?

주요 결과

• 제안된 VQ-VAE 모델은 3D 뇌 MRI 볼륨을 원본 크기의 0.825%로 압축하면서도 높은 재구성 품질을 유지한다.
• α-WGAN 기준선에 비해 MS-SSIM에서 유의미한 향상과 더 낮은 MMD를 기록했다(p < 0.01, Wilcoxon 부호순위 검정).
• 재구성 영상에서 GM, WM, CSF 세분화의 Dice 스코어가 α-WGAN 재구성보다 유의미하게 높았다.
• VBM 분석 결과, 제안된 방법은 원본 데이터와 재구성 데이터 간의 형태학적 차이를 상당히 줄였으며, 특히 알츠하이머병 코hort에서 두드러졌다.
• 적응형 손실을 사용한 VQ-VAE가 Baur 손실 버전보다 VBM T-검정 지도에서 더 우수한 성능을 보여, 인구 수준의 형태학적 패턴 유지 능력이 뛰어났다.
• 모델은 이식 가능성을 입증했으며, 사전학습된 가중치를 새로운 데이터셋에서 미세조정함으로써 편향이나 성능 저하 없이 효과적으로 활용될 수 있었다.