[논문 리뷰] Self-supervised learning for skin cancer diagnosis with limited training data

Hamish Haggerty, Rohitash Chandra|arXiv (Cornell University)|2024. 01. 01.

Cutaneous Melanoma Detection and Management인용 수 4

한 줄 요약

이 논문은 제한된 레이블이 부여된 데이터를 바탕으로 피부암 진단을 위한 자기지도학습(SSL)을 Barlow Twins를 활용해 제안하며, ImageNet에서의 SSL 사전학습이 전통적인 지도학습 기반 사전학습(ResNet-50)보다 높은 70%의 테스트 정확도를 기록해 66%를 기록한 것보다 뛰어나다는 것을 입증한다. 추가로, 레이블이 없는 피부 병변 데이터를 활용한 두 번째 SSL 사전학습 후 미세조정을 통해 성능 향상이 이루어지며, 비용이 많이 드는 데이터 레이블링에 대한 확장 가능한 대안을 제공한다.

ABSTRACT

Early cancer detection is crucial for prognosis, but many cancer types lack large labelled datasets required for developing deep learning models. This paper investigates self-supervised learning (SSL) as an alternative to the standard supervised pre-training on ImageNet for scenarios with limited training data using a deep learning model (ResNet-50). We first demonstrate that SSL pre-training on ImageNet (via the Barlow Twins SSL algorithm) outperforms supervised pre-training (SL) using a skin lesion dataset with limited training samples. We then consider extit{further} SSL pre-training (of the two ImageNet pre-trained models) on task-specific datasets, where our implementation is motivated by supervised transfer learning. This approach significantly enhances initially SL pre-trained models, closing the performance gap with initially SSL pre-trained ones. Surprisingly, further pre-training on just the limited fine-tuning data achieves this performance equivalence. Linear probe experiments reveal that improvement stems from enhanced feature extraction. Hence, we find that minimal further SSL pre-training on task-specific data can be as effective as large-scale SSL pre-training on ImageNet for medical image classification tasks with limited labelled data. We validate these results on an oral cancer histopathology dataset, suggesting broader applicability across medical imaging domains facing labelled data scarcity.

연구 동기 및 목표

• 희귀하거나 과소진단되는 암에 대해 레이블이 부족하고 균형이 깨진 의료 영상 데이터셋이 존재하는 문제를 해결한다.
• 피부암 분류의 저자료 환경에서 전통적인 지도학습 기반 전이학습 대비 자기지도학습 기반 사전학습이 우수한 성능을 내는지 조사한다.
• 레이블이 없는 타겟 도메인 데이터(피부 병변)에서 추가적인 자기지도학습 사전학습이 레이블 추가 없이 성능 향상에 기여하는지 탐색한다.
• 다양한 모델 아키텍처와 데이터 스케일에서 SSL과 지도학습 사전학습의 효과를 평가한다.
• 다른 저자원 의료 영상 분류 작업에 적용 가능한 확장 가능한, 자원 효율적인 프레임워크를 제공한다.

제안 방법

• ImageNet에서 레이블 없이 일반화 가능한 특징을 학습하기 위해 ResNet-50 기반 모델을 Barlow Twins 대비 자기지도학습 알고리즘을 사용해 사전학습한다.
• 이중 또는 다중 분류를 위한 작은 레이블이 부여된 피부 병변 데이터셋(ISIC)에서 사전학습된 SSL 모델을 미세조정한다.
• 동일한 ResNet-50를 사용해 ImageNet에서 지도학습 레이블을 사용해 사전학습한 표준 지도학습 기반 전이학습 베이스라인과 성능을 비교한다.
• 최종 미세조정 이전에 레이블이 없는 피부 병변 이미지를 사용해 Barlow Twins를 활용한 두 번째 자기지도학습 사전학습 단계를 수행한다.
• 이전 SSL 연구에서 입증된 효과성에 기반해 표준 데이터 증강 기법(예: 자르기, 색상 왜곡)을 SSL 사전학습 중에 사용한다.
• 표준 지표인 정확도와 균형 잡힌 정확도를 사용해 대규모 보류된 테스트 세트에서 모델 성능을 평가한다.

Figure 1: A residual block [ 24 ]

실험 결과

연구 질문

• RQ1 제한된 레이블이 부여된 데이터에서 피부암 진단에 있어 Barlow Twins를 활용한 자기지도학습 사전학습이 ImageNet에서의 지도학습 사전학습보다 우수한가?
• RQ2 레이블이 없는 피부 병변 이미지에서 추가적인 자기지도학습 사전학습이 ImageNet 사전학습을 초월해 분류 성능 향상에 기여하는가?
• RQ3 ImageNet에서 SSL을 통해 사전학습한 모델의 크기를 증가시킬 경우(예: ResNet-50보다 10배 이상 많은 파라미터), 전이 성능 향상이 이루어지는가?
• RQ4 클래스 불균형이 심한 의료 데이터셋 환경에서 SSL 전이 성능가 지도학습 전이 성능와 비교해 어떻게 다를까?
• RQ5 제안된 SSL 프레임워크는 다른 저자원 의료 영상 분류 문제에 일반화 가능한가?

주요 결과

• Barlow Twins를 활용한 자기지도학습 사전학습은 ISIC 피부 병변 데이터셋에서 평균 테스트 정확도 70%를 기록하여 지도학습 기반 전이학습의 66%를 초월했다.
• 레이블이 없는 피부 병변 이미지에서 두 번째 자기지도학습 사전학습을 수행함으로써 추가적인 성능 향상이 이루어졌으며, 이는 더 많은 레이블이 없는 데이터 수집이 비용이 많이 드는 레이블링 작업의 대안이 될 수 있음을 시사한다.
• ImageNet에서 SSL을 통해 사전학습한 더 큰 모델(ResNet-50보다 10배 이상 많은 파라미터)은 전이 성능 향상에 기여하지 않았으며, 이는 이 설정에서 모델 크기만으로는 SSL 전이 성능 향상이 이루어지지 않음을 의미한다.
• SSL과 지도학습 간의 정확도 격차는 뚜렷했으며, 저자원 조건에서 SSL이 지도학습 대비 정확도에서 4%p 향상된 성능을 보였다.
• 결과적으로, 특히 타겟 도메인 자기지도학습과 결합된 경우, 저레이블 데이터 환경에서 SSL 전이가 지도학습 전이보다 우수한 전략임을 시사한다.
• 이 프레임워크는 구강암 또는 췌장암 진단과 같이 레이블이 제한된 다른 의료 영상 분류 작업에 일반화 가능하고 적용 가능하다.

Figure 2: Joint embedding architecture