[논문 리뷰] Translation Consistent Semi-supervised Segmentation for 3D Medical Images

Yuyuan Liu, Yu Tian외 5명|arXiv (Cornell University)|2022. 03. 28.

Advanced Neural Network Applications인용 수 7

한 줄 요약

이 논문은 레이블이 없는 데이터의 이동된 시각 간 일관성을 강제하여 공간적 맥락에 대한 의존도를 줄이고 전경 객체의 일반화 성능을 향상시키는 새로운 반감독 3D 의료 영상 분할 방법인 TraCoCo를 제안한다. 높은 수준의 수렴 성능을 보장하기 위해 Confident Regional Cross-Entropy (CRC) 손실을 도입하였으며, 3D CutMix 증강 기법을 확장하여, 최소한의 레이블 데이터로도 LA, Pancreas, BRaTS19 데이터셋에서 최신 기준(SOTA) 성능을 달성한다.

ABSTRACT

3D medical image segmentation methods have been successful, but their dependence on large amounts of voxel-level annotated data is a disadvantage that needs to be addressed given the high cost to obtain such annotation. Semi-supervised learning (SSL) solve this issue by training models with a large unlabelled and a small labelled dataset. The most successful SSL approaches are based on consistency learning that minimises the distance between model responses obtained from perturbed views of the unlabelled data. These perturbations usually keep the spatial input context between views fairly consistent, which may cause the model to learn segmentation patterns from the spatial input contexts instead of the segmented objects. In this paper, we introduce the Translation Consistent Co-training (TraCoCo) which is a consistency learning SSL method that perturbs the input data views by varying their spatial input context, allowing the model to learn segmentation patterns from visual objects. Furthermore, we propose the replacement of the commonly used mean squared error (MSE) semi-supervised loss by a new Cross-model confident Binary Cross entropy (CBC) loss, which improves training convergence and keeps the robustness to co-training pseudo-labelling mistakes. We also extend CutMix augmentation to 3D SSL to further improve generalisation. Our TraCoCo shows state-of-the-art results for the Left Atrium (LA) and Brain Tumor Segmentation (BRaTS19) datasets with different backbones. Our code is available at https://github.com/yyliu01/TraCoCo.

연구 동기 및 목표

• 기존의 반감독 학습(SSL) 방법이 3D 의료 영상에서 공간적 입력 맥락에 지나치게 의존하는 데서 비롯되는 한계를 해결하기 위해.
• 공간적으로 이동된 시각 간의 공훈련 프레임워크를 도입하여 확신된 예측의 일관성을 유지함으로써 확인 편향과 국소 최적점 수렴 문제를 줄이기 위해.
• 학생 모델이 높은 신뢰도로 예측한 긍정 및 부정 예측에 집중하는 새로운 손실 함수를 통해 훈련 수렴성과 정확도를 향상시키기 위해.
• 3D CutMix 증강 기법을 반감독 3D 분할에 확장하여 일반화 성능을 향상시키기 위해.
• 제한된 애너테이션 데이터로도 주요 3D 의료 분할 벤치마크에서 최신 기준(SOTA) 성능을 달성하기 위해.

제안 방법

• TraCoCo는 동일한 레이블이 없는 3D 볼륨의 공간적으로 이동된 시각 간 일관성을 강제하는 공훈련 전략을 사용하며, 두 개의 독립적으로 초기화된 모델이 훈련된다.
• 이동된 시각 간의 일관성은 이동된 시각의 겹치는 영역 내에서만 세그멘테이션 일치도를 계산함으로써 유지되며, 배경 패atters에 대한 의존도를 감소시킨다.
• 새로운 Confident Regional Cross-Entropy (CRC) 손실이 도입되었으며, 이는 학생 모델이 높은 신뢰도로 예측한 예측값(긍정 및 부정)에 집중하여 수렴성과 정확도를 향상시킨다.
• CRC 손실은 높은 신뢰도로 예측된 영역에만 적용되므로, 노이즈가 있거나 잘못된 가짜 레이블의 영향을 최소화한다.
• 3D CutMix 증강 기법은 훈련 중 서로 다른 볼륨의 패치를 혼합함으로써 3D SSL에 확장되었으며, 일반화 성능과 데이터 효율성을 향상시킨다.
• 연결된 성분 후처리 기법을 적용하여 입력 볼륨의 1/1500 이하인 고립된 거짓 양성 영역을 걸러내어 세그멘테이션 정확도를 추가로 향상시켰다.

실험 결과

연구 질문

• RQ1 레이블이 없는 3D 의료 영상의 공간적으로 이동된 시각 간 일관성을 강제하면 배경 패턴에 대한 모델 의존도를 줄이고 전경 객체로의 일반화 성능을 향상시킬 수 있는가?
• RQ2 높은 신뢰도의 가짜 레이블에 집중하는 신뢰 기반 손실 함수가 기존의 MSE 기반 손실 함수보다 훈련 수렴성과 정확도를 향상시키는가?
• RQ3 3D CutMix 증강 기법이 제한된 애너테이션 데이터로도 반감독 3D 의료 영상 분할에서 일반화 성능을 향상시킬 수 있는가?
• RQ4 제안된 TraCoCo 프레임워크는 다양한 3D 의료 영상 작업에서 최신 기준(SOTA) 방법들과 비교해 세그멘테이션 정확도 측면에서 어떻게 성능을 내는가?
• RQ5 연결된 성분 필터링과 같은 후처리 기법이 모델의 최종 세그멘테이션 성능에 미치는 영향은 무엇인가?

주요 결과

• TraCoCo는 모든 평가 지표(Dice, mIoU, 95HD)에서 Left Atrium (LA), Pancreas-CT (Pancreas), Brain Tumor Segmentation (BRaTS19) 데이터셋에서 최신 기준(SOTA) 성능을 달성하였다.
• Confident Regional Cross-Entropy (CRC) 손실은 전경 및 배경 볼륨에 대해 더 좁은 신뢰도 클러스터링을 만들어내어, 더 높은 세그멘테이션 정확도와 더 빠른 수렴을 이끌어냈다.
• Pancreas-CT 데이터셋에서 CRC 손실은 MSE 손실 대비 1.5% 향상된 Dice 점수와 2.1% 향상된 mIoU 점수를 기록하였다.
• 입력 볼륨의 1/1500 비율 이하인 영역을 기준으로 연결된 성분 후처리를 적용하면 거짓 양성 영역이 감소하고 모든 지표가 향상되며, CRC 손실과 조합했을 때 가장 우수한 성능을 기록하였다.
• 3D CutMix 증강 기법의 추가로 일반화 성능이 향상되었으며, 특히 데이터가 적은 환경에서 일관된 성능 향상이 관찰되었다.
• TraCoCo는 평균적으로 볼륨당 1.386초(후처리 포함 시 1.412초)에 실행되며, BraTS19에서 후처리에 추가로 0.098초의 오버헤드만 발생시켜 실용적인 효율성을 입증하였다.