[논문 리뷰] A Teacher-Student Framework for Semi-supervised Medical Image Segmentation From Mixed Supervision

Liyan Sun, Jianxiong Wu외 4명|arXiv (Cornell University)|2020. 10. 23.

Medical Imaging and Analysis참고 문헌 24인용 수 16

한 줄 요약

이 논문은 교사-학생 프레임워크를 제안하며, 부분적인 밀도 레이블과 쉽게 확보할 수 있는 루프 경계 상자(Loose bounding boxes)를 혼합된 감독 방식으로 사용하여 반감독 학습 기반의 의료 영상 분할을 수행한다. 교사 모델은 고급장치의 장기-병변 주의 메커니즘을 통해 고품질의 가짜 레이블을 생성한다. 학생 모델은 이러한 가짜 레이블과 희소한 수동 레이블, 그리고 딥 디코더 기반의 국소화 브랜치를 통합하여 LiTS 데이터셋에서 최신 기술 수준(SOTA)의 성능을 달성하며, 경계 상자 노이즈에 대해 뛰어난 내성과 더불어 일부 데이터 분할 조건에서는 전체 감독 기반 기준보다도 뛰어난 성능을 보인다.

ABSTRACT

Standard segmentation of medical images based on full-supervised convolutional networks demands accurate dense annotations. Such learning framework is built on laborious manual annotation with restrict demands for expertise, leading to insufficient high-quality labels. To overcome such limitation and exploit massive weakly labeled data, we relaxed the rigid labeling requirement and developed a semi-supervised learning framework based on a teacher-student fashion for organ and lesion segmentation with partial dense-labeled supervision and supplementary loose bounding-box supervision which are easier to acquire. Observing the geometrical relation of an organ and its inner lesions in most cases, we propose a hierarchical organ-to-lesion (O2L) attention module in a teacher segmentor to produce pseudo-labels. Then a student segmentor is trained with combinations of manual-labeled and pseudo-labeled annotations. We further proposed a localization branch realized via an aggregation of high-level features in a deep decoder to predict locations of organ and lesion, which enriches student segmentor with precise localization information. We validated each design in our model on LiTS challenge datasets by ablation study and showed its state-of-the-art performance compared with recent methods. We show our model is robust to the quality of bounding box and achieves comparable performance compared with full-supervised learning methods.

연구 동기 및 목표

• 의료 영상 분할에서 비용이 많이 드는 밀도 레이블(픽셀 단위의 정밀 레이블)에 대한 의존도를 줄이기 위해, 부분적인 밀도 레이블과 쉽게 확보할 수 있는 경계 상자 레이블을 혼합 감독 방식으로 활용하는 것.
• 장기 및 병변 분할에 있어 약한 감독 방식의 경계 상자 레이블에서 제공하는 형태 및 국소화 정보의 제한성 문제를 해결하는 것.
• 의료 영상에 특화된 반감독 학습 프레임워크를 개발하여, 병변이 장기에 존재하는 일관된 해부학적 사전 지식을 활용하는 것.
• 병변과 같은 소형 구조물의 분할 성능을 향상시키기 위해 고수준 특징을 통합하는 국소화 브랜치를 통합하는 것.

제안 방법

• 교사 네트워크 내에 계층적 장기-병변(O2L) 주의 모듈을 도입하여, 장기와 그 내부 병변 간의 공간적 관계를 모델링함으로써 가짜 레이블의 정밀도를 향상시킨다.
• 교사 모델은 O2L 주의 메커니즘을 활용해 약한 감독 데이터에 대해 가짜 분할 마스크를 생성함으로써, 형태 일관성과 국소화 정확도를 향상시킨다.
• 학생 네트워크는 수동으로 레이블링된 데이터와 교사로부터 생성된 가짜 레이블을 조합하여 학습함으로써 지식 전이(Knowledge Distillation)를 실현한다.
• 학생 모델 내에 국소화 브랜치를 도입하여 디코더에서 추출한 깊은 특징을 통합하여 객체 위치를 예측함으로써, 병변과 같은 소형 구조물의 정밀도를 향상시킨다.
• 경계 상자 레이블은 외곽점 탐지 및 외접 사각형 생성을 통해 가짜 레이블로 변환되며, 실제 레이블링 노이즈를 시뮬레이션하기 위해 무작위 변형이 적용된다.
• 강한/약한 데이터 비율의 다양한 조건에서 프레임워크의 유효성을 검증하고, 경계 상자 변형에 대한 내성 평가를 실시하며, 주요 평가 지표로 Dice 점수를 사용한다.

실험 결과

연구 질문

• RQ1 교사-학생 프레임워크가 부분적 밀도 레이블과 느슨한 경계 상자 레이블을 혼합 감독 방식으로 효과적으로 활용하여 반감독 학습 기반 의료 영상 분할을 수행할 수 있는가?
• RQ2 해부학적 사전 지식(예: 병변이 장기에 존재하는 구조)을 어떻게 모델링하여 가짜 레이블의 품질과 분할 정확도를 향상시킬 수 있는가?
• RQ3 깊은 특징 통합 기반의 국소화 브랜치가 약한 감독 조건 하에서 소형 객체 분할 성능을 얼마나 향상시키는가?
• RQ4 임상 실무에서 흔히 발견되는 노이즈가 많거나 정밀도가 떨어지는 경계 상자 레이블에 대해 모델의 내성은 어느 정도인가?
• RQ5 제한된 레이블링 데이터 조건에서 혼합 감독 방식을 사용하는 반감독 학습이 전체 감독 기반 기준보다 성능이 뛰어날 수 있는가?

주요 결과

• 경계 상자 노이즈가 존재하는 조건에서도 간 분할에서 전반적인 Dice 점수는 92.5%, 케이스별 Dice 점수는 92.2%를 기록하여 레이블링 노이즈에 대해 매우 강력한 내성을 입증하였다.
• 병변 분할의 경우 노이즈 조건에서도 전반적인 Dice 점수는 76.5%, 케이스별 Dice 점수는 59.6%를 기록하여 소형 병변임에도 불구하고 성능 저하가 최소한도로 유지됨을 보였다.
• 강한 데이터 대비 약한 데이터 비율이 70:20을 초과하는 경우, 반감독 모델이 전체 감독 기반 기준(황금 표준)을 초월하여 성능을 냄을 확인하여 고품질의 가짜 레이블이 성능 향상에 기여함을 입증함.
• 제거 실험을 통해 계층적 O2L 주의 모듈과 국소화 브랜치 모두가 분할 정확도 향상에 효과적이라는 것을 확인함.
• 다양한 데이터 분할 조건에서도 높은 성능 유지를 유지하며, 실제 경계 상자 품질의 다양성에 대해 강건한 성능을 보여, 임상 현장 적용 가능성을 뒷받침함.