[논문 리뷰] Cluster-to-Conquer: A Framework for End-to-End Multi-Instance Learning for Whole Slide Image Classification

Yash Sharma, Aman Shrivastava외 4명|arXiv (Cornell University)|2021. 03. 19.

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

한 줄 요약

C2C는 슬라이드당 패치를 클러스터링하고, 클러스터에서 샘플링하며 KL-발산 정규화를 포함한 적응형 주의(attention)으로 슬라이드 수준 예측을 개선하는 엔드-투-엔드 MIL 프레임워크를 제시한다.

ABSTRACT

In recent years, the availability of digitized Whole Slide Images (WSIs) has enabled the use of deep learning-based computer vision techniques for automated disease diagnosis. However, WSIs present unique computational and algorithmic challenges. WSIs are gigapixel-sized ($\sim$100K pixels), making them infeasible to be used directly for training deep neural networks. Also, often only slide-level labels are available for training as detailed annotations are tedious and can be time-consuming for experts. Approaches using multiple-instance learning (MIL) frameworks have been shown to overcome these challenges. Current state-of-the-art approaches divide the learning framework into two decoupled parts: a convolutional neural network (CNN) for encoding the patches followed by an independent aggregation approach for slide-level prediction. In this approach, the aggregation step has no bearing on the representations learned by the CNN encoder. We have proposed an end-to-end framework that clusters the patches from a WSI into ${k}$-groups, samples ${k}'$ patches from each group for training, and uses an adaptive attention mechanism for slide level prediction; Cluster-to-Conquer (C2C). We have demonstrated that dividing a WSI into clusters can improve the model training by exposing it to diverse discriminative features extracted from the patches. We regularized the clustering mechanism by introducing a KL-divergence loss between the attention weights of patches in a cluster and the uniform distribution. The framework is optimized end-to-end on slide-level cross-entropy, patch-level cross-entropy, and KL-divergence loss (Implementation: https://github.com/YashSharma/C2C).

연구 동기 및 목표

• 기가픽셀 WSIs의 도전 과제를 해결하기 위해 MIL 기반 WSI 분류를 위한 엔드-투-엔드 학습을 가능하게 한다.
• 다양한 판별 패치 특징에 모델을 노출시키기 위해 클러스터 기반 샘플링을 활용한다.
• 패치 인코딩, 주의 기반 집계, 그리고 KL-발산 정규화를 융합하여 패치 및 슬라이드 표현의 공동 학습을 개선한다.
• 위장 질환 및 유방암 데이터 세트에서 두 단계 MIL 방법에 비해 경쟁력 있거나 우수한 성능을 보여준다.

제안 방법

• 패치 임베딩에서 k-means를 사용해 각 WSI의 패치를 에포크당 k개 클러스터로 클러스터링한다.
• 각 클러스터에서 k'개의 패치를 샘플링해 각 WSI에 대해 관리 가능한 학습 부분집합을 구성한다.
• CNN 인코더로 패치를 인코딩하고 패치 표현 h를 계산한 후 두 계층의 어텐션 모듈을 적용해 인스턴스 가중치 a_n을 얻는다.
• 어텐션 풀링으로 패치 표현을 WSI 표현 z로 집계하고 슬라이드 수준 라벨을 예측한다.
• L = alpha*L_WSI + beta*L_Patch + gamma*L_KLD로 결합된 손실로 엔드 투 엔드 학습하며, L_KLD는 클러스터 내 어텐션 분산을 정규화한다.
• 각 클러스터 내 패치-어텐션 분포와 균등 분포 간의 KL-발산을 적용해 정규화한다.

실험 결과

연구 질문

• RQ1 클러스터 기반 샘플링이 MIL에서 WSIs에 대해 학습된 패치 표현의 다양성과 품질에 어떤 영향을 미치는가?
• RQ2 주의 기반 집계가 포함된 엔드-투-엔드 학습이 두 단계 MIL 방법보다 슬라이드 수준 정확도를 향상시킬 수 있는가?
• RQ3 KL-발산 정규화가 어텐션 분포와 모델 성능에 미치는 영향은 무엇인가?
• RQ4 GI 생검 WSIs와 CAMELYON16 유방암 데이터에서 C2C가 완전 감독 및 두 단계 방식과 비교하여 어떤 성능을 보이는가?

주요 결과

Method	Accuracy	Precision	Recall	F1-Score
Campanella-MIL	82.8	94.9	74.5	83.5
Campanella-MIL RNN	74.7	75.4	84.3	79.6
Two-Stage Mean	81.6	87.3	80.3	83.7
C2C (w WSI Loss)	81.6	80.7	90.1	85.2
C2C (w WSI+KLD Loss)	83.9	84.9	86.3	85.4
C2C (w WSI+Patch Loss)	85.1	86.5	88.2	87.4
C2C (w WSI+Patch+KLD Loss)	86.2	85.5	92.2	88.7

• C2C는 셀리악병 대 정상 분류를 위한 위장 데이터셋에서 두 단계 MIL 기준선보다 우수하다.
• WSI 손실, 패치 손실, 및 KL-발산 손실의 결합은 더 높은 F1 점수와 경쟁력 있는 정밀도/재현율을 제공합니다.
• C2C는 슬라이드 수준 감독 없이 ResNet-18 백본으로 CAMELYON16에서 강한 성능을 달성한다.
• KL-발산 정규화는 양성 인스턴스 클래스 간 어텐션을 안정화한다(예: MNIST 가방 실험).
• 클러스터 기반 샘플링은 다양한 판별 패치를 노출시키는 데 도움이 되어 엔드-투-엔드 학습을 강화한다.