[논문 리뷰] Semi-supervised Medical Image Segmentation through Dual-task Consistency

Xiangde Luo, Jieneng Chen외 4명|arXiv (Cornell University)|2020. 09. 09.

Advanced Neural Network Applications인용 수 29

한 줄 요약

이 논문은 이중작업 일致성 기반의 새로운 반감독 의료 영상 분할 프레임워크를 제안한다. 깊이 신경망이 픽셀 단위 분할 맵과 기하학적 인식 수준 집합 함수를 동시에 예측한다. 미분 가능한 작업 변환 레이어를 통해 수준 집합 표현을 분할 맵으로 변환하며, 두 출력 간의 일치성 정규화는 레이블이 없는 데이터를 활용해 성능을 햖저하여, 두 개의 3D 의료 영상 데이터셋에서 최신 기술 수준의 성능을 달성하면서 계산 비용도 감소시켰다.

ABSTRACT

Deep learning-based semi-supervised learning (SSL) algorithms have led to promising results in medical images segmentation and can alleviate doctors' expensive annotations by leveraging unlabeled data. However, most of the existing SSL algorithms in literature tend to regularize the model training by perturbing networks and/or data. Observing that multi/dual-task learning attends to various levels of information which have inherent prediction perturbation, we ask the question in this work: can we explicitly build task-level regularization rather than implicitly constructing networks- and/or data-level perturbation-and-transformation for SSL? To answer this question, we propose a novel dual-task-consistency semi-supervised framework for the first time. Concretely, we use a dual-task deep network that jointly predicts a pixel-wise segmentation map and a geometry-aware level set representation of the target. The level set representation is converted to an approximated segmentation map through a differentiable task transform layer. Simultaneously, we introduce a dual-task consistency regularization between the level set-derived segmentation maps and directly predicted segmentation maps for both labeled and unlabeled data. Extensive experiments on two public datasets show that our method can largely improve the performance by incorporating the unlabeled data. Meanwhile, our framework outperforms the state-of-the-art semi-supervised medical image segmentation methods. Code is available at: https://github.com/Luoxd1996/DTC

연구 동기 및 목표

• 의료 영상 분할에서 높은 표기 비용 문제를 해결하기 위해 레이블이 없는 데이터를 효과적으로 활용하는 것.
• 반감독 학습에서 전통적인 데이터 또는 네트워크 수준의 변형 기반 정규화보다 작업 수준 정규화가 더 우수한 성능을 낼 수 있는지 조사하는 것.
• 수준 집합 표현을 통해 기하학적 및 전반적인 형태 제약 조건을 반감독 분할에 통합하는 것.
• 복잡한 아키텍처나 반복적 추론을 피하는 단순하면서도 효과적인 프레임워크를 개발하는 것.
• 기존 방법들과 비교해 더 낮은 계산 비용으로 최신 기술 수준의 성능을 달성하는 것.

제안 방법

• 픽셀 단위 분할 맵과 분할 윤곽을 나타내는 전역 수준의 수준 집합 함수(LSF)를 동시에 예측하는 이중작업 딥 네트워크를 설계한다.
• 부드러운 헤비사이드 레이어를 사용해 수준 집합 함수를 분할 확률 맵으로 미분 가능하게 변환함으로써 엔드 투 엔드 학습을 가능하게 한다.
• 레이블이 있는 데이터와 없는 데이터 모두에 대해 직접 예측된 분할 맵과 수준 집합 함수에서 유도된 맵 간의 차이를 최소화하는 이중작업 일치 손실을 도입한다.
• 이 일치 손실은 두 개의 서로 다른 작업에서 유도된 예측 간의 일치를 강제하는 정규화 항으로 기능한다.
• 이 프레임워크는 이미지당 단일 프로퍼게이션만을 사용하여 반복적 개선이나 다중 추론 단계를 피한다.
• 이 방법은 다중 클래스 분할으로 확장 가능하며, 에지 검출이나 관건점 추정과 같은 추가적인 미분 가능한 작업을 통합할 수 있다.

실험 결과

연구 질문

• RQ1 반감독 의료 영상 분할에서 작업 수준의 일치 정규화가 데이터 수준 또는 네트워크 수준의 변형 기반 정규화보다 더 뛰어난 성능을 낼 수 있는가?
• RQ2 작업 수준의 정규화가 수준 집합 표현을 통해 기하학적 인식 특성을 통합함으로써 반감독 설정에서 분할 성능을 향상시킬 수 있는가?
• RQ3 다른 작업 브랜치(픽셀 단위 및 전역 수준)에서 유도된 예측 간의 일치를 강제하면 레이블이 없는 데이터에 대해 더 나은 일반화와 강건성을 확보할 수 있는가?
• RQ4 단순한 단일 프로퍼게이션 프레임워크가 복잡한 아키텍처나 학습 방식 없이 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ5 기존 최신 기술 수준의 반감독 분할 방법들과 비교했을 때, 제안된 방법은 정확도, 파rameter 효율성, 추론 시간 측면에서 어떻게 비교되는가?

주요 결과

• 제안된 방법은 3D 좌심방 MRI 데이터셋과 췌장 CT 데이터셋에서 모두 최신 기술 수준의 성능을 달성하였으며, 모든 평가 지표에서 기존 반감독 방법들을 능가하였다.
• 좌심방 데이터셋에서 이 방법은 DICE, ASD, 95HD 측면에서 최고 성능을 기록하였으며, 다음으로 우수한 방법보다 유의미하게 높은 DICE 스코어를 확보하였다.
• 췌장 데이터셋에서는 모든 비교 방법들 중에서 가장 높은 DICE 스코어와 가장 낮은 95퍼센트 백분위수 하우스도르프 거리를 기록하여 경계 정확도가 뛰어나다는 것을 보여주었다.
• SASSNet, UAMT, CCT와 같은 방법들보다 훈련 시간과 파라미터 수가 적지만 더 뛰어난 성능을 달성하였다.
• 시각화 결과는 기준 방법들에 비해 진짜 값과의 겹침이 더 높고, 거짓 양성 수가 적으며 해부학적 세부 정보를 더 잘 유지하고 있음을 보여주었다.
• 제거 실험 결과 이중작업 일치성이 성능 향상의 핵심 요소임을 확인하였으며, 다양한 해부학적 구조와 영상 모odalities에 걸쳐 잘 일반화됨을 입증하였다.