[논문 리뷰] Task-specific Fine-tuning via Variational Information Bottleneck for Weakly-supervised Pathology Whole Slide Image Classification

Honglin Li, Chenglu Zhu외 6명|arXiv (Cornell University)|2023. 03. 15.

AI in cancer detection인용 수 6

한 줄 요약

이 논문은 약 10,000개의 패치를 1,000개 미만의 핵심 인스턴스로 압축함으로써, 각 WSI당 계산 비용을 줄이는 변동 정보 봉쇄 기반 미세조정 프레임워크를 제안한다. WSI 수준의 레이블만을 사용하여 백본의 엔드 투 엔드 훈련을 가능하게 함으로써, ImageNet 미리 훈련된 모델이나 자기지도 학습 특징에 비해 높은 정확도와 일반화 성능을 달성하며, 패치 레이블의 1% 미만을 사용함으로써 완전히 지도 학습 기반 성능에 가까운 성능을 달성한다.

ABSTRACT

While Multiple Instance Learning (MIL) has shown promising results in digital Pathology Whole Slide Image (WSI) classification, such a paradigm still faces performance and generalization problems due to challenges in high computational costs on Gigapixel WSIs and limited sample size for model training. To deal with the computation problem, most MIL methods utilize a frozen pretrained model from ImageNet to obtain representations first. This process may lose essential information owing to the large domain gap and hinder the generalization of model due to the lack of image-level training-time augmentations. Though Self-supervised Learning (SSL) proposes viable representation learning schemes, the improvement of the downstream task still needs to be further explored in the conversion from the task-agnostic features of SSL to the task-specifics under the partial label supervised learning. To alleviate the dilemma of computation cost and performance, we propose an efficient WSI fine-tuning framework motivated by the Information Bottleneck theory. The theory enables the framework to find the minimal sufficient statistics of WSI, thus supporting us to fine-tune the backbone into a task-specific representation only depending on WSI-level weak labels. The WSI-MIL problem is further analyzed to theoretically deduce our fine-tuning method. Our framework is evaluated on five pathology WSI datasets on various WSI heads. The experimental results of our fine-tuned representations show significant improvements in both accuracy and generalization compared with previous works. Source code will be available at https://github.com/invoker-LL/WSI-finetuning.

연구 동기 및 목표

• 고정된 ImageNet 또는 자기지도 학습 표현을 사용하는 기존의 약한 지도 학습 WSI 분류 방법의 높은 계산 비용과 성능 한계를 해결하기 위해.
• 기하급수적인 패치를 포함하는 기가픽셀 WSI에서, 백본의 엔드 투 엔드 미세조정을 Multiple Instance Learning (MIL) 프레임워크 내에서 가능하게 하기 위해.
• WSI 수준의 레이블을 사용하여 임의의 작업에 관계없이 특화된 자기지도 학습 특징을 작업에 특화된 표현으로 변환함으로써, 특징의 작업 특화성과 일반화 성능을 향상시키기 위해.
• 약한 지도 학습 환경에서 다양한 훈련 시점 증강 기법을 지원하여 도메인 이동 상황에서도 강건성을 향상시키기 위해.

제안 방법

• 기존에 10,000개 이상의 패치를 포함하는 WSI를 1,000개 미만의 가장 대표적인 인스턴스로 압축함으로써, 훈련 비용을 10배 이상 감소시키는 변동 정보 봉쇄(VIB) 모듈을 도입한다.
• VIB 모듈을 사용하여 WSI 백에 대한 최소한의 충분통계를 학습함으로써, 기하급수적인 패치 수에도 불구하고 백본을 통해 기울기 역전파를 가능하게 한다.
• VIB 기반 미세조정을 자기지도 학습 미리 훈련과 결합하여, WSI 수준의 레이블만을 사용하여 작업에 관계없이 특화된 SSL 특징을 작업에 특화된 표현으로 변환한다.
• 정보 압축과 예측 정확도 사이의 균형을 확보하기 위해 라그랑주 최적화 프레임워크를 활용함으로써, 모델이 중요한 진단 특징을 유지하도록 한다.
• 미세조정 중 다양한 훈련 시점 증강 기법을 적용하여, 도메인 이동이 발생하는 데이터셋에서의 일반화 성능을 향상시킨다.
• 각 WSI 백 내에서 가장 지원도가 높은 인스턴스를 식별하기 위해 미분 가능한 top-K 선택 메커니즘을 사용하여 효율적인 훈련을 위한 압축된 백을 구성한다.

실험 결과

연구 질문

• RQ1 변동 정보 봉쇄 프레임워크는 기가픽셀 WSI에서 훈련의 계산 부담을 효과적으로 줄이면서도 진단 정보를 유지할 수 있는가?
• RQ2 WSI 수준의 약한 지도 학습은 자기지도 학습 또는 ImageNet 미리 훈련된 특징의 작업 특화성을 어느 정도 향상시킬 수 있는가?
• RQ3 압축된 패치 표현을 사용한 백본의 엔드 투 엔드 미세조정은 고정된 특징 추출에 비해 성능을 더 좋게 만들 수 있는가?
• RQ4 다양한 데이터 증강 기법과 결합했을 때, 제안된 방법은 도메인 이동 상황에서도 효과적으로 일반화할 수 있는가?

주요 결과

• 제안된 미세조정 방법은 Camelyon-16-C에서 ImageNet 미리 훈련된 특징을 고정한 경우보다 F1 스코어가 18.43% 향상되었고, Camelyon-17에서는 18.86% 향상되었다.
• 미세조정된 특징을 사용한 Max-pooling은 Camelyon-16-C에서 CLAM-SB보다 10.91% 높은 F1 스코어, TransMIL보다 6.28% 높은 F1 스코어를 기록했다.
• 도메인 이동 상황에서도 일반화 성능이 향상되었으며, DTFD-MIL은 미세조정을 통해 Camelyon-17에서 F1 스코어가 11.30% 향상되어 강건성을 잘 보여주었다.
• AUC에 비해 F1 지표에서 훨씬 더 큰 향상이 있었으며, 이는 이 방법이 랭킹 능력보다 분류 성능(민감도 및 특이도)을 더 크게 향상시키며 임상적으로 더 관련성이 높다는 것을 시사한다.
• T-SNE 시각화 결과, 미세조정된 특징은 ImageNet-1K나 SSL 특징에 비해 더 분리 가능하고 작업에 특화된 분포를 형성하며, 더 높은 v-스코어를 기록했다.
• CLAM-SB의 어텐션 맵을 분석한 결과, 미세조정된 특징을 사용하는 모델이 종양 영역에 더 집중하는 경향을 보였으며, 이는 개선된 해석 가능성과 진단 관련성을 보여주었다.