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

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

AI in cancer detection인용 수 501

한 줄 요약

TransMIL은 Transformer 기반의 상관된 MIL 프레임워크를 도입하여 형태학적 정보와 공간 정보를 모두 활용하고, 여러 데이터셋에서 최첨단 MIL 방법을 능가하며 수렴 속도와 해석 가능성을 향상시킵니다.

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. Implementation is available at: https://github.com/szc19990412/TransMIL.

연구 동기 및 목표

• 이상적 i.i.d. 인스턴스 가정에서 벗어나 인스턴스 간의 상관관계를 모델링하여 약한 감독 하의 WSI 분류 문제를 해결한다.
• 패치 간 형태학적 및 공간적 관계를 포착하기 위한 Transformer 기반 MIL(TransMIL)을 제안한다.
• 병리학에서 상관된 MIL에 대해 수렴 가능하고 해석 가능하며 효율적인 3단계 알고리즘을 제공한다.
• 여러 공공 조직병리 데이터셋에서 우수한 성능과 더 빠른 수렴을 입증한다.

제안 방법

• 수렴 증명을 포함한 일반적인 3단계 상관된 MIL 프레임워크를 제안한다.
• 패치 임베딩 간의 장기 상관을 모델링하기 위해 Transformer 기반의 Temporal Patch Transformer(TPT) 모듈을 갖춘 TransMIL을 개발한다.
• 조건부 다중 스케일 공간 위치 인코딩을 위한 Pyramid Position Encoding Generator(PPEG)을 도입한다.
• TPT에서 Nyström 기반의 근사 자기 주의(attention) 방법을 사용하여 긴 시퀀스를 효율적으로 처리한다.
• 형태학(f), 공간(h), 두 성분의 합(fh) 임베딩을 결합한 뒤 학습된 풀링 행렬 P와 최종 매핑 g를 통해 가방 라벨을 예측한다.
• 교차 엔트로피 손실과 Lookahead 옵티마이저를 이용한 엔드투엔드 학습 설정을 ResNet50 기반 패치 임베딩(1024-d, 512-d로 축소)과 함께 제공한다.

실험 결과

연구 질문

• RQ1 WSI 내 인스턴스 간 상관관계 모델링이 i.i.d.-기반 MIL 방법과 비교하여 MIL 기반 분류 성능을 향상시키는가?
• RQ2 Transformer 기반 아키텍처를 대규모 가변 길이의 WSI 패치 시퀀스에 어떻게 적응시켜 공간 맥락을 보존할 수 있는가?
• RQ3 조건부 다중 스케일 위치 인코딩(PPEG)이 병리학 MIL의 진단 성능과 해석 가능성을 개선하는가?
• RQ4 TransMIL은 이진 및 다중 클래스 WSI 분류 작업과 비균형 데이터셋에서 효과적인가?
• RQ5 상관 MIL 접근법이 기존 MIL 방법보다 더 빠르게 수렴하는가?

주요 결과

• TransMIL은 세 가지 공개 데이터셋(CAMELYON16, TCGA-NSCLC, TCGA-RCC)에서 최첨단 성능을 달성한다.
• CAMELYON16에서 TransMIL ACCURACY 0.8837 및 AUC 0.9309를 달성하며 i.i.d.-기반 및 일부 비국소(non-local) 방법을 능가한다.
• TCGA-NSCLC에서 TransMIL은 ACC 0.8835 및 AUC 0.9603으로 경쟁 방법들보다 더 높은 성능을 보인다.
• TCGA-RCC에서 TransMIL은 ACC 0.9466 및 AUC 0.9882로 보고된 방법들 중 최고 성능을 달성한다.
• PPEG 및 조건부 위치 인코딩이 사인파 포지션 인코딩이나 비인코딩 대비 성능을 크게 향상시키는 것으로 나타난다.
• TransMIL은 수렴 속도가 더 빠르며, 경쟁 MIL 방법에 비해 2–3배 더 적은 학습 에포크가 필요하다.
• 주의도 맵의 시각화 결과가 암 영역과 높은 중복을 보이며 해석 가능성을 뒷받침한다.