[논문 리뷰] Self-training with dual uncertainty for semi-supervised medical image segmentation

Zhanhong Qiu, Haitao Gan외 3명|arXiv (Cornell University)|2023. 04. 10.

Radiomics and Machine Learning in Medical ImagingMedicine인용 수 3

한 줄 요약

이 논문은 표본 수준 및 픽셀 수준의 이중 불확실성 추정을 통한 자기학습 프레임워크를 제안하여, 반감독 학습 기반 의료 영상 분할 성능을 향상시킨다. 표본 수준의 불확실성은 동일한 비라벨링 샘플에 대해 여러 체크포인트에서의 모델 예측 분산을 활용하고, 픽셀 수준의 불확실성은 서로 다른 업샘플링 방법(예: transposed convolution 대비 deconvolution)을 사용하는 두 개의 디코더 출력 간 차이를 기반으로 산정한다. 이로써 어려운 비라벨링 샘플만을 선택적으로 재학습하고, 의사라벨을 정교화함으로써 동일한 설정 하에서 ACDC 및 프로스타트 데이터셋에서 최신 기준 성능을 달성한다.

ABSTRACT

In the field of semi-supervised medical image segmentation, the shortage of labeled data is the fundamental problem. How to effectively learn image features from unlabeled images to improve segmentation accuracy is the main research direction in this field. Traditional self-training methods can partially solve the problem of insufficient labeled data by generating pseudo labels for iterative training. However, noise generated due to the model's uncertainty during training directly affects the segmentation results. Therefore, we added sample-level and pixel-level uncertainty to stabilize the training process based on the self-training framework. Specifically, we saved several moments of the model during pre-training, and used the difference between their predictions on unlabeled samples as the sample-level uncertainty estimate for that sample. Then, we gradually add unlabeled samples from easy to hard during training. At the same time, we added a decoder with different upsampling methods to the segmentation network and used the difference between the outputs of the two decoders as pixel-level uncertainty. In short, we selectively retrained unlabeled samples and assigned pixel-level uncertainty to pseudo labels to optimize the self-training process. We compared the segmentation results of our model with five semi-supervised approaches on the public 2017 ACDC dataset and 2018 Prostate dataset. Our proposed method achieves better segmentation performance on both datasets under the same settings, demonstrating its effectiveness, robustness, and potential transferability to other medical image segmentation tasks. Keywords: Medical image segmentation, semi-supervised learning, self-training, uncertainty estimation

연구 동기 및 목표

• 분할 작업에서 라벨가공된 의료 영상 데이터의 부족이라는 핵심 과제를 해결하기 위해.
• 모델의 불확실성으로 인한 노이즈를 줄임으로써 반감독 학습에서 자기학습 과정의 안정성을 높이기 위해.
• 불확실성 추정에 기반해 어려운 비라벨링 샘플을 선택적으로 재학습함으로써 분할 정확도를 향상시키기 위해.
• 이중 디코더를 활용한 픽셀 수준의 불확실성 추정을 통해 의사라벨 품질을 향상시키기 위해.
• 다양한 의료 영상 데이터셋에 걸쳐 일반화성과 강인성을 입증하기 위해.

제안 방법

• 표본 수준의 불확실성은 동일한 비라벨링 샘플에 대해 사전학습 단계에서 여러 체크포인트의 모델 예측 분산으로 추정된다.
• 비라벨링 샘플은 표본 수준의 불확실성 증가 순서로 점차적으로 학습에 통합되며, 쉬운 샘플에서 어려운 샘플로 진행된다.
• 픽셀 수준의 불확실성은 서로 다른 업샘플링 전략(예: transposed convolution 대비 deconvolution)을 사용하는 두 개의 디코더 출력 간 차이로 계산된다.
• 의사라벨은 픽셀 수준의 불확실성에 기반해 정교화되며, 불확실성이 낮은 영역일수록 더 높은 신뢰도가 부여된다.
• 모델는 자기학습을 통해 반복적으로 향상되며, 오직 가장 불확실한 샘플들만 재학습되어 오류 전파를 줄인다.
• 프레임워크는 아키텍처의 대대적 개선 없이도 표준 자기학습 파이프라인에 불확실성 인식 의사라벨링을 통합한다.

실험 결과

연구 질문

• RQ1 표본 수준의 불확실성 추정은 반감독 학습 기반 의료 영상 분할에서 의사라벨의 신뢰성 향상에 기여하는가?
• RQ2 이중 디코더에서 유도된 픽셀 수준의 불확실성은 의사라벨 품질 향상과 노이즈 감소에 기여하는가?
• RQ3 이중 불확실성 기반으로 어려운 비라벨링 샘플을 선택적으로 재학습하면 분할 성능 향상이 이루어지는가?
• RQ4 제안된 방법은 기준 의료 영상 데이터셋에서 기존 자기학습 접근법과 비교해 어떻게 성능을 내는가?
• RQ5 이 방법은 다양한 의료 영상 분할 작업 간에 얼마나 강인하고 이식 가능성이 있는가?

주요 결과

• 제안된 방법은 기준 방법들과 동일한 학습 프rotocol를 적용한 2017 ACDC 데이터셋에서 최신 기준 성능을 달성했다.
• 2018 프로스타트 데이터셋에서, 이 방법은 다섯 가지 기존 반감독 접근법보다 분할 정확도에서 뛰어난 성능을 보였다.
• 이중 불확실성 통합으로 인해 자기학습 과정에서 노이즈가 많은 의사라벨로 인한 오류 전파가 크게 감소했다.
• 오직 가장 불확실한 샘플들만 재학습함으로써 모델는 낮은 품질의 의사라벨로 인한 성능 저하를 피하며 개선된 강인성을 보였다.
• 이 방법은 두 가지 서로 다른 의료 영상 분할 작업에서 일관된 성능 향상을 보이며 강력한 일반화 잠재력을 보였다.
• 여러 모델 체크포인트와 이중 디코더의 활용은 추가적인 학습 또는 추론 비용 없이 신뢰할 수 있는 불확실성 추정을 가능하게 했다.