[논문 리뷰] Self-Supervised Pre-Training of Swin Transformers for 3D Medical Image Analysis

Yucheng Tang, Dong Yang외 6명|arXiv (Cornell University)|2021. 11. 29.

Medical Imaging and Analysis인용 수 46

한 줄 요약

이 논문은 Swin UNETR를 소개하며, 자기지도 프록시 작업으로 사전 학습된 5,050개의 CT 볼륨에서 3D 변환기 기반 인코더를 활용, 파인튜닝 후 BTCV 및 MSD 벤치마크에서 최첨단 세분화 성능을 달성한다.

ABSTRACT

Vision Transformers (ViT)s have shown great performance in self-supervised learning of global and local representations that can be transferred to downstream applications. Inspired by these results, we introduce a novel self-supervised learning framework with tailored proxy tasks for medical image analysis. Specifically, we propose: (i) a new 3D transformer-based model, dubbed Swin UNEt TRansformers (Swin UNETR), with a hierarchical encoder for self-supervised pre-training; (ii) tailored proxy tasks for learning the underlying pattern of human anatomy. We demonstrate successful pre-training of the proposed model on 5,050 publicly available computed tomography (CT) images from various body organs. The effectiveness of our approach is validated by fine-tuning the pre-trained models on the Beyond the Cranial Vault (BTCV) Segmentation Challenge with 13 abdominal organs and segmentation tasks from the Medical Segmentation Decathlon (MSD) dataset. Our model is currently the state-of-the-art (i.e. ranked 1st) on the public test leaderboards of both MSD and BTCV datasets. Code: https://monai.io/research/swin-unetr

연구 동기 및 목표

• 의료 영상 분석에 적합한 3D 변환기 기반 인코더(Swin UNETR) 개발.
• 해부학적 맥락에 맞춘 자기지도 프록시 작업 설계 및 통합(Inpainting, Rotation, Contrastive Learning).
• 대규모 비레이블 CT 코퍼스에서 효과적인 프리트레이닝을 수행하고 분할 작업으로의 이전(transfer) 입증.
• 공개 벤치마크 BTCV 및 MSD에서 성능을 검증하고 이전 상태의 최첨단과 비교.

제안 방법

• Swin UNETR를 계층적 3D Swin Transformer 인코더와 건너뛰기 연결이 있는 CNN 기반 디코더로 제안.
• 세 가지 자기지도 프록시 작업으로 인코더를 프리트레이닝: 마스킹된 볼륨 인페인팅, 3D 회전 예측, 그리고 대비 학습.
• 가중치가 동일한 다중 목적 손실 L_tot = λ1 L_inpaint + λ2 L_contrast + λ3 L_rot (λ1 = λ2 = λ3 = 1) 사용.
• 프리트레이닝 중 임의 하위 볼륨 자르기 및 증강을 수행하여 head/neck, chest, abdomen/pelvis 영역 전역에서 ROI 인식 표현 학습.
• BTCV 다기관 분할 및 MSD 작업에서 프리트레이닝된 인코더를 4-stage Swin Transformer 인코더와 CNN 디코더 및 skip connections로 파인튜닝.

실험 결과

연구 질문

• RQ1 자기지도 작업으로 프리트레이닝된 3D Swin Transformer 인코더가 CT 데이터에 대해 강건한 ROI-인식 표현을 학습할 수 있는가?
• RQ2 마스킹된 볼륨 인페인팅, 회전 예측 및 대비 코딩이 downstream 3D 의료 영상 분할을 시너지적으로 개선하는가?
• RQ3 대규모 비레이블 CT 데이터 세트에서의 프리트레이닝이 BTCV 및 MSD 벤치마크에서 성능, 데이터 효율성, 수렴 속도에 어떤 영향을 미치는가?

주요 결과

• Swin UNETR with self-supervised pre-training achieves state-of-the-art Dice scores on BTCV multi-organ segmentation.
• On MSD, Swin UNETR attains top performance in several tasks and the best overall Dice across ten tasks.
• Ablation shows combining all proxy tasks yields the best Dice (84.72% on BTCV in their study), with inpainting providing strong single-task gains.
• Pre-training reduces annotation effort, achieving higher performance with less labeled data (e.g., 10% of BTCV labels yields ~10% Dice improvement).
• Increasing pre-training data size and using all proxy tasks accelerates convergence and improves downstream accuracy.