[논문 리뷰] Correlation-Aware Mutual Learning for Semi-supervised Medical Image Segmentation

Shengbo Gao, Ziji Zhang외 3명|arXiv (Cornell University)|2023. 07. 12.

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 레이블이 부여된 데이터를 통해 비라벨 데이터의 특징 학습을 안내함으로써 모델 성능을 햖थ하는 새로운 반감성의 의료 영상 분할 프레임워크인 상관관계 인식 상호학습(CAML)을 제안한다. CAML은 두 모듈—표본 간 상호주의(Cross-sample Mutual Attention, CMA)는 조밀한 표본 간 상호주의를, 옴니-상관 일관성(Omni-Correlation Consistency, OCC)은 레이블이 부여된 및 비레이블 특징 간 일관성 정규화를 담당하는 공-teaching 기반의 구조를 채택하여, 특히 레이블 수가 적은 조건에서 최신 기술 수준(SOTA)의 성능을 달성한다. 5%의 레이블 데이터에서 87.34%의 Dice 스코어를 기록하고, 20%의 레이블 데이터에서 완전히 레이블이 부여된 모델 성능에 근접한다.

ABSTRACT

Semi-supervised learning has become increasingly popular in medical image segmentation due to its ability to leverage large amounts of unlabeled data to extract additional information. However, most existing semi-supervised segmentation methods only focus on extracting information from unlabeled data, disregarding the potential of labeled data to further improve the performance of the model. In this paper, we propose a novel Correlation Aware Mutual Learning (CAML) framework that leverages labeled data to guide the extraction of information from unlabeled data. Our approach is based on a mutual learning strategy that incorporates two modules: the Cross-sample Mutual Attention Module (CMA) and the Omni-Correlation Consistency Module (OCC). The CMA module establishes dense cross-sample correlations among a group of samples, enabling the transfer of label prior knowledge to unlabeled data. The OCC module constructs omni-correlations between the unlabeled and labeled datasets and regularizes dual models by constraining the omni-correlation matrix of each sub-model to be consistent. Experiments on the Atrial Segmentation Challenge dataset demonstrate that our proposed approach outperforms state-of-the-art methods, highlighting the effectiveness of our framework in medical image segmentation tasks. The codes, pre-trained weights, and data are publicly available.

연구 동기 및 목표

• 기존 반감성 의료 영상 분할 방법에서 레이블이 부여된 데이터가 초기 감독 신호로만 사용되는 문제점을 해결하기 위해.
• 레이블이 부여된 데이터와 비레이블 데이터 간의 표본 간 상관관계를 명시적으로 모델링하여 사전 지식을 전달함으로써 모델 일반화 능력을 향상시키기 위해.
• 레이블이 부여된 데이터와 비레이블 데이터의 공동 최적화를 통해 특징 표현을 향상시키는 상호학습 프레임워크를 개발하기 위해.
• 소수의 레이블이 부여된 영상만 존재하는 상황에서도 강력하고 일관된 성능을 달성하기 위해.
• 임상 응용에서 대규모 레이블이 부여된 데이터셋에 대한 의존도를 줄이기 위해, 제한된 레이블 데이터에서 최대한의 정보를 추출하는 데 기여하기 위해.

제안 방법

• CAML은 상호 허위 감독 신호를 통해 훈련되는 두 개의 학생 네트워크를 활용하는 공-teaching 프레임워크를 채택하여 일관성과 지식 정착을 장려한다.
• 표본 간 상호주의(Cross-sample Mutual Attention, CMA) 모듈은 그룹 내 표본 간 조밀하고 이방향의 주의를 설정하여, 레이블이 부여된 및 비레이블 데이터 간 공유된 주목 특징 학습을 가능하게 한다.
• 옴니-상관 일관성(Omni-Correlation Consistency, OCC) 모듈은 레이블이 부여된 특징의 메모리 백업을 구축하고, 비레이블 특징과 이 기준 집합 간의 옴니-상관을 계산하여 특징 수준의 유사도를 모델링한다.
• OCC는 두 하위 모델의 옴니-상관 행렬 간 일관성을 강제로 유지하여 학습 과정을 정규화하고 비레이블 데이터에서 통합된 특징 분포를 촉진한다.
• 이 프레임워크는 상호학습 루프 내에서 두 모듈을 통합하여, 레이블이 부여된 데이터에서 비레이블 표본으로 지식을 전파하면서 양 네트워크를 동시에 업데이트한다.
• 이 방법은 엔드 투 엔드로 훈련 가능하며, 추가적인 추론 비용이 발생하지 않아 임상 적용에 적합하다.

실험 결과

연구 질문

• RQ1 반감성 분할에서 비레이블 의료 영상의 표현 학습을 향상시키기 위해 레이블이 부여된 데이터를 효과적으로 활용할 수 있는가?
• RQ2 레이블이 부여된 데이터와 비레이블 데이터 간의 표본 간 상관관계를 어떻게 모델링할 수 있는가? 이는 특징 전달 및 일관성을 향상시키는 데 기여하는가?
• RQ3 표본 간 다중 상호주의를 도입할 경우, 기존의 표준 일관성 정규화 대비 분할 성능 향상이 이루어지는가?
• RQ4 레이블이 부여된 특징과 비레이블 특징 간 옴니-상관 일관성은 더 강력하고 일반화 능력이 뛰어난 특징 학습을 이끌 수 있는가?
• RQ5 제안된 CAML 프레임워크는 5% 또는 10%의 레이블 데이터와 같은 극도로 레이블이 부족한 상황에서 어떻게 성능을 발휘하는가?

주요 결과

• 5%의 레이블 데이터에서 CAML은 표준 편차 0.05를 동반한 87.34%의 Dice 스코어를 기록하며, 베이스라인을 크게 앞서며 높은 안정성을 보였다.
• 10%의 레이블 데이터에서 CAML은 89.62%의 Dice 스코어를 기록하여 완전히 레이블이 부여된 모델의 상한선(90.98%)에 근접하였다.
• 20%의 레이블 데이터에서 CAML은 90.78%의 Dice 스코어를 기록하여, 완전히 레이블이 부여된 상한선과의 격차가 극히 미미하였다.
• 제거 실험을 통해 CMA 및 OCC 모듈이 모두 필수적임을 확인하였으며, 이들의 조합은 5% 레이블 데이터에서 6.42%의 절대적인 Dice 스코어 향상을 이끌었다.
• 실험 전반에서 낮은 표준 편차를 유지하여, 특히 레이블 수가 적은 환경에서 높은 강건성과 일반화 능력을 보였다.
• 정성적 결과에서는 CAML이 모든 비교 방법보다 더 정밀하고 완전한 분할 마스크를 생성하였으며, 경계 및 소형 구조 영역에서 두드러진 성능 향상을 보였다.