[논문 리뷰] High- and Low-level image component decomposition using VAEs for improved reconstruction and anomaly detection

David Zimmerer, Jens Petersen외 1명|arXiv (Cornell University)|2019. 11. 27.

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

한 줄 요약

이 논문은 고해상도(粗) 및 저해상도(細) 성분으로 이미지를 분해하는 전용 복원 브랜치를 갖춘 새로운 변분 오토에인드어 아키텍처인 주된 구성 요소 조건부 계층적 VAE(pchVAE)를 제안한다. 수정된 손실 함수를 통해 분리성을 강제함으로써 pchVAE는 표준 VAE보다 더 선명한 복원을 달성하면서도, 의료 영상 생성에서 흔히 발생하는 과도한 부드러움 문제를 해결하며, 이상 탐지 성능을 유지하거나 略로 향상시킨다.

ABSTRACT

Variational Auto-Encoders have often been used for unsupervised pretraining, feature extraction and out-of-distribution and anomaly detection in the medical field. However, VAEs often lack the ability to produce sharp images and learn high-level features. We propose to alleviate these issues by adding a new branch to conditional hierarchical VAEs. This enforces a division between higher-level and lower-level features. Despite the additional computational overhead compared to a normal VAE it results in sharper and better reconstructions and can capture the data distribution similarly well (indicated by a similar or slightly better OoD detection performance).

연구 동기 및 목표

• 표준 VAE를 의료 영상 복원에 적용할 때 발생하는 과도한 부드러움 문제를 해결하기 위해.
• 비지도 표현 학습에서 고해상도(粗) 및 저해상도(細) 이미지 특징 간의 분리성을 향상시키기 위해.
• 아키텍처 수정에도 불구하고 이상 분포(out-of-distribution, OoD) 탐지 성능을 유지하거나 향상시키기 위해.
• 명시적인 구조 계층 모델링을 통해 뇌 MRI 데이터의 복원 정밀도를 향상시키기 위해.
• 의료 영상 분석을 위한 더 해석 가능하고 모듈화된 VAE 프레임워크를 제공하기 위해.

제안 방법

• 고해상도 및 저해상도 성분에 대해 별도로 복원 오차를 최소화하고, 상호정보 정규화 항을 추가한 새로운 손실 함수를 제안한다.
• 저해상도 디코더가 고해상도 잠재 코드와 저해상도 코드 양쪽에 조건부로 설정된 조건부 계층적 VAE 아키텍처를 도입하여 분리성을 촉진한다.
• 고해상도 및 저해상도 특징을 위한 공통 인코더 백본을 사용하고, 각각의 별도 디코더를 통해 군집 및 세밀한 구조 복원을 수행한다.
• 세 가지 복원 항(고해상도, 저해상도, 교차항)을 포함한 수정된 변분 추론 목표 함수를 적용하여 성분 간 독립성을 강제한다.
• 재파arameterization 및 몬테카를로 샘플링을 활용하여 계층 모델의 엔드 투 엔드 학습을 가능하게 한다.
• 각 잠재 공간에 대해 KL 발산 항을 도입하여 사후 분포가 표준 정규분포로 정규화되도록 한다.

실험 결과

연구 질문

• RQ1 전용 저해상도 복원 브랜치를 갖춘 계층적 VAE 아키텍처가 표준 VAE보다 이미지 복원의 선명함을 향상시킬 수 있는가?
• RQ2 고해상도 및 저해상도 특징 간의 분리성을 강제하면 의료 영상에서 세밀한 무늬를 더 잘 복원할 수 있는가?
• RQ3 제안된 pchVAE는 기준 VAE보다 이상 분포(OoD) 탐지 성능을 유지하거나 향상시키는가?
• RQ4 모델은 뇌 MRI 데이터에서 군집적 구조 성분과 세밀한 무늬 세부 정보를 효과적으로 분리할 수 있는가?
• RQ5 상호정보 정규화 항이 고해상도 및 저해상도 성분 간 독립성을 촉진하는 데 효과적인가?

주요 결과

• pchVAE는 표준 VAE 및 VQVAE-2보다 뇌 MRI 슬라이스에서 세부 정보와 무늬를 더 잘 유지하는 등, 특히 세밀한 디테일을 복원하는 데 있어 훨씬 더 선명한 이미지 복원을 달성한다.
• AUROC 및 평균 정밀도로 측정된 이상 탐지 성능을 유지하거나 略로 향상시켜 데이터 분포를 강건하게 모델링하고 있음을 시사한다.
• High VAE 및 Low VAE 변종보다 복원 품질에서 뛰어난 성능을 보이며, 양 성분의 공동 모델링의 이점을 입증한다.
• 상호정보 정규화 항(λ₃)의 포함으로 더 분리된 표현이 생성되었으며, 저해상도 성분이 고해상도 특징의 영향을 덜 받는다.
• 높은 계산 비용에도 불구하고, HCP 및 BraTS2017 데이터셋에서 일관된 성능을 보이며 일반화 능력에 손상 없이 더 높은 복원 정밀도를 달성한다.
• 제거 실험을 통해 세 항 복원 손실이 시각적 품질 향상과 분리성 향상 달성에 필수적임을 확인한다.