[논문 리뷰] Boosting Semi-supervised Image Segmentation with Global and Local Mutual Information Regularization

Jizong Peng, Marco Pedersoli외 1명|arXiv (Cornell University)|2021. 03. 08.

Advanced Neural Network Applications참고 문헌 59인용 수 22

한 줄 요약

이 논문은 전역 및 국소 특징 표현에 대한 상호정보(MI) 정규화를 사용하여 성능을 햖थ시키는 준지도 학습 이미지 분할 방법을 제안한다. 같은 이미지의 변환된 뷰들 간의 클러스터링된 특징 임베딩 간의 MI를 최대화하고 네트워크 단계 간의 일관성을 강제함으로써, 의료 영상 기준 평가에서 레이블이 5%인 경우에도 준지도 학습 성능에 근접한 성능을 달성한다.

ABSTRACT

The scarcity of labeled data often impedes the application of deep learning to the segmentation of medical images. Semi-supervised learning seeks to overcome this limitation by exploiting unlabeled examples in the learning process. In this paper, we present a novel semi-supervised segmentation method that leverages mutual information (MI) on categorical distributions to achieve both global representation invariance and local smoothness. In this method, we maximize the MI for intermediate feature embeddings that are taken from both the encoder and decoder of a segmentation network. We first propose a global MI loss constraining the encoder to learn an image representation that is invariant to geometric transformations. Instead of resorting to computationally-expensive techniques for estimating the MI on continuous feature embeddings, we use projection heads to map them to a discrete cluster assignment where MI can be computed efficiently. Our method also includes a local MI loss to promote spatial consistency in the feature maps of the decoder and provide a smoother segmentation. Since mutual information does not require a strict ordering of clusters in two different assignments, we incorporate a final consistency regularization loss on the output which helps align the cluster labels throughout the network. We evaluate the method on four challenging publicly-available datasets for medical image segmentation. Experimental results show our method to outperform recently-proposed approaches for semi-supervised segmentation and provide an accuracy near to full supervision while training with very few annotated images.

연구 동기 및 목표

• 딥 러닝 기반의 의미 분할에서 annotation이 제한된 의료 영상 문제를 해결하기 위해.
• 표현 불변성과 공간 일관성에 기반하여 무 supervision 데이터를 활용함으로써 준지도 학습을 향상시키기 위해.
• 큰 배치 크기가 필요하거나 고차원 연속 특징에서 어려움을 겪는 기존의 일관성 기반 및 대비 학습 방법의 한계를 극복하기 위해.
• 이산 클러스터링과 상호정보를 사용하여 전역 및 국소 특징 표현을 효과적으로 정규화하는 방법을 개발하기 위해.

제안 방법

• 동일한 이미지의 변환된 뷰들 간의 클러스터 할당 간의 상호정보(MI)를 최대화하여 인코더에서 전역 표현 불변성을 강제한다.
• 연속된 특징 임베딩을 이산 클러스터 할당으로 매핑하기 위해 프로젝션 헤드를 사용하여 연속 밀도 추정이 필요 없이 효율적인 MI 추정을 가능하게 한다.
• 이웃하는 공간 특징 간의 MI를 최대화함으로써 디코더 특징 맵에서 공간 일관성을 유도하는 국소 MI 손실을 도입한다.
• 다른 네트워크 단계 간의 클러스터 레이블을 정렬하기 위해 최종 일관성 정규화 손실을 적용하여 일관된 의미 그룹화를 보장한다.
• 두 단계의 사전 훈련-피팅 트레이닝 설정을 피하기 위해, 감독 손실과 MI 기반의 전역 및 국소 정규화를 통합 최적화 프레임워크에 통합한다.
• 허용도 기반 공식을 사용한 대칭적 MI 추정 방법을 적용하여 임의의 클러스터링 해법을 방지한다.

실험 결과

연구 질문

• RQ1 이산 클러스터 할당에 대한 상호정보 정규화가 준지도 학습 의료 영상 분할 성능을 향상시킬 수 있는가?
• RQ2 전역 및 국소 MI 정규화를 결합함으로써 특징 표현과 분할 정확도가 어떻게 향상되는가?
• RQ3 감독 손실과 MI 기반 손실을 동시에 최적화하는 것이 두 단계의 사전 훈련-피팅 트레이닝 전략을 능가하는가?
• RQ4 MI 기반 방법이 최소한의 레이블 데이터로도 준지도 학습 성능에 근접한 성능을 달성할 수 있는가?
• RQ5 표준 일관성 기반 및 대비 학습 방법에 비해 제안된 방법은 안정성과 효율성 측면에서 어떻게 비교되는가?

주요 결과

• 제안된 방법은 네 개의 공개 의료 영상 분할 데이터셋에서 최신 준지도 학습 방법인 Mean Teacher와 엔트로피 최소화 방법을 모두 능가한다.
• ACDC 데이터셋에서 레이블이 5%일 경우 평균 DSC가 86.7%에 도달하여 전체 감독 학습 성능에 근접한다.
• 두 단계의 사전 훈련-피팅 트레이닝 기반 모델 대비 통합 최적화 전략이 분할 정확도를 5.10% 향상시켜, 레이블 데이터와 함께 종단 간 MI 정규화의 이점을 입증한다.
• 세 번의 독립적인 실행에서 안정적이고 일관된 성능을 기록하며, 표준편차가 낮아 초기화의 임의성에 대해 강건함을 보였다.
• 제거 분석 결과 전역 및 국소 MI 손실 모두 성능 향상에 기여하며, 전역 손실은 불변성을 향상시키고 국소 손실은 공간 스무쓰니스를 향상시킨다.
• 프로젝션 헤드를 통한 이산 클러스터 할당 사용은 효율적인 MI 추정을 가능하게 하며 연속적 MI 추정의 계산 부담을 피한다.