[논문 리뷰] Self-supervised Feature Learning for 3D Medical Images by Playing a Rubik's Cube

Xinrui Zhuang, Yuexiang Li외 4명|arXiv (Cornell University)|2019. 10. 05.

Medical Imaging and Analysis참고 문헌 12인용 수 15

한 줄 요약

이 논문은 3D 볼륨 큐브를 뒤섞고 돌려 재구성하는 새로운 프록시 작업인 루빅스 큐브 복구를 사용하여 의료 영상에 대한 자기지도 학습 3D 특징 학습 프레임워크를 제안한다. 네트워크는 이 과정을 통해 이동 및 회전 불변성을 학습한다. 이 방법은 뇌 출혈 분류(83.8% 정확도)와 뇌 종양 세분화(77.3% mIoU)에서 최신 기술을 상회하는 성능을 달성하며, 추가 데이터 없이도 초기 학습 및 자연 영상 데이터셋에서 사전 훈련한 것보다 뚜렷하게 뛰어난 성능을 보인다.

ABSTRACT

Witnessed the development of deep learning, increasing number of studies try to build computer aided diagnosis systems for 3D volumetric medical data. However, as the annotations of 3D medical data are difficult to acquire, the number of annotated 3D medical images is often not enough to well train the deep learning networks. The self-supervised learning deeply exploiting the information of raw data is one of the potential solutions to loose the requirement of training data. In this paper, we propose a self-supervised learning framework for the volumetric medical images. A novel proxy task, i.e., Rubik's cube recovery, is formulated to pre-train 3D neural networks. The proxy task involves two operations, i.e., cube rearrangement and cube rotation, which enforce networks to learn translational and rotational invariant features from raw 3D data. Compared to the train-from-scratch strategy, fine-tuning from the pre-trained network leads to a better accuracy on various tasks, e.g., brain hemorrhage classification and brain tumor segmentation. We show that our self-supervised learning approach can substantially boost the accuracies of 3D deep learning networks on the volumetric medical datasets without using extra data. To our best knowledge, this is the first work focusing on the self-supervised learning of 3D neural networks.

연구 동기 및 목표

• 의료 진단을 위한 딥 러닝 모델 훈련에 있어 제한된 애너테이션 처리된 3D 의료 영상 문제를 해결하기 위해.
• 비정상적인 3D 부피 데이터를 활용하여 최종 작업 성능을 향상시키는 자기지도 사전 훈련 프레임워크를 개발하기 위해.
• 3D 컨볼루션 네트워크가 이동 및 회전 불변성을 학습하도록 강제하는 새로운 프록시 작업—루빅스 큐브 복구—를 제안하기 위해.
• 외부 데이터셋 없이 의료 데이터만으로 사전 훈련하는 것이 이미지넷 또는 UCF101 사전 훈련보다 의료 작업에서 더 뛰어난 성능을 낼 수 있음을 입증하기 위해.

제안 방법

• 방법은 3D 의료 부피를 3D 격자(예: 2×2×2)로 분할하여 부분 부피(큐브)로 나눈다.
• 원본 부피의 손상된 버전을 만들기 위해 큐브의 무작위 순서 뒤섞기와 무작위 회전을 적용한다.
• 프록시 작업은 큐브의 원래 구성—공간적 순서와 방향—을 재구성하는 것이다.
• 대trastive 학습 또는 분류 손실을 사용하여 3D CNN이 올바른 순서와 회전을 예측하도록 훈련한다.
• 한정된 애너테이션 데이터로 뇌 출혈 분류 및 종양 세분화와 같은 최종 작업에 대해 사전 훈련된 인코더를 미세 조정한다.
• 기준선과 비교하기 위해, 초기 학습 및 UCF101 사전 훈련된 가중치에서의 미세 조정을 포함한 프레임워크를 평가한다.

실험 결과

연구 질문

• RQ1 3D 공간 조작 기반의 자기지도 프록시 작업이 부피 의료 영상의 특징 학습을 향상시킬 수 있는가?
• RQ2 외부 데이터셋 없이 의료 데이터만으로 사전 훈련하면 자연 영상 데이터(예: UCF101)에서 사전 훈련한 것보다 의료 작업에서 더 높은 성능을 낼 수 있는가?
• RQ3 프록시 작업에 큐브의 회전을 포함시키면 큐브 순서만 고려한 경우에 비해 학습된 특징의 일반화 능력과 강건성에 어떤 영향을 미치는가?
• RQ4 루빅스 큐브 복구 작업이 분류 및 세분화와 같은 최종 3D 의료 영상 작업의 성능을 어느 정도 향상시키는가?

주요 결과

• 제안된 루빅스 큐브 복구 사전 훈련은 뇌 출혈 분류에서 83.8%의 정확도를 달성했으며, 초기 학습(72.6%)보다 11.2% 향상된 성능을 보였다.
• 뇌 종양 세분화의 경우, 3D DUC를 사용해 77.3% mIoU를 달성했으며, 큐브 순서만 고려한 기준선보다 2.3% 향상된 성능을 보였다.
• 루빅스 큐브 사전 훈련된 가중치에서의 미세 조정은 분류 작업에서 UCF101 사전 훈련된 가중치에서의 미세 조정보다 +2.7% 높은 성능을 보였고, 세분화 작업에서는 +1.1% 향상되었다.
• 프록시 작업에 큐브의 회전을 포함시킴으로써 더 강건하고, 이동 및 회전 불변성을 갖춘 특징이 학습되어 큐브 순서만 고려한 경우보다 일반화 능력이 향상되었다.
• 3D DUC 아키텍처는 모든 사전 훈련 프로토콜에서 U-Net을 능가했으며, 무작위 디코더 초기화로 인한 성능 저하가 감소함을 시사했다.
• 이 방법은 외부 데이터셋 없이도 비정상적인 의료 데이터만으로 효과적인 자기지도 사전 훈련을 수행하는 3D 컨볼루션 네트워크를 위한 첫 번째 사례이다.