[논문 리뷰] Self-Supervised Learning for 3D Medical Image Analysis using 3D SimCLR and Monte Carlo Dropout

Yamen Ali, Aiham Taleb외 2명|arXiv (Cornell University)|2021. 09. 29.

Medical Imaging Techniques and Applications참고 문헌 2인용 수 8

한 줄 요약

이 논문은 3D MRI 및 CT 스캔과 같은 레이블이 없는 데이터를 사전학습하여 강력한 특징을 학습하는 3D 자기지도 학습 프레임워크인 3D SimCLR를 제안한다. 또한 추론 중 몬테카를로 드롭아웃을 통합하여 불확실성 추정을 수행함으로써, 특히 레이블이 제한된 경우에 뇌 및 췌장 종양 분할 작업에서 정확도와 데이터 효율성을 크게 향상시킨다.

ABSTRACT

Self-supervised learning methods can be used to learn meaningful representations from unlabeled data that can be transferred to supervised downstream tasks to reduce the need for labeled data. In this paper, we propose a 3D self-supervised method that is based on the contrastive (SimCLR) method. Additionally, we show that employing Bayesian neural networks (with Monte-Carlo Dropout) during the inference phase can further enhance the results on the downstream tasks. We showcase our models on two medical imaging segmentation tasks: i) Brain Tumor Segmentation from 3D MRI, ii) Pancreas Tumor Segmentation from 3D CT. Our experimental results demonstrate the benefits of our proposed methods in both downstream data-efficiency and performance.

연구 동기 및 목표

• 대규모 레이블이 없는 데이터를 활용해 자기지도 학습을 통해 3D 의료 영상의 부족한 레이블링 문제를 해결하기 위해 모델을 사전학습하는 것.
• 전문가 레이블링된 훈련 데이터의 양에 의존도를 줄임으로써 최종 분할 작업의 데이터 효율성을 향상시키기 위한 것.
• 몬테카를로 드롭아웃을 베이지안 근사로 사용하여 추론 중 불확실성 추정을 통합함으로써 모델 신뢰도를 향상시키기 위한 것.
• 두 가지 도전적인 3D 의료 영상 분할 작업인 뇌 종양 및 췌장 종양 분할 작업에서 제안된 방법을 평가하기 위한 것.
• 특히 데이터가 적은 환경에서 몬테카를로 드롭아웃을 통한 불확실성 인식 추론이 측정 가능한 성능 향상을 이끌어내는지 입증하기 위한 것.

제안 방법

• 사전학습은 레이블이 없는 MRI 및 CT 스캔에서 추출한 3D 패치를 대상으로 3D SimCLR를 사용하여 대비 학습을 수행하며, 랜덤한 증강 기법(예: 회전, 노이즈, 블러)을 적용하여 양성 쌍을 생성한다.
• 모델 아키텍처는 3D-CNN 인코더와 비선형 투사 헤드로 구성되며, 동일한 패치의 증강된 시각 간 유사도를 최대화하기 위해 정규화된 온도 조정 교차 엔트로피 손실(NT-Xent)을 사용하여 훈련한다.
• 최종 분할 작업을 위해 사전학습된 인코더는 U-Net 디코더와 함께 미세조정되며, 종양 분할에서의 클래스 불균형을 처리하기 위해 가중치가 부여된 Dice 손실을 사용한다.
• 몬테카를로 드롭아웃은 추론 중에 여러 번의 순방향 전파를 수행할 수 있도록 적용되어 불확실성 추정과 앙상블 평균화를 통한 강건성 향상을 가능하게 한다.
• 다양한 집계 전략—다수결정, Borda 순위, 가중 다수결정—을 평가하여 몬테카를로 드롭아웃 예측을 통합하고 분할 일관성을 향상시킨다.
• 몬테카를로 드롭아웃 예측의 백분위수 히트맵을 시각화하여 예측된 종양 영역 전반의 신뢰 수준을 시각화한다.

실험 결과

연구 질문

• RQ1 레이블이 없는 3D 의료 영상에서 3D SimCLR 사전학습이 무작위 또는 감독 사전학습 대비 최종 분할 성능을 향상시키는가?
• RQ2 추론 중 몬테카를로 드롭아웃을 통합하면 특히 데이터가 적은 환경에서 분할 정확도와 강건성이 향상되는가?
• RQ3 몬테카를로 드롭아웃을 통한 불확실성 추정은 예측 신뢰도에 어떤 영향을 미치며, 임상 의사결정 지원을 위해 의미 있는 방식으로 시각화될 수 있는가?
• RQ4 몬테카를로 드롭아웃의 성능 향상은 네트워크의 인코더, 디코더, 또는 둘 다에 적용되었을 때에 의존하는가?
• RQ5 다양한 몬테카를로 드롭아웃 설정(예: 드롭아웃 비율, 샘플 수)이 최종 분할 성능에 어떤 영향을 미치는가?

주요 결과

• 췌장 데이터셋에서 제안된 3D SimCLR 방법은 레이블이 5% 밖에 없는 경우 평균 Dice 스코어에서 기준 모델 대비 최대 8.5% 향상되었다.
• 추론 중 몬테카를로 드롭아웃을 적용함으로써 췌장 데이터셋에서 평균 Dice 스코어가 결정론적 기준 모델 대비 3.2% 향상되었으며, 특히 인코더에 드롭아웃을 적용했을 때 가장 큰 향상이 관찰되었다.
• 학습 데이터의 25%만으로 미세조정한 3D SimCLR 모델이 100% 데이터로 훈련된 기준 모델보다 높은 평균 Dice 스코어를 기록하여 강력한 데이터 효율성을 입증했다.
• 몬테카를로 드롭아웃 예측에 대해 가중 다수결정 전략이 다른 집계 방법보다 우수했으며, 레이블이 5% 이상 존재할 경우에 특히 두드러졌다.
• 몬테카를로 드롭아웃 예측의 히트맵 분석 결과, 5번째 백분위수 맵은 종양 핵심 영역에서 높은 신뢰도 영역을 나타냈고, 95번째 백분위수 맵은 증가한 가짜 양성 결과를 보여, 이 방법이 불확실성을 올바르게 반영할 수 있음을 확인했다.
• 절단 분석 결과 몬테카를로 드롭아웃의 성능 향상 요인은 주로 인코더에서의 불확실성 추정에 기인하며, 디코더에만 드롭아웃을 적용한 경우는 성능 향상이 미미한 것으로 확인되었다.