[논문 리뷰] Video Pretraining Advances 3D Deep Learning on Chest CT Tasks

Alexander Ke, Shih-Cheng Huang외 4명|arXiv (Cornell University)|2023. 04. 02.

Radiomics and Machine Learning in Medical ImagingMedicine인용 수 3

한 줄 요약

이 논문은 Kinetics와 같은 대규모 외부 도메인 데이터셋에서의 비디오 사전 훈련이 흉부 CT 작업에서 3D 딥 러닝 성능을 크게 향상시킨다는 것을 입증한다. 이는 2D 모델과 내부 도메인 사전 훈련보다도 뛰어나며, 일곱 종류의 3D 모델과 두 가지 임상 작업(폐색전증 및 종양 탐지) 전반에 걸쳐 일관되게 성능 향상을 이끌어내며, 특히 소규모 데이터셋에서 두드러진다. 이는 자료가 부족한 3D 의료 영상 분석에 있어 대규모 비디오 사전 훈련이 더 유망한 전략임을 입증한다.

ABSTRACT

Pretraining on large natural image classification datasets such as ImageNet has aided model development on data-scarce 2D medical tasks. 3D medical tasks often have much less data than 2D medical tasks, prompting practitioners to rely on pretrained 2D models to featurize slices. However, these 2D models have been surpassed by 3D models on 3D computer vision benchmarks since they do not natively leverage cross-sectional or temporal information. In this study, we explore whether natural video pretraining for 3D models can enable higher performance on smaller datasets for 3D medical tasks. We demonstrate video pretraining improves the average performance of seven 3D models on two chest CT datasets, regardless of finetuning dataset size, and that video pretraining allows 3D models to outperform 2D baselines. Lastly, we observe that pretraining on the large-scale out-of-domain Kinetics dataset improves performance more than pretraining on a typically-sized in-domain CT dataset. Our results show consistent benefits of video pretraining across a wide array of architectures, tasks, and training dataset sizes, supporting a shift from small-scale in-domain pretraining to large-scale out-of-domain pretraining for 3D medical tasks. Our code is available at: https://github.com/rajpurkarlab/chest-ct-pretraining

연구 동기 및 목표

• 대규모 외부 도메인 비디오 데이터셋(예: Kinetics)에서의 사전 훈련이 흉부 CT 작업을 위한 3D 딥 러닝 모델 성능을 향상시키는지 평가하는 것, 특히 자료가 부족한 환경에서의 성능 향상 여부를 확인하는 것.
• 외부 도메인 비디오 사전 훈련(예: Kinetics)이 내부 도메인 CT 사전 훈련 및 순차적 사전 훈련과 비교하여 얼마나 효과적인지 평가하는 것.
• 비디오 사전 훈련이 다양한 3D 아키텍처와 임상 작업에 걸쳐 일반화되는지 평가하는 것.
• 특히 의료 분야에서 흔한 소규모 데이터 설정에서 사전 훈련 성능이 하류 데이터셋 크기에 따라 어떻게 변화하는지 이해하는 것.
• 미래의 3D 의료 AI 개발을 안내하기 위해 모델, 작업, 데이터 환경 전반에 걸쳐 사전 훈련 전략의 종합적 벤치마크를 제공하는 것.

제안 방법

• ImageNet과 Kinetics 데이터셋에서 비디오 사전 훈련을 수행한 일곱 종류의 3D 모델(예: MViT, Swin-T, SlowFast)과 세 종류의 2D 모델을 사전 훈련하였다.
• 내부 도메인인 스탠포드 PE 탐지 데이터셋에서 사전 훈련한 후 Kinetics에서 사전 훈련하는 순차적 사전 훈련을 수행하여 비디오 사전 훈련의 영향을 분리해내었다.
• 모든 모델을 두 개의 공개 흉부 CT 데이터셋인 RSNA(폐색전증 탐지용)와 LIDC-IDRI(폐 종양 탐지용)에서 미세 조정하였다.
• 사전 훈련 전략의 데이터 효율성과 확장성 평가를 위해 세 가지 훈련 데이터셋 크기(1%, 10%, 100%)에서 성능을 평가하였다.
• 모델 성능 비교를 위해 표준 지표(AUC)와 95% 신뢰구간을 사용하였다.
• 공정하고 일반화 가능한 비교를 위해 모델과 사전 훈련 프로토콜 전반에 걸쳐 초모수 최적화를 수행하였다.

Figure 1: Visual summary of our methods. We examined seven 3D (three 2D) models that were pre-trained either on Kinetics (ImageNet), or on the Stanford dataset for PE detection, or on both sequentially. We then transferred these weights and finetuned them on RSNA for PE detection and LIDC for lung n

실험 결과

연구 질문

• RQ1 대규모 외부 도메인 비디오(예: Kinetics)에서의 비디오 사전 훈련이 ImageNet 사전 훈련 또는 사전 훈련 없이도 3D 의료 영상 분류 성능을 향상시키는가?
• RQ2 특히 소규모 하류 데이터셋에서 비디오 사전 훈련이 내부 도메인 CT 사전 훈련과 비교하여 성능 향상에 얼마나 기여하는가?
• RQ3 비디오 사전 훈련의 이점이 3D 흉부 CT 분석에서 다양한 모델 아키텍처와 작업 다양성에 따라 확장되는가?
• RQ4 다양한 하류 데이터셋 크기에서 비디오 사전 훈련된 모델의 성능이 2D 베이스라인과 어떻게 비교되는가?
• RQ5 자료가 부족한 의료 영상 환경에서 사전 훈련 전략과 하류 데이터셋 크기 간의 상호작용은 어떠한가?

주요 결과

• Kinetics에서의 비디오 사전 훈련은 RSNA 및 LIDC 데이터셋 전반에서 일곱 종류의 3D 모델 평균 성능을 향상시켰으며, 하류 데이터셋 크기에 관계없이 동일하게 성능 향상을 보였다.
• 훈련 데이터의 1%만 사용한 RSNA 데이터셋(50건의 연구)에서 비디오 사전 훈련은 AUC 0.599를 기록하여 2D 베이스라인(AUC 0.517)과 내부 도메인 사전 훈련(AUC 0.510)을 뛰어넘었다.
• 훈련 데이터의 10%만 사용한 LIDC 데이터셋(71건의 연구)에서 비디오 사전 훈련은 AUC 0.831을 기록하여 2D 베이스라인(AUC 0.530)과 내부 도메인 사전 훈련(AUC 0.500)을 초월하였다.
• 비디오 사전 훈련은 모든 데이터셋 크기에서 3D 모델이 2D 모델을 능가하도록 했으며, 특히 자료가 부족한 환경(훈련 데이터의 1% 및 10%)에서 상대적 성능 향상이 가장 두드러졌다.
• 일반적으로 사용되는 내부 도메인 CT 데이터셋보다 Kinetics에서의 사전 훈련이 더 높은 성능을 보였으며, 이는 외부 도메인 비디오 사전 훈련이 3D 의료 작업에 더 효과적임을 시사한다.
• 내부 도메인 사전 훈련 후 Kinetics에서의 순차적 사전 훈련은 Kinetics 전용 사전 훈련보다 성능이 떨어지지 않았으며, 이는 비디오 사전 훈련의 효과가 누적되지 않으며, Kinetics 사전 훈련만으로도 충분하고 더 우수하다는 것을 의미한다.

Figure 2: Video pretraining consistently improves 3D models’ AUC on 10% and 100% of RSNA and LIDC, except MViT on 10% of LIDC ( $p=0.254$ )