[논문 리뷰] Unimodal-uniform Constrained Wasserstein Training for Medical Diagnosis

Xiaofeng Liu, Xu Han외 3명|arXiv (Cornell University)|2019. 11. 03.

Medical Image Segmentation Techniques인용 수 4

한 줄 요약

이 논문은 순서형 레이블을 가진 의료 진단을 위한 유일모달-균일 제약 워셔스타인 훈련 손실을 제안한다. 이는 유클리드 거리의 증가 함수 기반 기초 거리 척도를 활용하여 클래스 상관관계를 인코딩한다. 매끄러운 타겟 레이블과 정확한 폐쇄형 해를 조합하여, 교차엔트로피 손실, 회귀, 이전의 순서형 방법들을 능가하는 최신 기술 성능을 달성한다. Diabetic Retinopathy 및 Ultrasound BIRADS 데이터셋에서 성능을 향상시킨다.

ABSTRACT

The labels in medical diagnosis task are usually discrete and successively distributed. For example, the Diabetic Retinopathy Diagnosis (DR) involves five health risk levels: no DR (0), mild DR (1), moderate DR (2), severe DR (3) and proliferative DR (4). This labeling system is common for medical disease. Previous methods usually construct a multi-binary-classification task or propose some re-parameter schemes in the output unit. In this paper, we target on this task from the perspective of loss function. More specifically, the Wasserstein distance is utilized as an alternative, explicitly incorporating the inter-class correlations by pre-defining its ground metric. Then, the ground metric which serves as a linear, convex or concave increasing function w.r.t. the Euclidean distance in a line is explored from an optimization perspective. Meanwhile, this paper also proposes of constructing the smoothed target labels that model the inlier and outlier noises by using a unimodal-uniform mixture distribution. Different from the one-hot setting, the smoothed label endues the computation of Wasserstein distance with more challenging features. With either one-hot or smoothed target label, this paper systematically concludes the practical closed-form solution. We evaluate our method on several medical diagnosis tasks (e.g., Diabetic Retinopathy and Ultrasound Breast dataset) and achieve state-of-the-art performance.

연구 동기 및 목표

• 표준 교차엔트로피 손실이 순서형 의료 진단에서 이웃 클래스 간 유사성과 순차적 레이블 구조를 간과하는 한계를 해결하기 위해.
• 유일모달-균일 혼합 분포를 통해 내부자와 외부자 불확실성을 포함시켜 임상 데이터의 레이블 노이즈를 모델링하기 위해.
• 워셔스타인 거리와 구조화된 기초 거리 척도를 사용하여 순서형 관계를 명시적으로 인코딩하는 가역적이고 효율적인 손실 함수를 개발하기 위해.
• 일반화형 및 매끄러운 타겟 레이블 모두에 대해 폐쇄형 해를 제공하여 빠르고 안정적인 훈련을 가능하게 하기 위해.
• 이산적이고 연속적인 위험 수준을 가진 실제 의료 데이터셋에서 뛰어난 성능을 입증하기 위해.

제안 방법

• 예측된 레이블 분포와 타겟 레이블 분포 사이의 최적 운반 비용을 측정하기 위해 제1의 워셔스타인 거리를 손실 함수로 사용한다.
• 선상의 유클리드 거리의 증가 함수로 정의된 기초 거리 척도를 사용하여 클래스 간 순서형 유사성을 인코딩한다.
• 진짜 레이블이 근처 또는 먼 클래스로 흐트러질 수 있는 유일모달-균일 혼합 분포를 도입하여 레이블 노이즈를 모델링한다.
• 일반화형 및 매끄러운 타겟 레이블 설정 하에서 워셔스타인 손실에 대한 정확한 폐쇄형 해를 유도하여 효율적인 최적화를 가능하게 한다.
• 최적화 행동을 탐색하기 위해 유클리드 거리의 볼록 및 오목 증가 함수를 기초 거리 척도로 사용한다.
• 추가적인 훈련 비용 없이도 일반화 및 강인성을 향상시키기 위해 매끄러운 타겟 레이블 분포를 적용한다.

실험 결과

연구 질문

• RQ1 구조화된 기초 거리 척도를 갖춘 워셔스타인 기반 손실이 표준 교차엔트로피 및 회귀 대비 순서형 의료 진단에서 성능 향상에 기여하는가?
• RQ2 유일모달-균일 혼합 분포를 통해 레이블 노이즈를 모델링하면 모델의 강인성과 정확도에 어떤 영향을 미치는가?
• RQ3 정확한 폐쇄형 해를 사용하는 것과 근사 최적화를 사용하는 것의 차이가 순서형 분류를 위한 워셔스타인 훈련에 어떤 영향을 미치는가?
• RQ4 기초 거리 척도의 다양한 형태(예: 선형, 볼록, 오목)가 모델 성능과 수렴에 어떤 영향을 미치는가?
• RQ5 제안된 방법이 당뇨성 망막병변(Diabetic Retinopathy)을 초월하여 US-BIRADS와 같은 다른 순서형 의료 데이터셋으로 일반화되는가?

주요 결과

• 유클리드 거리의 볼록 증가 함수(D²)를 사용한 제안된 워셔스타인 손실이 당뇨성 망막병변 데이터셋에서 최고 성능을 기록했으며, 정확도 80.6%와 QWK 0.716을 달성했다.
• 유일모달-균일 매끄러운 타겟 레이블이 모든 베이스라인에서 성능 향상을 이끌었으며, 특히 US-BIRADS와 같은 노이즈가 많은 환경에서는 뚜렷한 기여를 했다.
• 정확한 폐쇄형 해가 근사해보다 일관되게 뛰어난 성능을 보였으며, 더 빠른 수렴과 더 나은 일반화 능력을 입증했다.
• 당뇨성 망막병변 데이터셋에서 80.6%의 정확도와 0.716의 QWK를 기록하여, 다중태스크, 스틱브레이킹, 포isson 기반 접근법을 포함한 이전 방법들을 능가했다.
• US-BIRADS 데이터셋에서는 80.6%의 정확도와 0.716의 QWK를 달성했으며, 높은 레이블 노이즈로 인해 매끄러운 레이블이 더 큰 향상을 이끌었다.
• 하이퍼파라미터 민감도 분석 결과, τ와 ξ의 변화에 비해 QWK가 안정되어 있어 파rameter 튜닝에 대해 강인함을 보였다.