[논문 리뷰] How Transferable Are Self-supervised Features in Medical Image Classification Tasks?

Tuan Truong, Sadegh Mohammadi외 1명|arXiv (Cornell University)|2021. 08. 23.

COVID-19 diagnosis using AI참고 문헌 11인용 수 12

한 줄 요약

이 논문은 의료 영상 분류에서 전이 학습을 위해 자기지도 학습(ImageNet) 사전 훈련 방법(SimCLR, SwAV, DINO)을 평가하며, 특히 데이터가 적은 상황에서 DINO 특징이 지도 학습 및 기타 자기지도 학습 기반 모델보다 뛰어난 성능을 보임을 입증한다. 또한, 다수의 사전 훈련된 특징을 동적 어텐션을 활용해 융합하는 모델에 종속적이지 않은 융합 방법인 동적 시각 메타임베딩(DVME)을 제안하며, 서로 보완적인 표현을 활용하여 네 가지 다양한 의료 영상 과제에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Transfer learning has become a standard practice to mitigate the lack of labeled data in medical classification tasks. Whereas finetuning a downstream task using supervised ImageNet pretrained features is straightforward and extensively investigated in many works, there is little study on the usefulness of self-supervised pretraining. In this paper, we assess the transferability of ImageNet self-supervisedpretraining by evaluating the performance of models initialized with pretrained features from three self-supervised techniques (SimCLR, SwAV, and DINO) on selected medical classification tasks. The chosen tasks cover tumor detection in sentinel axillary lymph node images, diabetic retinopathy classification in fundus images, and multiple pathological condition classification in chest X-ray images. We demonstrate that self-supervised pretrained models yield richer embeddings than their supervised counterpart, which benefits downstream tasks in view of both linear evaluation and finetuning. For example, in view of linear evaluation at acritically small subset of the data, we see an improvement up to 14.79% in Kappa score in the diabetic retinopathy classification task, 5.4% in AUC in the tumor classification task, 7.03% AUC in the pneumonia detection, and 9.4% in AUC in the detection of pathological conditions in chest X-ray. In addition, we introduce Dynamic Visual Meta-Embedding (DVME) as an end-to-end transfer learning approach that fuses pretrained embeddings from multiple models. We show that the collective representation obtained by DVME leads to a significant improvement in the performance of selected tasks compared to using a single pretrained model approach and can be generalized to any combination of pretrained models.

연구 동기 및 목표

• 낮은 데이터에서 의료 영상 분류 과제에 대해 자기지도 학습 ImageNet 특징의 전이 가능성 평가.
• 최신 자기지도 학습 모델(SimCLR, SwAV, DINO)이 지도 학습 ImageNet 사전 훈련과 비교하여 성능를 어떻게 냅니다.
• 다양한 의료 영상 과제에서 다수의 사전 훈련된 특징을 융합하면 최종 성능 향상이 가능한지 조사.
• 의료 영상에서 전이 학습을 향상시키기 위해 모델에 종속적이지 않은 융합 방법인 동적 시각 메타임베딩(DVME)을 제안하고 검증.

제안 방법

• 네 가지 의료 과제에서 지도 학습 및 자기지도 학습 ImageNet 모델의 고정된 특징을 사용하여 선형 평가 및 미세조정(fine-tuning)을 평가.
• 다양한 사전 훈련된 모델의 특징을 연결하고, 가속화 가능한 자기주의 어텐션 레이어를 사용해 기여도를 적응적으로 가중하는 동적 시각 메타임베딩(DVME)을 적용.
• 백본 모델을 미세조정하지 않고 고정된 특징 추출기 위에 메타임베딩 헤드만을 훈련시켜 엔드 투 엔드 전이 학습을 가능하게 함.
• 융합된 표현을 해석하고 특징의 보완성 여부를 평가하기 위해 t-SNE 시각화 및 어텐션 행렬 분석을 수행.
• DVME에서 자기주의 어텐션의 기여도를 평가하기 위해 아블레이션 스터디를 수행하며, 비어느 어텐션 융합 기반 베이스라인과 비교.
• 각 과제의 소규모(S), 중간규모(M), 전체(F) 데이터 서브세트에서 표준 지표인 AUC, 카파 점수, 정확도를 사용.

실험 결과

연구 질문

• RQ1 낮은 데이터 환경에서 의료 영상 분류 과제에 대해 자기지도 학습 ImageNet 특징이 지도 학습 특징보다 전이 가능성 면에서 어떻게 비교되는가?
• RQ2 SimCLR, SwAV, DINO 중 어느 자기지도 학습 사전 훈련 방법이 다양한 의료 영상 과제에서 가장 전이 가능한 특징을 제공하는가?
• RQ3 메타임베딩 접근법을 통해 다수의 사전 훈련된 특징을 융합하면 최고의 단일 모델 기반 성능을 초월할 수 있는가?
• RQ4 DVME의 자기주의 어텐션 메커니즘이 단순 연결 또는 고정 가중치와 비교해 성능 향상에 기여하는 방식은 무엇인가?

주요 결과

• DINO 특징가 대부분의 과제에서 가장 높은 성능를 기록했으며, 특히 데이터가 적은 상황에서 지도 학습 및 기타 자기지도 학습 기반 모델을 능가했다.
• 선형 평가에서, 흉부 X-ray에서 병리 상태 탐지 과제에서 DINO는 1%의 데이터만으로 AUC를 9.4% 향상시켰고, 당뇨성 망막병변 진단 과제에서 카파 점수를 14.79% 향상시켰다.
• 소규모 데이터 서브태스크 4개 전부에서 DVME는 최고의 단일 모델 기반 베이스라인을 초월했으며, 중간규모 데이터 서브태스크 4개 중 3개, 전체 데이터 서브태스크 4개 중 2개에서 승리했다.
• APTOS 당뇨성 망막병변 데이터셋에서 DVME는 소규모 및 중간규모 데이터 설정에서 최고의 단일 모델 기반 베이스라인 대비 카파 점수 6% 향상시켰다.
• DVME의 자기주의 어텐션은 성능 향상에 크게 기여했으며, APTOS 서브태스크 전반에서 카파 점수 4.4%에서 5.6%까지 향상시켜 보완적인 특징을 효과적으로 동적 가중치로 조합함을 시사했다.
• t-SNE 시각화 결과, DVME 임베딩은 특히 다중 클래스 설정에서 더 잘 분리된 클러스터를 생성함을 확인했으며, 이는 향상된 표현 품질을 뒷받침한다.