[논문 리뷰] TPMIL: Trainable Prototype Enhanced Multiple Instance Learning for Whole Slide Image Classification

Litao Yang, Deval Mehta외 4명|arXiv (Cornell University)|2023. 05. 01.

AI in cancer detection인용 수 4

한 줄 요약

TPMIL는 전체 패치 임베딩의 소프트 클러스터링을 통해 특징 공간을 개선하고 프로토타입 거리 기반 해석성을 향상시켜 약한 감독 하에 전체 슬라이드 이미지(WSI) 분류를 위한 학습 가능한 프로토타입 강화 다중 예비 학습 프레임워크를 제안한다. 이는 TCGA 뇌 및 폐암 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성하며, CLAM 및 DGMIL와 같은 이전 방법보다 AUC와 정확도에서 뛰어난 성능을 보인다.

ABSTRACT

Digital pathology based on whole slide images (WSIs) plays a key role in cancer diagnosis and clinical practice. Due to the high resolution of the WSI and the unavailability of patch-level annotations, WSI classification is usually formulated as a weakly supervised problem, which relies on multiple instance learning (MIL) based on patches of a WSI. In this paper, we aim to learn an optimal patch-level feature space by integrating prototype learning with MIL. To this end, we develop a Trainable Prototype enhanced deep MIL (TPMIL) framework for weakly supervised WSI classification. In contrast to the conventional methods which rely on a certain number of selected patches for feature space refinement, we softly cluster all the instances by allocating them to their corresponding prototypes. Additionally, our method is able to reveal the correlations between different tumor subtypes through distances between corresponding trained prototypes. More importantly, TPMIL also enables to provide a more accurate interpretability based on the distance of the instances from the trained prototypes which serves as an alternative to the conventional attention score-based interpretability. We test our method on two WSI datasets and it achieves a new SOTA. GitHub repository: https://github.com/LitaoYang-Jet/TPMIL

연구 동기 및 목표

• 고해상도와 패치 수준의 레이블 부족으로 인해 약한 감독 하에 전체 슬라이드 이미지(WSI) 분류를 수행하는 데 도전하는 문제를 해결하기 위해.
• 일부 극단적 패치에 의존하는 것 대신, 모든 인스턴스 임베딩을 활용하여 특징 공간 표현을 향상시키기 위해.
• 주의 점수를 넘어서 프로토타입 거리 기반 시각화를 도입하여 더 높은 해석성을 제공하기 위해.
• 학습된 프로토타입 거리 기반으로 종양 아형 간 상관관계를 모델링하기 위해.
• 특징 학습과 프로토타입 개선을 동시에 최적화하는 학습 가능한 종단 간 프레임워크를 개발하기 위해.

제안 방법

• 모든 패치 임베딩을 해당 프로토타입에 할당함으로써 소프트 클러스터링을 수행하는 학습 가능한 프로토타입 강화 다중 예비 학습 프레임워크(TPMIL)를 도입한다.
• 모든 인스턴스 특징의 분포에 기반해 프로토타입 중심을 갱신하는 미분 가능한 프로토타입 학습 메커니즘을 사용하여 전역 특징 공간 개선을 가능하게 한다.
• 패치의 해석성은 예측된 클래스의 가장 가까운 프로토타입까지의 L2 거리에 의해 결정되는 프로토타입 기반 주의 메커니즘을 활용한다.
• 정상 조직, 혈관, 잡음 등 진단에 기여하지 않는 영역을 포괄하기 위해 음성 프로토타입을 도입한다.
• 슬라이드 수준의 레이블을 사용하여 종단 간 학습을 수행하며, 특징 추출기, 프로토타입 중심, 분류기의 공동 최적화를 수행한다.
• 기존의 주의 점수를 프로토타입 거리 히트맵으로 대체하여 임상 시각화에서 더 뛰어난 강건성과 해석성을 확보한다.

Figure 1: Overview of TPMIL. The key novelty of our proposed framework is the integration of Prototype Learning Module for learning an optimal distribution of instances.

실험 결과

연구 질문

• RQ1 모든 패치 임베딩을 고려하는 프로토타입 기반 접근 방식이 극단적 패치에 의존하는 기존 방법에 비해 약한 감독 하에 WSI 분류의 특징 공간 표현을 향상시킬 수 있는가?
• RQ2 학습된 프로토타입 거리 간 상관관계는 WSI 내 종양 아형과 형태학적 차이와 어떻게 관련이 있는가?
• RQ3 프로토타입 거리 기반 히트맵이 주의 점수 기반 히트맵보다 더 강건하고 임상적으로 신뢰할 수 있는 해석성을 제공할 수 있는가?
• RQ4 제안된 방법이 다중 분류를 위한 벤치마크 WSI 데이터셋에서 최신 기술 수준 성능을 달성하는가?
• RQ5 프로토타입 간 거리 기반으로 종양 아형을 효과적으로 구분할 수 있는가?

주요 결과

• TPMIL는 TCGA 뇌 종양 데이터셋에서 기존 최고 성능(SOTA)을 경신하였으며, CLAM 대비 AUC가 0.24% 높고 정확도가 1.54% 높았다.
• TCGA 폐암 데이터셋에서 TPMIL는 CLAM 대비 AUC가 0.11% 높고 정확도가 1.41% 높은 성능을 기록하여 일관된 최신 기술 수준 성능을 입증하였다.
• 프로토타입 간 거리 행렬은 ACM(아스트로시토마)과 ODG(오ligodendroglioma)가 GBM(글리오블라스토마)보다 더 가까운 것을 보여주며, 이는 이들이 공통으로 저-grade glioma 상태를 공유하고 있음을 반영한다.
• 음성 프로토타입은 GBM보다 ACM과 ODG에 더 가까운 것으로 나타나, GBM의 형태학적 다양성과 더 높은 병리학적 복잡성이 정상 유사 영역과 더 뚜렷하게 다를 수 있음을 시사한다.
• TPMIL의 거리 기반 히트맵은 특히 글리오블라스토마의 괴사 영역을 강조하는 데 있어 주의 점수 기반 히트맵에 비해 잡음 및 아티팩트에 더 강건한 것으로 나타났다.
• 전문 임상의들은 프로토타입 거리 히트맵이 주의 기반 히트맵보다 더 신뢰할 수 있고 임상적으로 해석 가능한 시각화를 제공함을 확인하였으며, 특히 진단적 특징과 아티팩트를 구분하는 데서 유의미한 우월성을 보였다.

Figure 2: The distance matrix of TCGA Brain Tumor prototypes after training