[논문 리뷰] Attention-Challenging Multiple Instance Learning for Whole Slide Image Classification

Yunlong Zhang, Honglin Li외 4명|arXiv (Cornell University)|2023. 11. 13.

Advanced Image and Video Retrieval Techniques인용 수 4

한 줄 요약

이 논문은 주사형 슬라이드 이미지(WSI) 분류를 위한 새로운 MIL 프레임워크인 주의도전성 다중 인스턴스 학습(ACMIL)을 제안한다. 이는 주의 메커니즘이 더 도전적인 인스턴스에 집중하도록 유도함으로써 과적합을 방지한다. ACMIL은 다양한 구분 가능한 패턴을 포착하기 위해 다중 브랜치 주의(MBA)를 통합하고, 두드러진 인스턴스를 억제하기 위해 확률적 상위-K 인스턴스 마스킹(STKIM)을 도입하여, Camelyon16, BRACS, LBC 세 가지 WSI 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성하며 일반화 능력과 강인성을 향상시킨다.

ABSTRACT

In the application of Multiple Instance Learning (MIL) methods for Whole Slide Image (WSI) classification, attention mechanisms often focus on a subset of discriminative instances, which are closely linked to overfitting. To mitigate overfitting, we present Attention-Challenging MIL (ACMIL). ACMIL combines two techniques based on separate analyses for attention value concentration. Firstly, UMAP of instance features reveals various patterns among discriminative instances, with existing attention mechanisms capturing only some of them. To remedy this, we introduce Multiple Branch Attention (MBA) to capture more discriminative instances using multiple attention branches. Secondly, the examination of the cumulative value of Top-K attention scores indicates that a tiny number of instances dominate the majority of attention. In response, we present Stochastic Top-K Instance Masking (STKIM), which masks out a portion of instances with Top-K attention values and allocates their attention values to the remaining instances. The extensive experimental results on three WSI datasets with two pre-trained backbones reveal that our ACMIL outperforms state-of-the-art methods. Additionally, through heatmap visualization and UMAP visualization, this paper extensively illustrates ACMIL's effectiveness in suppressing attention value concentration and overcoming the overfitting challenge. The source code is available at \url{https://github.com/dazhangyu123/ACMIL}.

연구 동기 및 목표

• 제한된 데이터, 높은 해상도, 클래스 불균형으로 인해 WSI 분류의 다중 인스턴스 학습(MIL)에서 지속적인 과적합 문제를 해결한다.
• 기존 MIL 모델에서 과적합의 근본 원인을 분석하며, 특히 히트맵 시각화를 통해 일부 구분 가능한 인스턴스에 주의가 집중되는 경향을 규명한다.
• 특정 소수의 두드러진 패치에 의존하는 대신, 더 다양한 도전적인 인스턴스에 주의 메커니즘이 집중되도록 하는 방법을 개발한다.
• 주목적 인스턴스에 대한 과도한 의존도를 억제하고 주의 다양성을 증진시켜 WSI 분석에서 모델의 일반화 능력과 강인성을 향상시킨다.

제안 방법

• 다양한 패턴을 가진 구분 가능한 인스턴스를 포착하기 위해 병렬 주의 브랜치를 활용하는 다중 브랜치 주의(MBA)를 도입한다.
• 상위-K 주의 가중치가 부여된 인스턴스를 무작위로 마스킹하고 잔여 인스턴스에 주의 점수를 재분배하여 주의 집중을 방지하는 데이터 증강 유사 기법인 확률적 상위-K 인스턴스 마스킹(STKIM)을 제안한다.
• MBA와 STKIM을 통합한 유일한 ACMIL 프레임워크를 구성하여 주의 메커니즘이 더 도전적이고 명확하지 않은 인스턴스에 집중되도록 유도하고, 일반화 능력을 향상시킨다.
• 각 주의 브랜치가 서로 다른 구분 지능을 학습하도록 하는 다각도 손실 $\mathcal{L}_d$ 를 MBA에 통합하여 중복을 방지한다.
• UMAP 시각화와 히트맵 분석을 통해 모델의 개선된 주의 분포 및 특징 군집화 행동을 검증한다.
• ImageNet 및 SSL 사전학습을 사용한 두 가지 백본(ResNet18 및 ViT-S/16)을 활용해 세 가지 WSI 데이터셋(Camelyon16, BRACS, LBC)에서 ACMIL를 훈련 및 평가한다.

실험 결과

연구 질문

• RQ1 기존 MIL 모델에서 몇몇 구분 가능한 인스턴스에 주의가 집중되는 것이 WSI 분류에서 과적합에 얼마나 기여하는가?
• RQ2 더 두드러지지 않은, 더 도전적인 인스턴스에 주의를 집중시키는 것이 WSI 데이터셋에서 모델의 일반화 능력과 성능 향상에 기여하는가?
• RQ3 확률적 상위-K 인스턴스 마스킹(STKIM)이 주의 지배를 방해하고 주의 다양성을 증진하는 데 얼마나 효과적인가?
• RQ4 다양도 손실 $\mathcal{L}_d$ 를 갖춘 다중 브랜치 주의(MBA)가 단일 브랜치 주의보다 더 강인하고 일반화 능력이 뛰어난 특징 학습을 이끌어내는가?
• RQ5 히트맵 및 UMAP 시각화를 통해 ACMIL가 이전 최신 기술 수준(SOTA) 방법보다 더 구분력 있고 잘 분리된 특징을 학습하는가?

주요 결과

• ACMIL은 세 가지 WSI 데이터셋에서 최신 기술 수준(SOTA) MIL 방법을 크게 능가하며, ResNet18과 ImageNet 사전학습 조건에서 LBC 데이터셋에서 F1-score가 5.9% 향상되었다.
• MBA에 다양도 손실 $\mathcal{L}_d$ 를 포함시켰을 때 데이터셋 전반에서 5.3–8.0%의 성능 향상이 발생하여, 주의 다양성 증진에 핵심적인 역할을 한다는 것을 확인했다.
• 테스트 시점에서 STKIM은 성능을 약간 저하시키는 것으로 나타나, 추론 시에는 필요하지 않으며 테스트 시점에서의 정규화 기법으로서 작동하는 것으로 보이며, cutout과 유사하다.
• UMAP 시각화 결과 ACMIL는 더 구분력 있고 잘 분리된 특징을 학습하는 것으로 나타났으며, LBC 데이터셋에서 V-measure가 0.316으로 ABMIL의 0.224보다 높았다.
• 히트맵 시각화 결과 ACMIL는 일부 인스턴스에 대한 주의 집중을 줄이고 패치 전역에 걸쳐 주의를 더 넓게 분포시켜 과적합 감소와 관련이 있음을 확인했다.
• Camelyon16, BRACS, LBC에서 모든 평가 지표(F1-score, AUC)에서 ACMIL는 최신 기술 수준(SOTA) 성능을 달성했으며, ABMIL 및 기타 SOTA 기반 모델보다 일관된 향상이 이루어졌다.