[논문 리뷰] Unsupervised Brain Anomaly Detection and Segmentation with Transformers

Walter Hugo Lopez Pinaya, Petru-Daniel Tudosiu외 5명|arXiv (Cornell University)|2021. 02. 23.

Domain Adaptation and Few-Shot Learning참고 문헌 26인용 수 31

한 줄 요약

이 논문은 VQ-VAE를 활용하여 이산 잠재 코드를 얻고 자동회귀 Transformer(Performer) 앙상블을 통해 건강한 뇌 분포를 모델링함으로써 이상치를 내부 채우기(in-painting)하고 잔여 기반 분할을 가능하게 하는 비지도 뇌 MRI 이상 탐지 및 분할 방법을 제시하며 합성 데이터와 실제 데이터셋에서 최첨단 결과를 달성한다.

ABSTRACT

Pathological brain appearances may be so heterogeneous as to be intelligible only as anomalies, defined by their deviation from normality rather than any specific pathological characteristic. Amongst the hardest tasks in medical imaging, detecting such anomalies requires models of the normal brain that combine compactness with the expressivity of the complex, long-range interactions that characterise its structural organisation. These are requirements transformers have arguably greater potential to satisfy than other current candidate architectures, but their application has been inhibited by their demands on data and computational resource. Here we combine the latent representation of vector quantised variational autoencoders with an ensemble of autoregressive transformers to enable unsupervised anomaly detection and segmentation defined by deviation from healthy brain imaging data, achievable at low computational cost, within relative modest data regimes. We compare our method to current state-of-the-art approaches across a series of experiments involving synthetic and real pathological lesions. On real lesions, we train our models on 15,000 radiologically normal participants from UK Biobank, and evaluate performance on four different brain MR datasets with small vessel disease, demyelinating lesions, and tumours. We demonstrate superior anomaly detection performance both image-wise and pixel-wise, achievable without post-processing. These results draw attention to the potential of transformers in this most challenging of imaging tasks.

연구 동기 및 목표

• 건강한 뇌 분포를 학습하여 뇌 이상을 탐지하고 분할하는 완전한 비지도 프레임워크를 개발한다.
• VQ-VAE를 통한 이산 잠재 공간을 활용하여 뇌 이미지의 트랜스포머 모델링을 효율적으로 가능하게 한다.
• 여러 잠재 공간 순서를 사용하는 트랜스포머 앙상블로 이상 분할 및 탐지를 향상시킨다.

제안 방법

• VQ-VAE를 통해 뇌 MRI를 이산 잠재 공간으로 압축하여 z_q를 얻고 2D에서 시퀀스에 이르는 간결한 표현을 확보한다.
• 잠재 코드를 학습시키기 위해 자기회귀 Transformer(Performer) 앙상블을 학습시켜 p(s_i|s_<i)를 학습하고 건강한 데이터의 분포를 모델링한다.
• 저확률 잠재 코드를 임계값으로 식별하고 트랜스포머로 샘플링된 값으로 채워 넣어 이미지를 치유한 후 이상치를 식별한다.
• 원본 이미지와 치유된 이미지 간의 픽셀 단위 잔여값 |x - x_hat'|를 계산하고 임계값으로 분할하여 이상 분할을 얻는다.
• 잠재 공간 이상 위치에서 도출된 재샘플링 마스크를 이용해 잔여값에 공간 정보를 주입하고 이미지 해상도로 업스케일한다.
• 8개의 서로 다른 잠재 공간 순서를 사용하는 트랜스포머 앙상블로 견고성을 강화하고 평균으로 잔여값을 집계한다.

실험 결과

연구 질문

• RQ1 비지도 트랜스포머가 VQ-VAE로 학습된 이산 잠재 공간에 적용될 때, 라벨이 없는 병리 데이터 없이도 뇌 이상을 정확하게 탐지하고 분할할 수 있는가?
• RQ2 다양한 잠재 공간 순서를 갖춘 자기회귀 트랜스포머의 앙상블이 단일 순서보다 이상 탐지 및 분할을 향상시키는가?
• RQ3 잠재 공간 재샘플링 마스크를 도입하는 것이 고주파 영역의 이상 잔여물의 특이도에 어떤 영향을 미치는가?
• RQ4 현실 병리 데이터(WMH, MS 병변, BRATS)에서 이 방법이 최신 자동인코더와 비교해 어떻게 성능하는가?

주요 결과

실험 / 데이터 세트	지표	값 (예시)	메모
Synthetic Segmentation	Best achievable Dice (ceil Dice)	0.895	8 raster orderings ensemble (main result)
Synthetic Segmentation	Best achievable Dice (ceil Dice)	0.768	VQ-VAE + Transformer + Masked Residuals + Ordering (ablation)
Synthetic Segmentation	Best achievable Dice (ceil Dice)	0.675	VQ-VAE + Transformer (no mask)
Synthetic Segmentation	Best achievable Dice (ceil Dice)	0.457	VQ-VAE only
Synthetic Segmentation	Best achievable Dice (ceil Dice)	0.533	VAE Dense
Synthetic Segmentation	Best achievable Dice (ceil Dice)	0.213	AE Dense

• 합성 데이터에서 잔여 마스킹 및 다중 순서를 활용한 전체 방법은 래스터 순서를 포함한 앙상블에서 가장 높은 분할 Dice(최대 달성 가능) 0.895를 달성했으며, 순서를 없는 앙상블 비교에서는 0.768을 달성했다.
• 합성 데이터에서 이미지 단위 이상 탐지는 OOD 추정 능력이 강하게 나타나며 근사 OOD AUROC가 0.921(일반 목적 VQ-VAE에서는 0.932)이며 먼 거리에 있는 OOD AUROC가 1.000이다.
• 실제 UK Biobank 및 다른 데이터셋에서 제안된 방법은 이상 분할에서 최신 자동인코더보다 성능이 우수하였다(예: UKB: 0.232, MSLUB: 0.378, BRATS: 0.759, WMH: 0.429).
• 여덟 가지 순서를 사용하는 앙상블은 단일 순서 대비 현저한 향상을 보였고, 래스터 앙상블이 합성 데이터에서 가장 잘 보고된 분할 Dice를 달성했다.
• 재샘플링 마스크로 잔여를 마스킹하면 고주파 영역의 잘못된 레이블링이 감소하고 분할 신뢰도가 향상된다.