[논문 리뷰] Unsupervised anomaly localization using VAE and beta-VAE

Leixin Zhou, Wenxiang Deng외 1명|arXiv (Cornell University)|2020. 05. 19.

COVID-19 diagnosis using AI참고 문헌 21인용 수 13

한 줄 요약

이 논문은 VAE와 베타-VAE를 사용하여 의료 영상에서 비지도 이상 지역 규명을 조사한다. 재구성 오차, ELBO 기울기, KL 기울기, 에너지 기반 투영을 비교한다. 뇌 MRI에서 기울기 기반 예측자인 Rec-grad가 재구성 오차와 에너지 투영보다 뛰어나며, 성능은 VAE의 하이퍼파rameter, 특히 베타-VAE의 베타에 매우 민감하다는 것을 발견한다.

ABSTRACT

Variational Auto-Encoders (VAEs) have shown great potential in the unsupervised learning of data distributions. An VAE trained on normal images is expected to only be able to reconstruct normal images, allowing the localization of anomalous pixels in an image via manipulating information within the VAE ELBO loss. The ELBO consists of KL divergence loss (image-wise) and reconstruction loss (pixel-wise). It is natural and straightforward to use the later as the predictor. However, usually local anomaly added to a normal image can deteriorate the whole reconstructed image, causing segmentation using only naive pixel errors not accurate. Energy based projection was proposed to increase the reconstruction accuracy of normal regions/pixels, which achieved the state-of-the-art localization accuracy on simple natural images. Another possible predictors are ELBO and its components gradients with respect to each pixels. Previous work claimed that KL gradient is a robust predictor. In this paper, we argue that the energy based projection in medical imaging is not as useful as on natural images. Moreover, we observe that the robustness of KL gradient predictor totally depends on the setting of the VAE and dataset. We also explored the effect of the weight of KL loss within beta-VAE and predictor ensemble in anomaly localization.

연구 동기 및 목표

• 의료 영상에서 이상 지역 규명을 위한 다양한 예측자—재구성 오차, ELBO 기울기, KL 기울기, 에너지 기반 투영—의 효과성을 VAE 및 베타-VAE 프레임워크 내에서 평가하기.
• 자연 이미지에서 효과적인 에너지 기반 투영이 뇌 종양 지역 규명과 같은 도전적인 의료 영상 작업으로 일반화되는지 조사하기.
• 특히 KL 가중치(베타)에 따라 베타-VAE 하이퍼파rameter의 이상 지역 규명 성능에 미치는 영향 분석하기.
• 다양한 예측자를 앙상블 학습을 통해 융합하여 개별 방법보다 높은 정확도를 달성하는 데 기여하는지 평가하기.

제안 방법

• 정상 T2 뇌 MRI 영상(BraTS2018)을 기반으로 VAE와 베타-VAE를 학습하여 건강한 조직의 데이터 분포를 학습한다.
• 다양한 이상 점수 계산: 픽셀 단위 재구성 오차, 입력에 대한 ELBO 및 KL 발산 기울기, 에너지 기반 투영 재구성 오차.
• 정상 다각형에 수렴하도록 반복적인 에너지 최소화를 적용하여 정상 영역의 정밀도를 향상시킨다.
• 다양한 예측자(Rec-Error, KL-grad, Rec-grad)를 앙상블 학습하기 위해 로지스틱 회귀를 사용하여 성능 향상.
• 다양한 잠재공간 차원과 베타 값에서 테스트 데이터의 AUROC를 사용해 모든 방법 평가.
• 모델 하이퍼파rameter 민감도 평가를 위해 다양한 VAE 구성 간 성능 비교.

실험 결과

연구 질문

• RQ1 에너지 기반 투영이 표준 재구성 오차에 비해 의료 영상에서 이상 지역 규명에 유의미하게 향상되는가?
• RQ2 KL 기울기 예측자의 강건성은 다양한 VAE 아키텍처와 하이퍼파rameter, 특히 베타-VAE의 베타에 따라 어떻게 달라지는가?
• RQ3 재구성 오차, ELBO 기울기, KL 기울기, Rec-grad 중 어느 이상 감지 예측자가 뇌 MRI 이상 지역 규명에서 가장 높은 AUROC를 달성하는가?
• RQ4 다양한 예측자를 앙상블 학습을 통해 융합했을 때 상호 보완성이 얼마나 높은가?
• RQ5 베타-VAE의 잠재공간 차원과 베타 값은 다양한 이상 지역 규명 방법의 성능에 어떻게 영향을 미치는가?

주요 결과

• 에너지 기반 투영은 정상 영역의 재구성 정확도를 향상시켰지만 기울기 기반 방법보다 떨어졌으며, AUROC 0.861을 기록한 반면 Rec-grad는 0.900을 달성했다.
• KL 기울기 예측자는 일관되게 강건하지 않았으며, VAE 하이퍼파rameter, 특히 베타와 잠재공간 크기에 매우 민감했다.
• 베타-VAE의 경우, 베타를 높일수록 대부분의 예측자에서 성능 향상이 있었으며, 베타=10에서 최고 성능을 기록했지만 Combi 성능은 높은 베타에서 저하되었다.
• Rec-grad 예측자가 가장 높은 AUROC 0.905를 기록했으며, 재구성 오차(0.856)와 에너지 기반 투영(0.861)을 모두 앞섰다.
• Rec-Error, KL-grad, Rec-grad를 사용한 앙상블 학습은 AUROC 0.903을 기록했으며, Rec-grad 단독 성능보다 略적으로 낮아 예측자 간 보완성이 제한적임을 시사했다.
• 표준 VAE에서는 정상 영역의 재구성 품질이 떨어졌으며, 특히 고베타 설정에서 두드러졌다. 이는 모델이 정상 다각형을 잘 학습하지 못했음을 의미하며, 투영 방법의 이점이 제한됨을 시사한다.