[논문 리뷰] Exploring Visual Prompts for Whole Slide Image Classification with Multiple Instance Learning

Yi Lin, Zhongchen Zhao외 4명|arXiv (Cornell University)|2023. 03. 23.

Digital Imaging for Blood Diseases인용 수 4

한 줄 요약

이 논문은 자연 이미지와 조직병리학적 이미지 간의 도메인 차이를 해결하기 위해, 이미지넷 사전 훈련된 모델을 전체 슬라이드 이미지(WSI) 분류에 적응시키기 위해 시각적 프롬프트를 도입한다. 다중 예제 학습(MIL)을 통해 수행되며, 백본을 동결하고 대표적인 패치에서만 경량 프롬프트 구성 요소를 미세조정함으로써, Camelyon16 및 TCGA-NSCLC 데이터셋에서 다양한 MIL 모델과 백본에 걸쳐 일관된 성능 향상을 달성한다.

ABSTRACT

Multiple instance learning (MIL) has emerged as a popular method for classifying histopathology whole slide images (WSIs). However, existing approaches typically rely on pre-trained models from large natural image datasets, such as ImageNet, to generate instance features, which can be sub-optimal due to the significant differences between natural images and histopathology images that lead to a domain shift. In this paper, we present a novel, simple yet effective method for learning domain-specific knowledge transformation from pre-trained models to histopathology images. Our approach entails using a prompt component to assist the pre-trained model in discerning differences between the pre-trained dataset and the target histopathology dataset, resulting in improved performance of MIL models. We validate our method on two publicly available datasets, Camelyon16 and TCGA-NSCLC. Extensive experimental results demonstrate the significant performance improvement of our method for different MIL models and backbones. Upon publication of this paper, we will release the source code for our method.

연구 동기 및 목표

• 이미지넷 사전 훈련된 모델과 조직병리학적 WSI 사이의 도메인 차이를 해결하여, MIL 기반 WSI 분류에서 특징 품질을 향상시키는 것.
• 제한된 패치 수준의 레이블이 존재할 때 전체 특징 추출기의 미세조정이 과적합 위험과 사전 훈련된 표현의 열화를 초래하는 문제를 해결하는 것.
• 시각 분야에서 프롬프트 학습을 조직병리학에 적용하기 위해, 사전 훈련된 특징을 병리학적 데이터에 적응시키는 시각적 프롬프트를 도입하는 것.
• 백본을 동결하고 프롬프트 구성 요소만 업데이트하는 경량 훈련 기법을 개발하여 계산 비용을 줄이는 것.
• 공개 WSI 데이터셋에서 다양한 백본(ResNet-18/50)과 MIL 모델(DTFD, ABMIL)에 대해 시각적 프롬프트의 효과를 검증하는 것.

제안 방법

• 사전 훈련된 이미지넷 백본을 동결하고, 학습 가능한 프롬프트 구성 요소를 추가함으로써 조직병리학적 데이터에 특징를 적응시키는 새로운 시각적 프롬프트 학습 프레임워크를 제안한다.
• 프롬프트 훈련에 사용되는 패치 수를 줄이고 효율성을 높이며 과적합 위험을 감소시키기 위해 대표성 있는 패치 선택(RPS) 전략을 사용한다.
• 백본 가중치를 동결한 채로 프롬프트 블록과 경량 MIL 분류기만 훈련함으로써 효율적인 도메인 적응을 가능하게 한다.
• 이중 단계 프로세스를 사용한다: 첫째, 특징 다양성을 기반으로 WSI에서 대표성 있는 패치를 선별하고, 둘째, 이러한 패치를 사용하여 프롬프트 구성 요소만 미세조정한다.
• 프롬프트 모듈을 특징 추출 헤드에 통합하여, 백본은 고정된 채로 프롬프트와 MIL 분류기의 엔드 투 엔드 훈련을 가능하게 한다.
• 주어진 프롬프트 적응 후에 패치 수준의 특징을 WSI 수준의 예측으로 집계하기 위해 주로 주의 기반 MIL 모델(DTFD, ABMIL 등)을 활용한다.

Figure 1: An illustration of three adapting schemes for WSI classification with MIL: (a) freezing the pretrained feature extractor while only training the MIL classifier; (b) fine-tuning the feature extractor and training the MIL classifier; (c) only training the prompt blocks and the MIL classifier

실험 결과

연구 질문

• RQ1 MIL 기반 분류에서 이미지넷과 조직병리학적 WSI 간의 도메인 차이를 효과적으로 줄일 수 있는가?
• RQ2 패치 수준의 레이블이 없을 때, 전체 특징 추출기의 미세조정보다 프롬프트 학습이 성능을 더 잘 향상시키는가?
• RQ3 프롬프트 블록의 수가 WSI 분류 성능과 안정성에 어떤 영향을 미치는가?
• RQ4 과적합을 피하면서 효과적인 프롬프트 훈련을 위해 최적의 대표 패치 수는 얼마인가?
• RQ5 제안된 방법이 다양한 백본과 MIL 아키텍처 간에 일반화 가능한가?

주요 결과

• ResNet-18를 사용할 때 Camelyon16 데이터셋에서 기준 모델 대비 AUC가 3.83% 향상되었고, TCGA-NSCLC에서는 1.28% 향상되었다.
• ResNet-50를 사용할 경우, 백본을 동결한 기준 모델과 전체 미세조정 방식(RPS-FT)보다도 일관되게 뛰어난 성능을 보이며, 더 뛰어난 일반화 능력과 안정성을 입증했다.
• 기준 모델 대비 프롬프트 블록을 하나 이상 추가하면 AUC가 약 1% 향상되어, 프롬프트 구성 요소의 효과성과 안정성을 입증했다.
• DTFD 및 ABMIL와 같은 다양한 MIL 모델과 ResNet-18, ResNet-50와 같은 다양한 백본에서 일관된 성능 향상을 유지함으로써 광범위한 적용 가능성을 보였다.
• 모든 K 값의 대표 패치 수에서 RPS-FT보다 뛰어난 성능을 달성하면서도 GPU 자원의 50% 미만을 사용하여 효율성을 입증했다.
• 제거 분석 결과, 전체 백본의 미세조정는 과적합과 성능 열화를 초래하는 반면, 프롬프트 기반 적응은 이러한 문제를 피함을 확인했다.

Figure 2: Overview of our method. ( \@slowromancap i@) MIL classifier training with all patch features. ( \@slowromancap ii@) Representative patch selection. ( \@slowromancap iii@) Prompt learning with representative patches.