[논문 리뷰] Bootstrap Representation Learning for Segmentation on Medical Volumes and Sequences

Zejian Chen, Wei Zhuo외 3명|arXiv (Cornell University)|2021. 01. 01.

AI in cancer detection참고 문헌 55인용 수 3

한 줄 요약

이 논문은 제한된 애너테이션을 사용하여 의료 볼륨 및 시퀀스 분할을 위한 부트스트랩 자기지도 학습 방법을 제안한다. 슬라이스 간 공간 연속성과 주의력 유도 예측자 및 배경-전경 프로토타입 校정을 갖춘 전역-국소 예측 프레임워크를 활용하여, 이 방법은 기존 방법보다 ACDC에서 4.5% DSC, 전립선에서 1.7%, CAMUS에서 2.3% 향상된 최신 기술 수준의 성능을 달성한다. 라벨이 부족한 볼륨 몇 개만으로도 성능을 확보한다.

ABSTRACT

In this work, we propose a novel straightforward method for medical volume and sequence segmentation with limited annotations. To avert laborious annotating, the recent success of self-supervised learning(SSL) motivates the pre-training on unlabeled data. Despite its success, it is still challenging to adapt typical SSL methods to volume/sequence segmentation, due to their lack of mining on local semantic discrimination and rare exploitation on volume and sequence structures. Based on the continuity between slices/frames and the common spatial layout of organs across volumes/sequences, we introduced a novel bootstrap self-supervised representation learning method by leveraging the predictable possibility of neighboring slices. At the core of our method is a simple and straightforward dense self-supervision on the predictions of local representations and a strategy of predicting locals based on global context, which enables stable and reliable supervision for both global and local representation mining among volumes. Specifically, we first proposed an asymmetric network with an attention-guided predictor to enforce distance-specific prediction and supervision on slices within and across volumes/sequences. Secondly, we introduced a novel prototype-based foreground-background calibration module to enhance representation consistency. The two parts are trained jointly on labeled and unlabeled data. When evaluated on three benchmark datasets of medical volumes and sequences, our model outperforms existing methods with a large margin of 4.5\% DSC on ACDC, 1.7\% on Prostate, and 2.3\% on CAMUS. Intensive evaluations reveals the effectiveness and superiority of our method.

연구 동기 및 목표

• 딥 러닝 기반 분할에서 제한된 애너테이션을 가진 의료 볼륨 및 시퀀스의 과제를 해결하기 위해.
• 3D 및 시간적 의료 데이터의 표현 학습을 향상시키기 위해 내재된 공간적 및 시간적 연속성을 활용하기 위해.
• 두 단계 학습 접근 방식과 달리, 라벨이 없는 데이터의 기억 상실을 방지하는 공동 학습 프레임워크를 개발하기 위해.
• 예측 가능한 슬라이스 전이 기반의 조밀한 자기지도 학습을 통해 국소 및 전역 표현 학습을 향상시키기 위해.
• 라벨이 없는 데이터의 의미 기반 프로토타입 기반 전경-배경 校정을 도입하여 대비 학습의 안정성을 제고하기 위해.

제안 방법

• 현재 슬라이스의 전역적 맥락을 활용하여 다른 슬라이스의 조밀한 특징을 예측하는 비대칭 인코더-예측자 네트워크를 도입하며, 거리별로 특화된 주의 메커니즘을 적용한다.
• 볼륨 간 및 볼륨 내에서 예측에 대한 자기지도 학습을 강화하기 위해 조밀한 유사도 손실을 적용하여, 공동으로 전역 및 국소 표현 탐색을 가능하게 한다.
• 예측의 가중치를 슬라이스 간 거리(내부 및 외부 볼륨)에 따라 적응적으로 조정하기 위해 예측자 내에서 거리 보정 주의 메커니즘을 설계한다.
• 라벨이 없는 디코더 특징을 라벨이 있는 데이터의 프로토타입과 일치시키는 프로토타입 기반 전경-배경 校정 모듈을 제안하여 특징 일관성을 향상시킨다.
• 라벨이 있는 데이터의 전경 마스크와 라벨이 없는 데이터의 확률 맵에서 고활성 영역을 의미론적 기반으로 대비 학습에 활용한다.
• 라벨이 있는 데이터와 라벨이 없는 데이터를 모두 사용하여 전체 모델을 엔드 투 엔드로 학습하며, 예측 및 校정 목표를 공동으로 최적화한다.

실험 결과

연구 질문

• RQ1 의료 볼륨에서 인접 슬라이스 간 예측 가능한 공간 연속성이 자기지도 표현 학습에 효과적으로 활용될 수 있는가?
• RQ2 전역-국소 예측 프레임워크는 3D 및 순차적 의료 데이터에서 국소 및 전역 특징 학습을 어떻게 향상시킬 수 있는가?
• RQ3 프로토타입 기반 전경-배경 校정은 준지도 학습 분할에서 특징 일관성과 일반화 능력을 향상시키는가?
• RQ4 제안된 방법은 볼륨 및 순차적 데이터에서 랜덤 크롭 또는 패치 수준 증강에 의존하는 대비 학습 기반 자기지도 학습 방법을 초월할 수 있는가?
• RQ5 이 방법은 다양한 백본 아키텍처와 다양한 애너테이션 예산을 가진 데이터셋에서 어떻게 일반화되는가?

주요 결과

• 제안된 방법은 ACDC 데이터셋에서 라벨이 2개 뿐인 볼륨만으로도 이전 최신 기술 수준(CGL)보다 4.5% 높은 DSC를 달성한다.
• 전립선 데이터셋에서 라벨이 2개 뿐인 볼륨을 사용할 때, CGL 대비 평균 DSC가 1.7% 향상되어 도전적인 해부학적 구조에서 뛰어난 성능을 입증한다.
• CAMUS 에코카디오그래피 데이터셋에서 라벨이 8개의 시퀀스만으로도 CGL 대비 2.3% 높은 DSC를 달성하며, 전체 라벨링 기준 대비 89.3% 적은 애너테이션을 사용한다.
• 다양한 인코더 백본 아키텍처에서 잘 일반화되며, VGG13 및 ResNet18 아키텍처에서 모두 무작위 초기화 및 CGL보다 뛰어난 성능을 보인다.
• 전립선 데이터셋에서 라벨이 8개 뿐인 제안 방법과 전체 라벨링 기준(20개 볼륨) 간 DSC 성능 격차는 뿐만 아니라 0.7%에 불과하여 높은 데이터 효율성을 보여준다.
• 절단 실험 결과, 주의력 유도 예측자와 전경-배경 校정 모듈이 최적 성능을 내기 위해 필수적임을 확인하였으며, 각 구성 요소가 최종 성능 향상에 크게 기여한다.