[논문 리뷰] CycleMix: A Holistic Strategy for Medical Image Segmentation from Scribble Supervision

Ke Zhang, Xiahai Zhuang|arXiv (Cornell University)|2022. 03. 03.

Radiomics and Machine Learning in Medical Imaging인용 수 5

한 줄 요약

CycleMix는 주로 스크래치(annotation)만을 사용하여 의료 영상 분할을 위한 새로운 비지도 학습 프레임워크를 제안한다. 사이클 일관성 있는 믹스업 증강을 활용하여 모델의 일반화 능력과 일관성을 향상시킨다. 이미지 혼합을 통해 주어진 픽셀의 커버리지를 극대화하고, 전역 및 국소적 사이클 일관성을 강제함으로써, 완전 지도 학습 모델과 비슷하거나 그 이상의 성능을 달성하며, 애너테이션 부담을 크게 줄인다.

ABSTRACT

Curating a large set of fully annotated training data can be costly, especially for the tasks of medical image segmentation. Scribble, a weaker form of annotation, is more obtainable in practice, but training segmentation models from limited supervision of scribbles is still challenging. To address the difficulties, we propose a new framework for scribble learning-based medical image segmentation, which is composed of mix augmentation and cycle consistency and thus is referred to as CycleMix. For augmentation of supervision, CycleMix adopts the mixup strategy with a dedicated design of random occlusion, to perform increments and decrements of scribbles. For regularization of supervision, CycleMix intensifies the training objective with consistency losses to penalize inconsistent segmentation, which results in significant improvement of segmentation performance. Results on two open datasets, i.e., ACDC and MSCMRseg, showed that the proposed method achieved exhilarating performance, demonstrating comparable or even better accuracy than the fully-supervised methods. The code and expert-made scribble annotations for MSCMRseg are publicly available at https://github.com/BWGZK/CycleMix.

연구 동기 및 목표

• 스크래치 수준의 지도 학습만을 사용하여 의료 영상 분할의 애너테이션 부담을 줄이기.
• 애너테이션된 픽셀의 효과적 사용을 극대화하여 제한된 지도 학습 하에서 모델의 일반화 능력과 강건성을 향상시키기.
• 데이터 증강과 일관성 정규화를 조합한 통합 전략을 개발하여 분할 성능을 향상시키기.
• 밀도 높은 애너테이션에 의존하지 않고 최소한의 인간 애너테이션으로 정확한 전체 분할 예측을 가능하게 하기.

제안 방법

• CycleMix는 두 이미지를 학습된 이동 행렬을 통해 혼합하여 혼합된 이미지 내에서 애너테이션된 픽셀 수를 극대화하는 사이클 일관성 있는 믹스업 전략을 사용한다.
• 모델은 혼합 및 원본 샘플에 대한 지도 학습 손실과 함께 전역 및 국소 일관성 손실을 조합한 가중 손실 함수를 사용하여 구조적 일관성을 강제한다.
• 이동 행렬 Π₁ᵀ는 PuzzleMix에서 영감을 얻은 교대 최적화 기법을 사용하여 혼합 이미지 내에서의 선명도를 극대화하도록 최적화된다.
• 전역 일관성 손실은 원본 및 재구성된 이미지 간의 일관하지 않은 예측을 방지하며, 국소 일관성 손실은 분할 구조의 단절을 방지한다.
• 손실 곡선을 통해 200 에포크 내에 안정적인 수렴을 보이며, 종합적인 엔드 투 엔드 학습이 가능하다.

(a) The curves of weighted loss functions.

실험 결과

연구 질문

• RQ1 데이터 혼합과 사이클 일관성을 조합한 통합 전략이 스크래치 지도 하에서 분할 성능을 향상시킬 수 있는가?
• RQ2 믹스업을 통해 애너테이션된 픽셀 비율을 증가시키면 모델의 일반화 능력과 정확도에 어떤 영향을 미치는가?
• RQ3 전역 및 국소 일관성 손실이 학습 안정성과 분할 품질 향상에 어떤 역할을 하는가?
• RQ4 제안된 방법이 CRF를 사용하는 기존의 자기지도 학습 방법(TS-UNet)을 능가하는 이유는 무엇인가?

주요 결과

• CycleMix는 완전 지도 학습 모델과 비슷하거나 그 이상의 성능을 달성하며, 스크래치 애너테이션만으로도 학습이 가능하다.
• 학습 과정이 안정적이며, 모든 손실 성분이 200 에포크 내에 수렴하여 신뢰할 수 있는 최적화를 보여준다.
• 전역 일관성 손실 (ℒcon-g)은 초기에 비정확한 무라벨 예측으로 인해 증가하지만, 50 에포크 이후 감소하여 무라벨 데이터로부터 효과적인 지식 전이가 이루어지고 있음을 시사한다.
• 국소 일관성 손실 (ℒcon-l)은 점진적으로 감소하며, 예측의 구조적 일관성이 향상되고 있음을 나타낸다.
• TS-UNet에 CRF를 적용한 기존 방법은 초기 제안이 정확하지 않아 자기학습 과정에서 오류가 누적되어 성능이 열 劣하다.

(b) $\mathcal{L}_{\text{mix}}$