[논문 리뷰] Embedding Task Knowledge into 3D Neural Networks via Self-supervised Learning

Jiuwen Zhu, Yuexiang Li외 2명|arXiv (Cornell University)|2020. 06. 10.

Domain Adaptation and Few-Shot Learning참고 문헌 13인용 수 4

한 줄 요약

이 논문은 3D 의료 영상 분류 성능을 햖을 수 있도록 병변 관련 사전 지식을 3D CNN에 통합하는 자기지도 학습 방법인 작업 관련 대비 예측 코드(TCPC)를 제안한다. SLIC를 이용한 초음소 클러스터링을 통해 후보 병변을 국소화한 후, 이러한 영역 주변의 부분 부피에 대해 校정된 'U'자형 대비 예측 코드 방식을 적용하여, 최소한의 레이블 데이터로도 뇌 출혈 및 폐암 데이터셋에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Deep learning highly relies on the amount of annotated data. However, annotating medical images is extremely laborious and expensive. To this end, self-supervised learning (SSL), as a potential solution for deficient annotated data, attracts increasing attentions from the community. However, SSL approaches often design a proxy task that is not necessarily related to target task. In this paper, we propose a novel SSL approach for 3D medical image classification, namely Task-related Contrastive Prediction Coding (TCPC), which embeds task knowledge into training 3D neural networks. The proposed TCPC first locates the initial candidate lesions via supervoxel estimation using simple linear iterative clustering. Then, we extract features from the sub-volume cropped around potential lesion areas, and construct a calibrated contrastive predictive coding scheme for self-supervised learning. Extensive experiments are conducted on public and private datasets. The experimental results demonstrate the effectiveness of embedding lesion-related prior-knowledge into neural networks for 3D medical image classification.

연구 동기 및 목표

• 딥 러닝에서 제한된 레이블이 부여된 3D 의료 영상 데이터 문제를 해결하기 위해.
• 자기지도 사전 훈련에 작업 특화 사전 지식을 통합하여 3D 의료 영상 분류 성능을 향상시키기 위해.
• 일반적인 이미지 증강 작업과는 달리 임상 진단과 더 관련성이 높은 프록시 작업을 개발하기 위해.
• 병변 인식 자기지도 사전 훈련이 특히 레이블이 적은 데이터셋에서 후행 분류 성능을 향상시킬 수 있음을 입증하기 위해.

제안 방법

• 3D 의료 볼륨에서 잠재적 병변 영역를 식별하기 위해 단순 선형 반복 클러스터링(SLIC)을 사용하여 초음소를 생성한다.
• 초음소 영역 중심의 부분 부피를 추출하여 진단적으로 관련성이 높은 영역에 집중한다.
• 대비 예측 코드에서 지원 특징 생성을 위한 校정된 'U'자형 경로를 제안하여, 문맥 특징 학습을 향상시킨다.
• 부분 부피에 대해 대비 예측 코드(CPC)를 적용하여, 공간적으로 구조화된 지원 세트의 특징을 사용해 타겟 큐브의 특징을 예측한다.
• 특징 추출 및 TCPC 프록시 작업을 사용한 레이블이 없는 데이터에 대한 사전 훈련을 위해 3D CNN 기반 아키텍처를 활용한다.
• UCF101에서의 전이 학습과 TCPC 사전 훈련을 조합하여 의료 데이터셋에서의 성능을 더욱 향상시킨다.

실험 결과

연구 질문

• RQ1 병변 관련 사전 지식을 통합한 자기지도 학습 방법이 3D 의료 영상 분류 성능을 향상시킬 수 있는가?
• RQ2 병변 국소화에 기반한 작업 특화 프록시 작업은 일반적인 프록시 작업보다 자기지도 3D 학습에서 어떻게 비교되는가?
• RQ3 소량의 레이블 데이터만 존재할 경우 TCPC는 성능을 얼마나 향상시키는가?
• RQ4 자연 이미지에서의 전이 학습과 작업 인식 자기지도 사전 훈련을 조합하면 분류 정확도가 더욱 향상되는가?

주요 결과

• 뇌 출혈 데이터셋에서 TCPC는 평균 78.38% 정확도와 0.921 AUC를 기록하여 3D CPC 및 루빅스 큐브와 같은 모든 베이스라인을 초월했다.
• 폐암 데이터셋에서 TCPC는 3D VGG를 사용해 95.86%의 정확도를 달성했고, 3D ResNet-18를 사용해 98.88%의 정확도를 기록하여 3D 재배열 및 루빅스 큐브 방법을 모두 능가했다.
• UCF101 + TCPC 캐스케이드 방법은 폐 결절에서 3D VGG와 ResNet-18에 대해 평균 1%의 정확도 향상을 기록했으며, 각각 97.85%와 99.40%의 정확도를 달성했다.
• TCPC는 하이퍼파rameter에 대해 강건했으며, 큐브 크기가 커질수록 더 풍부한 문맥 정보를 제공함에 따라 성능이 향상되었다.
• 최소한의 레이블 데이터로도 TCPC는 초기 학습보다 정확도와 AUC를 일관되게 향상시켰으며, 특히 데이터가 적은 환경에서 두드러진 성능 향상을 보였다.