[논문 리뷰] Towards Large-Scale Training of Pathology Foundation Models

kaiko. ai, Nanne Aben외 8명|arXiv (Cornell University)|2024. 03. 24.

AI in cancer detection인용 수 10

한 줄 요약

이 논문은 온라인 패칭을 사용하여 대규모 병리학 파운데이션 모델을 학습하는 확장 가능한 파이프라인을 제시하고, TCGA 데이터에서 여러 하이퍼파라미터와 배율을 평가하며, 표준화된 다운스트림 작업을 위한 평가 프레임워크 eva를 공개합니다.

ABSTRACT

Driven by the recent advances in deep learning methods and, in particular, by the development of modern self-supervised learning algorithms, increased interest and efforts have been devoted to build foundation models (FMs) for medical images. In this work, we present our scalable training pipeline for large pathology imaging data, and a comprehensive analysis of various hyperparameter choices and training techniques for building pathology FMs. We release and make publicly available the first batch of our pathology FMs (https://github.com/kaiko-ai/towards_large_pathology_fms) trained on open-access TCGA whole slide images, a commonly used collection of pathology images. The experimental evaluation shows that our models reach state-of-the-art performance on various patch-level downstream tasks, ranging from breast cancer subtyping to colorectal nuclear segmentation. Finally, to unify the evaluation approaches used in the field and to simplify future comparisons of different FMs, we present an open-source framework (https://github.com/kaiko-ai/eva) designed for the consistent evaluation of pathology FMs across various downstream tasks.

연구 동기 및 목표

• 대규모 WSIs에서 병리학 파운데이션 모델을 위한 확장 가능한 학습 파이프라인을 시연한다.
• 초기화, 배율 혼합, 데이터 크기 등의 하이퍼파라미터가 다운스트림 성능에 미치는 영향을 분석한다.
• 온라인 패칭이 오프라인 사전 생성 없이도 고처리량 패치 로딩을 가능하게 함을 보여준다.
• 다운스트림 작업 전반에 걸친 일관된 평가를 위한 오픈 프레임워크 (eva)를 제공한다.

제안 방법

• Blob 저장소에 저장된 WSIs로부터 고처리량의 패치 수준 로딩을 위한 온라인 패칭 개발.
• 다수의 배율에서 TCGA의 패치를 사용하여 DINO 및 DINOv2로 ViT 기반 파운데이션 모델을 사전 학습한다.
• ImageNet SSL 가중치에서 초기화하고 수렴 이점을 연구한다.
• 선형 프로빙을 사용하여 여러 패치 수준 다운스트림 작업(BACH, CRC, MHIST, PCam, TP53, CoNSeP)에서 모델을 평가한다.
• 강건성과 일반화 능력을 평가하기 위한 모델 크기와 배율 전략을 비교한다.

실험 결과

연구 질문

• RQ1 온라인 패칭이 성능 저하 없이 확장 가능하고 다양한 패치를 샘플링할 수 있게 하는가?
• RQ2 ImageNet에서의 초기화 및 사전 학습이 파운데이션 모델의 수렴 및 다운스트림 정확도에 어떤 영향을 미치는가?
• RQ3 여러 배율로 학습하는 것이 강건성과 작업 성능에 어떤 영향을 미치는가?
• RQ4 학습 데이터 크기(슬라이드 및 패치)가 인디스트리뷰션(IID) 및 아웃 오브 디스트리뷰션(OOD) 성능에 어떤 영향을 미치는가?

주요 결과

모델	학습 데이터	BACH	CRC	MHIST	PCam	TP53 **	CoNSeP **
ViT-S16 (rand.)	None	0.410	0.617	0.501	0.728	0.500	0.583
DINO ViT-S16 [24]	ImageNet	0.695	0.935	0.831	0.849	0.519	0.611
DINO ViT-B8 [24]	ImageNet	0.710	0.939	0.814	0.856	0.548	0.710
Lunit [29]	TCGA (21k WSIs)	0.801	0.934	0.768	0.895	0.571	0.654
Phikon [30]	TCGA (6k WSIs)	0.725	0.935	0.777	0.915	0.630	0.666
DINO ViT-S16 (ours)	TCGA (29k WSIs)	0.797	0.943	0.828	0.893	0.633	0.649
DINO ViT-S8 (ours)	TCGA (29k WSIs)	0.834	0.946	0.832	0.887	0.621	0.724
DINO ViT-B16 (ours)	TCGA (29k WSIs)	0.810	0.960	0.826	0.898	0.651	0.658
DINOv2 ViT-L14 (ours)	TCGA (29k WSIs)	0.870	0.930	0.809	0.898	0.656	0.679
Virchow [33]	Private (1.5M WSIs)	n/a	0.962 *	0.830 *	0.933 *	n/a	n/a

• 온라인 패칭은 최신 방법과 비교해 패치-수준 성능이 경쟁력 있거나 우수하면서도 확장 가능한 데이터 처리를 가능하게 한다.
• ImageNet 사전 학습 가중치에서 초기화하는 것이 수렴을 가속하고 다운스트림 성능을 향상시킨다.
• 다중 배율로 학습하는 것이 강건성을 향상시키고 단일 배율 모델보다 우수하다.
• 학습 슬라이드 수를 늘리면 일반적으로 성능이 향상되지만 수익 체감이 있으며, OOD 데이터에 대한 일반화를 개선하려면 다양한 데이터가 필요하다.
• 다양한 학습 패치의 수를 늘리면 ID 성능이 향상되고, 패치 다양성이 크게 증가하지 않는 한 OOD 작업의 이득은 제한적이다.