[논문 리뷰] Co-training with High-Confidence Pseudo Labels for Semi-supervised Medical Image Segmentation

Zhiqiang Shen, Peng Cao외 4명|arXiv (Cornell University)|2023. 01. 11.

Advanced Neural Network Applications인용 수 7

한 줄 요약

이 논문은 제한된 레이블이 부여된 데이터에서 상태의 성능을 달성하기 위해 불확실성 유도 영역 혼합(UMIX)을 통해 고신뢰도 가짜 레이블을 생성하고, 협동형 평균티처(CMT) 아키텍처를 통해 모델 간 이견을 유지함으로써 공훈련(co-training)을 향상시키는, 불확실성 유도 협동형 평균티처(UCMT)라는 반감독 의료 영상 분할 프레임워크를 제안한다. UCMT는 네 가지 공개 2D 및 3D 의료 영상 데이터셋에서 상태의 성능를 기록한다.

ABSTRACT

Consistency regularization and pseudo labeling-based semi-supervised methods perform co-training using the pseudo labels from multi-view inputs. However, such co-training models tend to converge early to a consensus, degenerating to the self-training ones, and produce low-confidence pseudo labels from the perturbed inputs during training. To address these issues, we propose an Uncertainty-guided Collaborative Mean-Teacher (UCMT) for semi-supervised semantic segmentation with the high-confidence pseudo labels. Concretely, UCMT consists of two main components: 1) collaborative mean-teacher (CMT) for encouraging model disagreement and performing co-training between the sub-networks, and 2) uncertainty-guided region mix (UMIX) for manipulating the input images according to the uncertainty maps of CMT and facilitating CMT to produce high-confidence pseudo labels. Combining the strengths of UMIX with CMT, UCMT can retain model disagreement and enhance the quality of pseudo labels for the co-training segmentation. Extensive experiments on four public medical image datasets including 2D and 3D modalities demonstrate the superiority of UCMT over the state-of-the-art. Code is available at: https://github.com/Senyh/UCMT.

연구 동기 및 목표

• 공훈련 기반 반감독 의료 영상 분할에서 조기 수렴하여 낮은 신뢰도의 가짜 레이블을 생성하는 모델의 한계를 해결하기 위해.
• 학습 중 불확실성을 줄이면서도 하위 네트워크 간 모델 이견을 유지함으로써 가짜 레이블의 품질을 향상시키기 위해.
• 제한된 레이블 데이터에서 일반화 능력을 향상시키기 위해 모델 이견과 가짜 레이블 신뢰도를 동시에 향상시키는 통합 프레임워크를 개발하기 위해.
• 공훈련 기반 반감독 학습에서 모델 이견과 가짜 레이블 불확실성 간의 상호작용을 조사하기 위해.

제안 방법

• UCMT는 공훈련을 두 개의 학생 네트워크 간에 강제하기 위해, 모멘터리 평균화된 교사 모델을 사용하여 훈련을 정규화하고 공 consensus 붕괴를 방지하는 협동형 평균티처(CMT) 구성요소를 도입한다.
• 불확실성 유도 영역 혼합(UMIX) 모듈은 CMT 예측에서 고불확실성 영역을 선택적으로 혼합하여 고신뢰도 가짜 레이블 생성을 촉진한다.
• UMIX는 CMT의 불확실성 맵을 사용해 모호한 영역을 식별하고 증강하며, 낮은 불확실성 영역만 유지하고 고불확실성 영역은 능동적으로 개선한다.
• 이 프레임워크는 CMT와 UMIX를 공동 훈련 체계에 통합하여, CMT가 UMIX로 증강된 입력에서 생성된 가짜 레이블을 활용해 예측의 강건성과 다양성을 향상시킨다.
• 이 방법은 학생 네트워크 간 교차 가짜 레이블 감독과 평균티처 정규화를 활용하여 훈련 안정성과 일반화 능력을 향상시킨다.
• UMIX의 하이퍼파라미터, 예를 들어 상위-k 영역 수와 패치 크기 등이 안정성 분석을 통해 다양한 설정에서 안정적인 성능을 보임을 확인하였다.

실험 결과

연구 질문

• RQ1 공훈련 기반 반감독 의료 영상 분할에서 모델 이견과 가짜 레이블 신뢰도를 동시에 최적화할 수 있는 방법은 무엇인가?
• RQ2 불확실성 인식 데이터 증강이 가짜 레이블 품질과 모델 성능에 미치는 영향은 무엇인가?
• RQ3 모멘터리 정규화된 교사 모델을 갖춘 협동형 평균티처 프레임워크는 자가학습 행동으로의 조기 수렴을 방지할 수 있는가?
• RQ4 낮은 데이터 환경에서 불확실성 유도 영역 혼합은 CutMix와 같은 무작위 증강 방법보다 어떻게 비교되는가?
• RQ5 공훈련 환경에서 모델 이견, 가짜 레이블 불확실성, 분할 성능 간의 관계는 어떠한가?

주요 결과

• UCMT는 다양한 2D 및 3D 모odalities를 갖춘 네 가지 공개 의료 영상 분할 데이터셋(ISIC, LA, ACDC, Polyp)에서 상태의 성능를 기록한다.
• 5%의 레이블이 부여된 ISIC 데이터셋에서 UCMT는 딱시 스코어 88.22를 기록하여 베이스라인 및 이전 SOTA 방법을 모두 초월한다.
• 제거 실험 결과, CMT와 UMIX를 결합한 것이 최고의 성능를 기록하며, 전체 UCMT 모델이 각각의 구성요소가 없는 변형보다 유의미하게 뛰어난 성능를 보였다.
• UMIX는 낮은 데이터 환경(예: 2% 레이블)에서 CutMix를 항상 능가하며, 고신뢰도 가짜 레이블 생성의 효과성을 입증한다.
• 하이퍼파라미터 민감도 분석 결과, UMIX는 패치 크기 및 상위-k 영역 선택에 관계없이 성능 저하가 최소한이었으며, 매우 강건함을 확인하였다.
• 정성적 결과는 UCMT가 평균티처 및 교차 가짜 레이블 감독 베이스라인에 비해 더 매끄럽고 정확한 분할 경계를 생성하며 미세한 세부 사항을 더 잘 복원함을 보여준다.