[논문 리뷰] Self-Supervised Out-of-Distribution Detection in Brain CT Scans

Abinav Ravi Venkatakrishnan, Seong Tae Kim외 3명|arXiv (Cornell University)|2020. 11. 10.

Anomaly Detection Techniques and Applications참고 문헌 15인용 수 7

한 줄 요약

이 논문은 재구성 기반 오토인코딩과 기하학적 변환 예측을 조합하여 분포 외 탐지 성능을 햖थ한 뇌 CT 스캔을 위한 자기지도 학습 이상 탐지 프레임워크를 제안한다. 문맥 복원으로 사전 훈련하고 재구성 및 공간적 변환 예측을 동시에 최적화함으로써 임상 데이터에서 최신 기준인 AUROC (0.822)와 AUPR (0.868)를 달성하여 이전의 비지도 방법들을 능가한다.

ABSTRACT

Medical imaging data suffers from the limited availability of annotation because annotating 3D medical data is a time-consuming and expensive task. Moreover, even if the annotation is available, supervised learning-based approaches suffer highly imbalanced data. Most of the scans during the screening are from normal subjects, but there are also large variations in abnormal cases. To address these issues, recently, unsupervised deep anomaly detection methods that train the model on large-sized normal scans and detect abnormal scans by calculating reconstruction error have been reported. In this paper, we propose a novel self-supervised learning technique for anomaly detection. Our architecture largely consists of two parts: 1) Reconstruction and 2) predicting geometric transformations. By training the network to predict geometric transformations, the model could learn better image features and distribution of normal scans. In the test time, the geometric transformation predictor can assign the anomaly score by calculating the error between geometric transformation and prediction. Moreover, we further use self-supervised learning with context restoration for pretraining our model. By comparative experiments on clinical brain CT scans, the effectiveness of the proposed method has been verified.

연구 동기 및 목표

• 임상 스캔에서 정상 케이스가 비정상 케이스에 비해 현저히 많고, 레이블이 부족한 3D 의료 영상 데이터 문제를 해결한다.
• 수동 레이블링에 의존하지 않고 데이터의 내재적 구조를 활용하여 비용이 많이 드는 수동 레이블링을 피하는 자기지도 학습 이상 탐지 방법을 개발한다.
• 세부적인 특징과 맥락적 특징을 모두 포착하기 위해 局부 재구성 오차와 전역 기하학적 변환 예측을 조합하여 탐지 성능을 향상시킨다.
• 다양한 비정상 케이스를 포함한 실제 임상 뇌 CT 스캔에서 제안된 방법의 효과성을 검증한다.

제안 방법

• 건강한 스캔에서 일반화 가능한 특징을 학습하기 위해 입력 영상의 패치를 교환하고, 원본을 재구성하도록 VAE를 훈련시켜 문맥 복원을 통한 사전 훈련을 수행한다.
• 두 개의 헤드를 가진 다중 작업 학습 프레임워크 도입: 하나는 이미지 재구성(VAE)을 위한 것이고, 다른 하나는 적용된 기하학적 변환을 예측하기 위한 것이다.
• 기하학적 변환 헤드를 교차 엔트로피 손실을 사용하여 훈련시키며, 가능한 변환들(0°, 90°, 180°, 270°, ±1/8번 이동) 중에서 적용된 변환을 예측한다.
• 정규화된 가중 평균을 통해 두 헤드의 이상도 점수를 조합하며, 하이퍼파rameter λ = 0.5를 사용한다. 여기서 재구성 점수는 α로 스케일링되고, 기하학적 예측 점수는 가능한 변환들에 대해 평균화된다.
• 두 작업 모두에 공통 인코더를 사용하여 재구성과 변환 예측에 모두 유익한 특징 학습을 가능하게 한다.
• 테스트 시점에 재구성 오차와 변환 예측 오차를 기반으로 이상도 점수를 할당하며, 오차가 클수록 분포 외일 가능성이 높아진다.

실험 결과

연구 질문

• RQ1 문맥 복원을 통한 자기지도 사전 훈련이 뇌 CT 스캔에서 이상 탐지의 특징 학습에 기여하는가?
• RQ2 재구성과 기하학적 변환 예측을 동시에 훈련하는 것이 단일 작업 접근법보다 탐지 성능을 향상시키는가?
• RQ3 국소적(재구성) 및 전역적(변환 예측) 이상 신호의 통합이 탐지 정확도에 어떤 영향을 미치는가?
• RQ4 제안된 다중 작업 프레임워크에서 분류 성능와 세그멘테이션 국소화 간의 상충 관계는 어떠한가?

주요 결과

• 제안된 다중 작업 프레임워크는 AUROC 0.822와 AUPR 0.868를 달성하여 기존 방법들인 VAE (AUROC: 0.668)와 Context-encoding VAE (AUROC: 0.766)를 뛰어넘었다.
• 문맥 복원으로 사전 훈련한 결과 표준 VAE 대비 AUROC가 0.668에서 0.673으로 향상되었고, AUPR는 0.704에서 0.772로 상승했다.
• 사전 훈련 후 재구성 손실만을 사용하는 모델(VAE with Pretraining)은 기준 VAE보다 성능 향상을 보였으며, 자기지도 사전 훈련의 유용성을 확인했다.
• 분류 성능 향상에도 불구하고 세그멘테이션 성능(DSC: 0.086 ± 0.024)은 다른 방법들보다 약간 열 劣하여 탐지와 국소화 간의 상충 관계가 있음을 시사했다.
• 기하학적 변환 헤드는 재구성 오차와 보완되는 강력한 전역 이상 신호를 제공하여 이상 탐지에 기여했다.
• 비록 정상 스캔만으로 훈련되었지만, 뇌 위축, intracranial bleeding, 허혈성 병변, 종양 등 다양한 비정상 케이스에 대해 뛰어난 일반화 성능를 보였다.