[논문 리뷰] Supervised Transfer Learning at Scale for Medical Imaging

Basil Mustafa, Aaron Loh외 13명|arXiv (Cornell University)|2021. 01. 14.

COVID-19 diagnosis using AI참고 문헌 49인용 수 23

한 줄 요약

이 논문은 대규모 자연 이미지 데이터셋(예: JFT 및 ImageNet-21k)에서 사전 훈련된 모델을 사용한 대규모 지도형 전이 학습이, 흉부 X선 촬영, 유방 조영술, 피부과 영상과 같은 의료 영상 작업에서 성능 향상, 데이터 효율성 향상 및 분포 외 일반화 능력 향상에 크게 기여함을 보여준다. 이는 자연 이미지와 의료 영상 간의 상당한 도메인 이탈이 존재하더라도 성립한다.

ABSTRACT

Transfer learning is a standard technique to improve performance on tasks with limited data. However, for medical imaging, the value of transfer learning is less clear. This is likely due to the large domain mismatch between the usual natural-image pre-training (e.g. ImageNet) and medical images. However, recent advances in transfer learning have shown substantial improvements from scale. We investigate whether modern methods can change the fortune of transfer learning for medical imaging. For this, we study the class of large-scale pre-trained networks presented by Kolesnikov et al. on three diverse imaging tasks: chest radiography, mammography, and dermatology. We study both transfer performance and critical properties for the deployment in the medical domain, including: out-of-distribution generalization, data-efficiency, sub-group fairness, and uncertainty estimation. Interestingly, we find that for some of these properties transfer from natural to medical images is indeed extremely effective, but only when performed at sufficient scale.

연구 동기 및 목표

• 자연 이미지와 의료 영상 간의 도메인 이탈 문제를 현대적인 대규모 사전 훈련 기술이 극복할 수 있는지 조사하기 위해.
• 모델 규모와 사전 훈련 데이터셋 크기의 증가가 의료 영상에서 전이 학습 성능에 미치는 영향을 평가하기 위해.
• 분포 외 일반화, 데이터 효율성, 하위군 공정성, 불확실성 캘리브레이션과 같은 임상적 배포에 중요한 특성들을 평가하기 위해.
• 딥 레이어에서의 특징 재사용 및 가중치 갱신을 분석함으로써 성능 향상의 메커니즘을 이해하기 위해.

제안 방법

• ResNet 기반 아키텍처를 사용한 자연 이미지 데이터셋(예: JFT, ImageNet-21k)에서 100만~3억 장의 이미지로 사전 훈련된 Big Transfer(BiT) 모델을 활용하였다.
• 세 가지 다양한 의료 영상 작업인 흉부 레이저 촬영(CheXpert), 유방 조영술(UK 데이터셋), 피부과 영상(피부과 전문의가 레이블링한 이미지)에 대해 이러한 대규모 사전 훈련된 모델을 미세 조정하였다.
• 표준 평가 지표(AUC, 정확도) 외에도 분포 외 일반화, 데이터 효율성, 하위군 공정성, 기대 캘리브레이션 오차(ECE)와 같은 임상적으로 관련성이 높은 지표를 사용해 성능을 평가하였다.
• 사전 훈련된 모델과 미세 조정된 모델 간의 특징 유사도를 분석하기 위해 선형 중심화된 커널 정렬(CKA)을 적용하였다.
• ResNet 블록 내에서 파라미터 이동을 추적하여, 다양한 사전 훈련 스케일이 미세 조정 중 가중치 갱신에 미치는 영향을 연구하였다.
• 특정 작업에서 모델의 캘리브레이션과 공정성을 향상시키기 위해 부드러운 확률적 레이블과 포칼 손실을 사용하였다.

실험 결과

연구 질문

• RQ1 JFT 및 ImageNet-21k와 같은 자연 이미지 데이터셋에서 대규모 지도형 사전 훈련이 상당한 도메인 이탈이 존재하는 의료 영상 작업으로 효과적으로 전이될 수 있는가?
• RQ2 모델 규모와 사전 훈련 데이터셋 크기를 증가시키면 의료 영상에서 분포 이탈 상황에서도 데이터 효율성과 일반화 능력 향상에 기여하는가?
• RQ3 대규모 전이 학습이 하위군 공정성과 모델 캘리브레이션에 미치는 영향은 무엇인가—이러한 요소들은 임상적 배포에 있어 핵심 요소이다.
• RQ4 대규모 사전 훈련을 통해 얻은 깊이 있는 특징들이 미세 조정된 모델에서 얼마나 효과적으로 재사용될 수 있는가?
• RQ5 사전 훈련 스케일, 모델 아키텍처 규모, 그리고 미세 조정 중 파라미터 갱신 패턴 간의 관계는 어떠한가?

주요 결과

• 자연 이미지에서 대규모 사전 훈련(JFT, ImageNet-21k)을 수행한 결과, 표준 ImageNet 사전 훈련보다 의료 영상 작업에서 성능 향상이 뚜렷하게 향상되었다.
• 더 큰 데이터셋(JFT 등)에서 사전 훈련된 모델들은 전체 데이터 기반 기준선 성능을 30~60%의 훈련 데이터로도 달성하여 강력한 데이터 효율성을 보였다.
• 더 큰 사전 훈련 데이터셋을 사용할수록 분포 외 일반화 능력이 크게 향상되어 분포 이탈에 대한 강건성이 향상됨을 시사했다.
• 대규모 사전 훈련을 통해 하위군 공정성과 모델 캘리브레이션은 유지되거나 약간 향상되었으며, 인구 통계적 하위군 간 공정성 지표에 변화가 없었다.
• 선형 CKA 분석 결과, 더 큰 데이터셋에서 사전 훈련된 더 큰 모델일수록 사전 훈련 초기화 상태에서 유도된 고수준 특징을 더 잘 유지하고 재사용하는 것으로 나타났다.
• 파라미터 이동 분석 결과, 더 큰 사전 훈련 데이터셋을 사용할수록 초기 레이어에서 더 큰 가중치 갱신이 발생하고, 깊은 레이어에서는 더 작은 갱신이 일어나, 고수준 표현의 더 효과적인 재사용이 이루어짐을 시사했다.