[논문 리뷰] Extreme Consistency: Overcoming Annotation Scarcity and Domain Shifts

Gaurav Fotedar, Nima Tajbakhsh외 2명|arXiv (Cornell University)|2020. 04. 15.

Retinal Imaging and Analysis참고 문헌 26인용 수 6

한 줄 요약

이 논문은 극단적인 데이터 변환을 통해 비라벨 데이터를 활용하는 강력한 일致성 정규화를 적용한 단순하고 일반적이며 유연한 준지도학습 프레임워크인 Extreme Consistency를 제안한다. 원본 입력에 대한 교사의 예측과, 중간 강도로 증강된 입력에 대한 학생 네트워크의 예측이 일致하도록 강제함으로써, 라벨이 제한된 조건과 도메인 이동 벤치마크에서 병변 및 망막 혈관 세분화 분야에서 최신 기술을 크게 앞서는 성능을 달성한다.

ABSTRACT

Supervised learning has proved effective for medical image analysis. However, it can utilize only the small labeled portion of data; it fails to leverage the large amounts of unlabeled data that is often available in medical image datasets. Supervised models are further handicapped by domain shifts, when the labeled dataset, despite being large enough, fails to cover different protocols or ethnicities. In this paper, we introduce \emph{extreme consistency}, which overcomes the above limitations, by maximally leveraging unlabeled data from the same or a different domain in a teacher-student semi-supervised paradigm. Extreme consistency is the process of sending an extreme transformation of a given image to the student network and then constraining its prediction to be consistent with the teacher network's prediction for the untransformed image. The extreme nature of our consistency loss distinguishes our method from related works that yield suboptimal performance by exercising only mild prediction consistency. Our method is 1) auto-didactic, as it requires no extra expert annotations; 2) versatile, as it handles both domain shift and limited annotation problems; 3) generic, as it is readily applicable to classification, segmentation, and detection tasks; and 4) simple to implement, as it requires no adversarial training. We evaluate our method for the tasks of lesion and retinal vessel segmentation in skin and fundus images. Our experiments demonstrate a significant performance gain over both modern supervised networks and recent semi-supervised models. This performance is attributed to the strong regularization enforced by extreme consistency, which enables the student network to learn how to handle extreme variants of both labeled and unlabeled images. This enhances the network's ability to tackle the inevitable same- and cross-domain data variability during inference.

연구 동기 및 목표

• 고비용 전문가 라벨링으로 인해 라벨이 제한되는 의료 영상 분석에서의 라벨 부족 문제를 해결한다.
• 촬영 프로토콜, 기기, 인종적 배경의 다양성으로 인한 도메인 이동 문제를 해결하여 모델의 일반화 능력을 향상시킨다.
• 추가적인 라벨 없이도 대량의 비라벨 데이터를 효과적으로 활용할 수 있는 준지도학습 방법을 개발한다.
• 분류, 세분화, 검출 작업에 적용 가능한 일반적이고 아키텍처에 종속되지 않는 프레임워크를 구축하며, 최소한의 구현 부담을 유도한다.
• 다양하고 확률적인 변환을 통해 불변성과 등변성을 강제함으로써, 데이터 변동성에 대한 모델의 강건성을 향상시킨다.

제안 방법

• 교사-학생 준지도학습 파라다임을 활용하여, 교사 네트워크는 원본이자 최소한의 증강을 가한 이미지를 처리하고, 학생 네트워크는 동일한 이미지의 극단적인 변환을 처리한다.
• 학생 네트워크의 입력에 대해 강도 기반, 기하학적, 이미지 혼합 변환의 다양한 조합을 적용하며, 변환 확률은 사전 분포에서 유도된다.
• 학생의 변환된 입력에 대한 출력과 교사의 원본 입력에 대한 출력 간 예측 불일치를 최소화함으로써 극단적 일치성을 강제한다.
• 교사 네트워크의 예측을 시간에 따라 안정화시키기 위해 메인 테이처 업데이트 전략을 사용하여 일치성 감독 성능을 향상시킨다.
• 비라벨 데이터의 다양성을 늘리고 도메인 간 일반화 능력을 강화하기 위해 교차 데이터셋 이미지 혼합 변환을 도입한다.
• 일치성 손실를 로짓(분류), 세분화 마스크(세분화), 그리고 민감도 맵(검출)에 적용함으로써, 다양한 작업에 광범위하게 적용 가능하도록 한다.

실험 결과

연구 질문

• RQ1 제한된 라벨과 도메인 이동 조건에서, 단순하고 일반적인 준지도학습 방법이 의료 영상 세분화 분야에서 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ2 극도로 변환된 입력을 사용하는 극단적 일치성 정규화가, 경미한 데이터 증강이나 표준 일치성 학습에 비해 얼마나 효과적인가?
• RQ3 강도, 기하학적, 혼합 변환 등의 다양한 변환 유형이 모델 성능과 학습 안정성에 기여하는 정도는 어떠한가?
• RQ4 교사의 입력은 그대로 유지하면서 학생 네트워크에만 극단적 변환을 적용하는 것이 대칭적 증강과 비교해 성능 향상에 뚜렷한 기여를 하는가?
• RQ5 실제 의료 영상 환경에서 데이터 다양성과 라벨 부족이 존재하는 상황에서, Extreme Consistency가 지도 학습 기반 모델과 최신 기술 준지도학습 모델을 모두 능가할 수 있는가?

주요 결과

• ISIC’17 데이터셋에서 Extreme Consistency는 최고 성능을 보인 SSL 방법 대비 2점의 Dice 점수 향상을 기록했으며, 평균 Dice는 0.836으로, 다음으로 좋은 방법의 0.816을 상회한다.
• STARE 데이터셋에서 이 방법은 평균 Dice(0.828 대 0.823)와 표준편차(0.002 대 0.004) 모두에서 두 번째로 좋은 SSL 방법을 앞서며, 더 높은 안정성을 보였다.
• 도메인 이동 실험에서 Extreme Consistency는 U-Net 지도 학습 기반 모델 대비 평균 6점의 Dice 향상을 기록했고, 최고의 SSL 방법 대비 3점의 향상을 기록했으며, 모든 향상이 통계적으로 유의미했다(p < 0.05).
• 제거 실험 결과, 다양한 변환을 제거하면 Dice 점수는 1.3–1.6% 감소하고, 이미지 혼합을 제외하면 1.3% 감소 및 변동성이 증가하여, 이들이 성능과 안정성에 결정적인 역할을 한다는 것을 입증했다.
• 비라벨 데이터를 제거하거나 Π-모델에서처럼 교사를 학생으로 대체하면 각각 Dice 점수는 1.4%와 0.6% 감소하여, 비라벨 데이터와 교사-학생 아키텍처의 필수성을 입증했다.
• 극단적 증강을 적용한 지도 학습 U-Net을 학습한 결과 유의미한 향상이 없었으며(Dice: 0.807), 이는 성능 향상의 원인이 데이터 증강이 아니라 일치성 정규화 메커니즘이라는 점을 확인한다.