[논문 리뷰] TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classication

Zhuchen Shao, Hao Bian외 5명|arXiv (Cornell University)|2021. 06. 02.

AI in cancer detection참고 문헌 38인용 수 26

한 줄 요약

TransMIL은 전체 슬라이드 이미지(WSIs) 내 패치 간의 공간적 및 형태학적 상관관계를 모델링하는 Transformer 기반의 상관 다중 예측 학습 프레임워크를 제안하며, 더 빠른 수렴과 향상된 해석 가능성으로 최신 기술 수준의 성능을 달성한다. 이는 이진 종양 분류에 대해 CAMELYON16에서 93.09% AUC를 기록하고, TCGA-RCC에서 최대 98.82% AUC를 달성한다.

ABSTRACT

Multiple instance learning (MIL) is a powerful tool to solve the weakly supervised classification in whole slide image (WSI) based pathology diagnosis. However, the current MIL methods are usually based on independent and identical distribution hypothesis, thus neglect the correlation among different instances. To address this problem, we proposed a new framework, called correlated MIL, and provided a proof for convergence. Based on this framework, we devised a Transformer based MIL (TransMIL), which explored both morphological and spatial information. The proposed TransMIL can effectively deal with unbalanced/balanced and binary/multiple classification with great visualization and interpretability. We conducted various experiments for three different computational pathology problems and achieved better performance and faster convergence compared with state-of-the-art methods. The test AUC for the binary tumor classification can be up to 93.09% over CAMELYON16 dataset. And the AUC over the cancer subtypes classification can be up to 96.03% and 98.82% over TCGA-NSCLC dataset and TCGA-RCC dataset, respectively.

연구 동기 및 목표

• 기존 다중 예측 학습(MIL) 방법이 독립적이고 동일하게 분포된 인스턴스를 가정하는 한계를 해결하기 위해 전체 슬라이드 이미지 내 패치 간의 상관관계를 모델링한다.
• WSI 내 형태학적 및 공간적 관계를 효과적으로 포착할 수 있는 약한 감독 학습 프레임워크를 개발한다.
• 불균형 및 다중 클래스 WSI 데이터셋에서 분류 성능를 향상시키면서도 높은 해석 가능성을 유지한다.
• 제안된 상관 다중 예측 학습 프레임워크의 수렴성에 대한 이론적 증명을 제공한다.

제안 방법

• 독립적이고 동일하게 분포된 인스턴스를 가정하는 기존 접근 방식을 대체하기 위해, 인스턴스 패치 간의 종속성을 모델링하기 위해 구조적 어텐션을 도입한 상관 다중 예측 학습 프레임워크를 제안한다.
• WSI 내 다수의 패치 인스턴스 간의 장거리 공간적 및 형태학적 종속성을 포착하기 위해 Transformer 인코더를 활용한다.
• 패치 수준의 특징을 글로벌 이미지 수준의 표현으로 집계하기 위해 학습 가능한 클래스 토큰과 다중 헤드 자기어텐션 메커니즘을 적용한다.
• 약한 감독 레이블을 사용한 엔드 투 엔드 훈련을 위해 글로벌 평균 풀링 레이어 이후 분류기 헤드를 사용한다.
• 패치 임bedding에 공간적 맥락을 유지하기 위해 위치 인코딩을 통합한다.
• 제안된 상관 다중 예측 학습 프레임워크에 대한 이론적 수렴 증명을 제공한다.

실험 결과

연구 질문

• RQ1 WSI 내 인스턴스 간 상관관계를 모델링하면 기존 표준 다중 예측 학습 방법에 비해 분류 성능 향상이 이루어지는가?
• RQ2 Transformer 기반 어텐션 메커니즘은 약한 감독 WSI 분류에서 특징 표현 학습을 어떻게 향상시키는가?
• RQ3 제안된 방법은 병리학적 이미지 분석에서 어텐션 맵의 해석 가능성과 시각화를 어느 정도 향상시키는가?
• RQ4 상관 다중 예측 학습 프레임워크는 이론적 보장 하에 수렴하는가? 기존 방법과 비교해 훈련 안정성은 어떻게 되는가?

주요 결과

• TransMIL은 이진 종양 분류에 대해 CAMELYON16 데이터셋에서 93.09% AUC를 기록하여 최신 기술 수준의 방법들을 능가한다.
• TCGA-NSCLC 데이터셋에서 TransMIL는 암 아형 분류에 대해 96.03% AUC를 달성하여 다중 클래스 문제에서 뛰어난 성능을 보였다.
• TCGA-RCC 데이터셋에서 TransMIL는 98.82% AUC를 기록하여 복잡한 신장암 아형 분류 과제에서 뛰어난 성능을 보였다.
• 기존의 다중 예측 학습 접근 방식에 비해 균형 및 불균형 데이터셋 모두에서 더 빠른 수렴과 더 나은 일반화 성능를 보였다.
• TransMIL가 생성한 어텐션 맵은 높은 해석 가능성을 보였으며, WSIs 내 관련 종양 영역을 명확히 국소화하였다.