[논문 리뷰] Revisiting Rubik's Cube: Self-supervised Learning with Volume-wise Transformation for 3D Medical Image Segmentation

Tao Xing, Yuexiang Li외 3명|arXiv (Cornell University)|2020. 07. 17.

Advanced Neural Network Applications참고 문헌 19인용 수 6

한 줄 요약

이 논문은 3D 병변 영상 분할을 위한 새로운 비지도 학습 프레임워크인 Rubik’s cube++를 제안한다. 이 방법은 3D 바vox 회전과 같은 볼륨 단위 변환을 사용하여 3D U-Net 아키텍처의 인코더 및 디코더 구성 요소를 사전학습한다. 공간적으로 일관된 부분 결정체 순서를 활용해 맥락 복원을 사전 과제로 설정함으로써, 추가적인 레이블이 없는 상태에서도 췌장 및 뇌 조직 데이터셋에서 분할 정확도를 크게 향상시키며, 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Deep learning highly relies on the quantity of annotated data. However, the annotations for 3D volumetric medical data require experienced physicians to spend hours or even days for investigation. Self-supervised learning is a potential solution to get rid of the strong requirement of training data by deeply exploiting raw data information. In this paper, we propose a novel self-supervised learning framework for volumetric medical images. Specifically, we propose a context restoration task, i.e., Rubik's cube++, to pre-train 3D neural networks. Different from the existing context-restoration-based approaches, we adopt a volume-wise transformation for context permutation, which encourages network to better exploit the inherent 3D anatomical information of organs. Compared to the strategy of training from scratch, fine-tuning from the Rubik's cube++ pre-trained weight can achieve better performance in various tasks such as pancreas segmentation and brain tissue segmentation. The experimental results show that our self-supervised learning method can significantly improve the accuracy of 3D deep learning networks on volumetric medical datasets without the use of extra data.

연구 동기 및 목표

• 딥 러닝에서 제한된 레이블이 부여된 3D 의료 영상 데이터라는 핵심적 한계를 해결하기 위해.
• 자신의 3D 해부학적 구조를 활용한 비지도 사전 과제를 통해 딥 네트워크를 사전학습시켜 3D 의료 영상 분할 성능을 향상시키기 위해.
• 이전의 비지도 학습 방법이 인코더 레이어만 사전학습하는 데 그친다는 한계를 극복하기 위해, 맥락 복원 과제를 통해 인코더와 디코더를 함께 사전학습시키기 위해.
• 볼륨 단위 변환(예: 3D 바vox 회전)이 패치 수준 또는 격자 기반 변환보다 더 강력한 특징 표현을 제공하는지 검증하기 위해.

제안 방법

• 3D 볼륨 의료 영상을 n×n×n 크기의 부분 결정체로 격자 구조로 분할하여 루빅스 큐브 유사한 구조를 형성한다.
• 전체 결정체 층을 Sagittal, Coronal, Axial 평면을 따라 90°, 180°, 270°로 회전시켜 볼륨 단위 변환을 적용하며, 이는 공간 일관성을 유지하면서 과제의 난이도를 증가시킨다.
• 사전 과제는 맥락 복원이다: 네트워크는 변환된 3D 볼륨의 원래 공간적 구성 요소를 복원해야 하며, 이는 3D 구조적 사전 지식 학습을 장려한다.
• 이 프레임워크는 3D U-Net과 같은 완전 컨volution 네트워크(Fully Convolutional Networks, FCNs)에 적용되어 인코더 및 디코더 경로를 함께 사전학습할 수 있도록 한다.
• 변환된 볼륨에서 원래 볼륨을 정확히 재구성하도록 유도하는 대비 손실(contrastive loss)을 사용해 엔드 투 엔드로 학습한다.
• 공개 데이터셋인 Pancreas-101 및 MRBrainS18을 사용해 교차 검증을 수행하며, 주요 평가 지표로는 Dice Similarity Coefficient(DSC)를 사용한다.

실험 결과

연구 질문

• RQ1 맥락 복원 과제에서 볼륨 단위 변환은 패치 수준 또는 격자 기반 변환보다 3D 해부학적 특징 학습에 더 나은 성능을 보이는가?
• RQ2 맥락 복원 과제를 통해 인코더와 디코더를 함께 사전학습시키는 것이 인코더만 사전학습시키는 것보다 3D 의료 영상 분할 성능을 향상시키는가?
• RQ3 하나의 부분 결정체 크기 n으로 제어되는 사전 과제의 난이도가 학습된 표현의 질과 후속 분할 정확도에 어떤 영향을 미치는가?
• RQ4 제안된 Rubik’s cube++ 방법은 기존의 비지도 기반 기준 모델들인 Models Genesis 및 Rubik’s cube [20]보다 3D 의료 영상 분할 과제에서 통계적으로 유의미하게 뛰어난 성능을 보이는가?

주요 결과

• 뇌 조직 분할에서 Rubik’s cube++는 초기 학습(From Scratch) 대비 평균 DSC를 5.34% 향상시켰으며, UCF101 사전학습 및 기타 비지도 기반 기준 모델들을 모두 능가했다.
• 췌장 분할 과제에서 100% 훈련 데이터를 사용할 경우 DSC가 84.08%를 기록하여, 이어지는 최고의 기준 모델(83.72%)을 뛰어넘었으며, 모든 데이터 설정에서 일관된 성능 향상을 보였다.
• t-검정 결과, Models Genesis 대비 성능 향상은 유의미하게 5% 수준에서 유의미하다(p-value = 3.42%), 이는 방법의 강건성을 검증한다.
• 최적의 부분 결정체 크기 n=7에서 과제 난이도와 특징 학습 간의 균형이 가장 우수했으며, 50% 훈련 데이터 시 DSC는 82.80%, 100% 훈련 데이터 시 84.08%로 최고 성능 기록했다.
• n을 9로 증가시킬 경우 성능 저하가 발생하여, 3D 구조가 과도하게 손상되어 해부학적 정보 추출에 악영향을 주고 표현의 강건성이 감소함을 시사한다.
• 이 방법은 맥락 복원 과제를 통해 인코더와 디코더를 함께 사전학습시키는 것이, 인코더만 사전학습하는 방법과는 달리 더 뛰어난 분할 성능을 이끌어낸다는 점을 입증한다.