[논문 리뷰] Constrained Optimization to Train Neural Networks on Critical and Under-Represented Classes

Sara Sangalli, Ertunç Erdil외 3명|arXiv (Cornell University)|2021. 02. 21.

Artificial Intelligence in Healthcare인용 수 10

한 줄 요약

이 논문은 클래스 불균형 데이터셋에서 부족하게 표현된 중요한 클래스에 대해 향상된 성능을 달성하기 위해 보완된 라그랑주 곱수 방법(ASL)을 사용한 제약 최적화 프레임워크를 제안한다. 높은 진성 양성률에서 거짓 양성률 감소를 비대칭적으로 우선시함으로써, AUC를 향상시키고 소수의(중요한) 클래스의 잘못 분류를 줄이며 감도를 유지함으로써 기존 손실 함수보다 의료 영상 및 비전 벤치마크에서 뛰어난 성능을 발휘한다.

ABSTRACT

Deep neural networks (DNNs) are notorious for making more mistakes for the classes that have substantially fewer samples than the others during training. Such class imbalance is ubiquitous in clinical applications and very crucial to handle because the classes with fewer samples most often correspond to critical cases (e.g., cancer) where misclassifications can have severe consequences. Not to miss such cases, binary classifiers need to be operated at high True Positive Rates (TPRs) by setting a higher threshold, but this comes at the cost of very high False Positive Rates (FPRs) for problems with class imbalance. Existing methods for learning under class imbalance most often do not take this into account. We argue that prediction accuracy should be improved by emphasizing reducing FPRs at high TPRs for problems where misclassification of the positive, i.e. critical, class samples are associated with higher cost. To this end, we pose the training of a DNN for binary classification as a constrained optimization problem and introduce a novel constraint that can be used with existing loss functions to enforce maximal area under the ROC curve (AUC) through prioritizing FPR reduction at high TPR. We solve the resulting constrained optimization problem using an Augmented Lagrangian method (ALM). Going beyond binary, we also propose two possible extensions of the proposed constraint for multi-class classification problems. We present experimental results for image-based binary and multi-class classification applications using an in-house medical imaging dataset, CIFAR10, and CIFAR100. Our results demonstrate that the proposed method improves the baselines in majority of the cases by attaining higher accuracy on critical classes while reducing the misclassification rate for the non-critical class samples.

연구 동기 및 목표

• 희귀하지만 심각한 경우(예: 암)를 잘못 분류할 경우 임상적 비용이 큰 의료 영상에서의 클래스 불균형 문제에 대한 핵심 과제를 해결한다.
• 높은 진성 양성률에서 작동하는 지점에서 거짓 양성률 감소에 초점을 맞춤으로써 소수의(중요한) 클래스 성능을 향상시킨다.
• 기존 손실 함수와 통합 가능한 원칙적이고 미분 가능한 제약 조건을 개발하여, 높은 진성 양성률에서 낮은 거짓 양성률을 우선시하는 비대칭적인 AUC 최적화를 가능하게 한다.
• 다중 클래스 상황에 광범위하게 적용 가능한 두 가지 제안된 공식화를 통해 이진 분류 제약 조건을 다중 클래스로 확장한다.
• 의료 영상 및 표준 비전 벤치마크를 포함한 다양한 데이터셋에서, 다양한 클래스 불균형 비율 조건 하에서도 일관된 성능 향상을 입증한다.

제안 방법

• 이진 DNN 학습을 AUC 최대화를 위해 비대칭적으로 높은 진성 양성률에서 거짓 양성률 감소에 중점을 두어 제약 최적화 문제로 공식화한다.
• Mann-Whitney U 통계량을 사용해 AUC 목표를 정의하고, 이를 최적화 프레임워크 내의 제약 조건으로 통합한다.
• 보완된 라그랑주 곱수 방법(ASL)을 적용하여 제약 최적화 문제를 확률적 경사 하강법으로 해결할 수 있는 이중 비제약 문제로 변환한다.
• ASL 기반 제약 조건을 BCE, 포칼 손실, 클래스 밸런스 BCE, LDAM을 포함한 9종의 기존 기준 손실 함수와 통합한다.
• 다중 클래스 문제를 위한 두 가지 확장인 ALM m,1 및 ALM m,2를 제안하여 다수의 중요 클래스를 다룰 수 있도록 제약 조건을 적응시킨다.
• 기존 아키텍처 및 학습 절차와의 호환성을 유지하면서 표준 딥 러닝 학습 파ip라인을 사용해 최종 손실를 최적화한다.

실험 결과

연구 질문

• RQ1 제약 최적화 프레임워크는 불균형 데이터셋에서 표현이 부족한 중요한 클래스에 대해 딥 네트워크 성능을 향상시킬 수 있는가?
• RQ2 높은 진성 양성률에서 거짓 양성률 감소에 초점을 맞춘 비대칭 AUC 최대화는 이진 분류 과제에서 더 높은 임상적 유용성을 이끌어낼 수 있는가?
• RQ3 기존의 클래스 불균형 손실 함수와 비교할 때, 제안된 ASL 기반 제약 조건은 거짓 양성률 및 진성 양성률 간의 트레이드오프에서 어떻게 성능을 내는가?
• RQ4 이진 제약 조건은 다중 클래스 설정으로 효과적으로 확장될 수 있는가? 이때 중요한 클래스의 성능 향상은 유지되는가?
• RQ5 이 방법은 의료 영상 및 표준 비전 벤치마크를 포함한 다양한 데이터셋에서 일관되게 성능 향상을 이끌어낼 수 있는가?

주요 결과

• 내부 MRI 데이터셋에서 ALM + BCE는 0% 거짓 음성률에서 54.0%의 거짓 양성률을 기록했으며, 기준 BCE(80.0%) 및 CB-BCE(100.0%)보다 뚜렷한 성능 향상을 보였다.
• ALM + BCE는 평균 AUC 76.8%와 앙상블 AUC 85.4%를 기록했으며, 최고의 기준 모델(MBAUC)의 82.4% 앙상블 AUC를 뛰어넘었다.
• 장꼬리형 CIFAR10(불균형 비율 1:100)에서 ALM m,2 + LDAM은 80% TPR에서 24.76% 오차, 90% TPR에서 31.02% 오차를 기록했으며, LDAM 및 CB-BCE를 포함한 모든 기준 모델을 능가했다.
• ALM m,2 + LDAM는 80% TPR에서 73.32%의 전체 정확도, 90% TPR에서 67.41%의 정확도를 기록하여 다양한 지표에서 일관된 성능 향상을 보였다.
• 대부분의 실험에서 ALM m,2 확장이 ALM m,1보다 더 큰 폭으로 전체 정확도를 향상시켰으며, 특히 LDAM 및 CB-BCE와 조합했을 때 두드러진 성능 향상을 보였다.
• 모든 데이터셋과 불균형 비율 조건에서 제안된 방법은 대부분의 경우 기준 모델을 향상시켰으며, 특히 높은 TPR에서의 거짓 양성률 감소 측면에서 효과를 입증하여 중요한 클래스 우선순위 부여에 효과적임을 확인했다.