[논문 리뷰] MS-KD: Multi-Organ Segmentation with Multiple Binary-Labeled Datasets

Shixiang Feng, Yuhang Zhou외 3명|arXiv (Cornell University)|2021. 08. 05.

Advanced Neural Network Applications참고 문헌 13인용 수 7

한 줄 요약

이 논문은 여러 사전 훈련된 단일 기관 분할 네트워크(교사)와 단일 다기관 학생 네트워크를 사용하여 다기관 분할 모델을 훈련하는 Multi-Teacher Single-Student Knowledge Distillation 프레임워크인 MS-KD를 제안한다. 각 교사가 자신의 기관 영역과 모든 교사가 일치하는 배경 영역에서 학생을 지도하는 영역 기반의 로짓 수준 및 특징 수준 지도를 적용함으로써, 기존 개별 모델 및 최신 기술을 초월하는 뛰어난 성능을 달성하며, 기준 데이터셋에서 평균 DSC 기준 최대 0.43% 향상된다.

ABSTRACT

Annotating multiple organs in 3D medical images is time-consuming and costly. Meanwhile, there exist many single-organ datasets with one specific organ annotated. This paper investigates how to learn a multi-organ segmentation model leveraging a set of binary-labeled datasets. A novel Multi-teacher Single-student Knowledge Distillation (MS-KD) framework is proposed, where the teacher models are pre-trained single-organ segmentation networks, and the student model is a multi-organ segmentation network. Considering that each teacher focuses on different organs, a region-based supervision method, consisting of logits-wise supervision and feature-wise supervision, is proposed. Each teacher supervises the student in two regions, the organ region where the teacher is considered as an expert and the background region where all teachers agree. Extensive experiments on three public single-organ datasets and a multi-organ dataset have demonstrated the effectiveness of the proposed MS-KD framework.

연구 동기 및 목표

• 3D 다기관 분할의 높은 주석 비용 문제를 해결하기 위해, 이진 레이블을 가진 기존의 단일 기관 데이터셋을 활용하고자 한다.
• 각각 다른 단일 기관 데이터셋에서 훈련된 작업 특화 교사 모델들로부터 지식을 전달할 수 있는 지식 정복 프레임워크를 개발하고자 한다.
• 기존 방법의 한계(조건 기반 접근 방식의 계산 비효율성 또는 저품질 허위 레이블에 대한 의존성)를 극복하기 위해, 다수의 교사로부터의 소프트 지도를 사용하고자 한다.
• 다양한 단일 기관 전문가들로부터의 지식을 영역 기반 지도를 통해 통합함으로써, 다기관 분할의 일반화 능력과 성능을 향상시키고자 한다.

제안 방법

• K개의 사전 훈련된 단일 기관 분할 네트워크를 교사로 사용하고, 하나의 다기관 분할 네트워크를 학생으로 사용한다.
• 영역 기반 지도를 적용한다: 각 교사는 자신의 기관 영역과 모든 교사가 일치하는 배경 영역에서 학생을 지도한다.
• 로짓 수준 지도는 교사의 이진 로짓을 학생의 K+1개 클래스 출력 공간에 맞추기 위해 차원 변환을 거친 후, 정복 손실을 적용한다.
• 특징 수준 지도는 분포 일치 손실을 사용하여 학생의 중간 특징 분포를 교사의 것과 일치시킨다.
• 학습은 로짓 수준 및 특징 수준 손실의 조합을 통해 엔드 투 엔드로 수행되며, 이는 다수의 작업 특화 교사로부터의 지식 전달을 가능하게 한다.
• 반복적 추론이나 동적 헤드 생성을 피함으로써, 훈련 및 추론 시 모두 낮은 계산 비용을 확보한다.

실험 결과

연구 질문

• RQ1 이진 주석을 가진 다수의 단일 기관 데이터셋만을 사용하여 단일 다기관 분할 모델을 효과적으로 훈련시킬 수 있는가?
• RQ2 작업 특화 조건 또는 반복적 추론 없이, 다수의 작업 특화 교사 모델에서 지식을 효과적으로 단일 학생 모델로 전달할 수 있는가?
• RQ3 로짓과 특징을 모두 사용하는 영역 기반 지도가 표준 허위 레이블 또는 일致성 기반 정복보다 더 높은 성능을 내는가?
• RQ4 제안된 프레임워크는 단일 기관 데이터만으로 훈련되었음에도 불구하고, 예측되지 않은 다기관 데이터셋에 잘 일반화되는가?

주요 결과

• MS-KD는 세 개의 단일 기관 데이터셋에서 훈련된 결과, 개별 단일 기관 모델 대비 평균 DSC에서 0.80% 절대 향상(91.10% 대 90.30%)을 달성한다.
• 특징 수준 지도 추가로 평균 DSC가 0.22% 향상되어 91.32%에 도달하며, 중간 특징 정복의 유용성을 입증한다.
• BTCV 다기관 벤치마크에서 MS-KD는 개별 모델 대비 평균 DSC가 1.32% 높아, 강력한 일반화 능력을 보여준다.
• 추가적인 추론 시간이나 모델 파rameter가 필요 없음에도 불구하고, Co-training(+0.43% 평균 DSC) 및 DoDNet과 같은 최신 기술을 뛰어넘는 성능을 기록한다.
• 정성적 결과에서는 MS-KD가 다수의 교사로부터의 부드러운 지도를 통해, 췌장과 같이 도전적인 기관에서 더 정확하고 부드러운 경계를 생성함을 보여준다.
• 불확실성 분석 결과, 교사들은 기관 경계에서 가장 불확실하며, MS-KD는 이러한 소프트 지도를 효과적으로 활용해 경계 검출 성능을 향상시킨다.