[논문 리뷰] Chest X-rays Classification: A Multi-Label and Fine-Grained Problem

Zongyuan Ge, Dwarikanath Mahapatra외 3명|arXiv (Cornell University)|2018. 07. 19.

COVID-19 diagnosis using AI참고 문헌 11인용 수 28

한 줄 요약

이 논문은 흉부 X선 영상 분류에서 다중 레이블, 미세한 분류, 클래스 불균형 문제를 해결하기 위해 새로운 다중 레이블 소프트맥스 손실(MSML)을 제안한다. 이중선형 풀링을 사용하는 미세한 특징 학습 프레임워크와 MSML을 통합함으로써, ChestX-ray14 데이터셋에서 기존의 모델들(예: DenseNet-121)을 능가하는 최신 기술 수준의 AUC-ROC 스코어 0.8537을 달 đạt한다.

ABSTRACT

The widely used ChestX-ray14 dataset addresses an important medical image classification problem and has the following caveats: 1) many lung pathologies are visually similar, 2) a variant of diseases including lung cancer, tuberculosis, and pneumonia are present in a single scan, i.e. multiple labels and 3) The incidence of healthy images is much larger than diseased samples, creating imbalanced data. These properties are common in medical domain. Existing literature uses stateof- the-art DensetNet/Resnet models being transfer learned where output neurons of the networks are trained for individual diseases to cater for multiple diseases labels in each image. However, most of them don't consider relationship between multiple classes. In this work we have proposed a novel error function, Multi-label Softmax Loss (MSML), to specifically address the properties of multiple labels and imbalanced data. Moreover, we have designed deep network architecture based on fine-grained classification concept that incorporates MSML. We have evaluated our proposed method on various network backbones and showed consistent performance improvements of AUC-ROC scores on the ChestX-ray14 dataset. The proposed error function provides a new method to gain improved performance across wider medical datasets.

연구 동기 및 목표

• 의료 영상에서 흉부 X선 영상 분류의 다중 레이블, 미세한 분류, 클래스 불균형 특성을 다루기 위해.
• 기본 이진 교차 엔트로피 손실을 사용한 표준 전이 학습을 넘어서 ChestX-ray14 데이터셋에서 딥 러닝 성능을 향상시키기 위해.
• 레이블 간 상관관계를 명시적으로 모델링하고, 부정적 클래스에 대한 과적합을 억제하는 손실 함수를 개발하기 위해.
• MSML과 이중선형 풀링을 조합했을 때, 더 작은 아키텍처(예: ResNet-18)가 더 큰 아키텍처(예: DenseNet-121)를 능가할 수 있음을 보여주기 위해.
• 유사한 데이터 특성을 지닌 더 넓은 의료 영상 작업에 적용 가능한 새로운 학습 메커니즘을 제공하기 위해.

제안 방법

• 양성 레이블에 대해 모든 부정 레이블 대비 소프트맥스를 계산하여 레이블 간 의존성을 모델링하는 새로운 오차 함수인 다중 레이블 소프트맥스 손실(MSML)을 제안한다.
• MSML는 $ E^{MSML} = \frac{1}{|Y_i|} \sum_{l \in Y_i} \frac{\exp(x_l^i)}{\exp(x_l^i) + \sum_{k \in \bar{Y}_i} \exp(x_k^i)} $ 로 정의되며, 여기서 $ Y_i $ 는 샘플 $ i $ 의 양성 레이블 집합이다.
• 시각적으로 유사한 병변에 대해 구분 능력을 향상시키기 위해 미세한 분류에서 유래한 이중선형 풀링을 통합한다.
• 공유된 특징 추출기와 이중선형 융합을 사용하는 이중 스트림 네트워크 아키텍처를 통해 클래스 간 상호관계를 모델링한다.
• 모든 구성 요소(특징 추출기, MSML 헤드)가 함께 최적화되는 글로벌 훈련 전략을 사용한다.
• 최종 레이어에서 시그모이드 활성화 함수와 MSML을 조합하여 캘리브레이션된 출력을 갖는 다중 레이블 예측을 가능하게 한다.

실험 결과

연구 질문

• RQ1 명시적으로 레이블 상관관계를 모델링하는 새로운 손실 함수가 기존의 표준 이진 교차 엔트로피를 넘어서 다중 레이블 흉부 X선 영상 분류 성능을 향상시킬 수 있는가?
• RQ2 미세한 특징 학습 기법을 통합하면 시각적으로 유사한 폐 병변에 대한 성능 향상이 이루어지는가?
• RQ3 제안된 MSML과 이중선형 풀링을 사용할 때, ResNet-18과 같은 경량 모델이 DenseNet-121과 같은 더 큰 모델을 능가할 수 있는가?
• RQ4 다중 손실 학습에서 글로벌 훈련 전략은 국소 또는 고정된 미세조정 전략에 비해 어떻게 비교되는가?
• RQ5 MSML이 의료 영상 데이터셋에서 클래스 불균형과 부정적 레이블에 대한 과적합을 어느 정도 완화하는가?

주요 결과

• 제안된 방법은 DenseNet-121 기반으로 ChestX-ray14 데이터셋에서 AUC-ROC 스코어 0.8537을 달성하여 이전 최고 성능인 0.8413을 초월한다.
• 더 작은 ResNet-18 기반에서도 AUC 0.8438을 기록하여 더 큰 DenseNet-121 기반의 베이스라인을 능가한다.
• 글로벌 훈련 전략이 가장 높은 성능(정확도 AUC 0.8388)을 보였으며, 국소 및 고정된 미세조정 전략을 모두 능가한다.
• MSML는 W-AUC와 D-AUC 스코어를 크게 향상시켜 소수의 클래스와 어려운 클래스에서의 성능 향상을 나타낸다.
• 이중선형 풀링과 MSML의 조합은 표준 전이 학습 대비 ResNet-18 기반에서 AUC 1.5% 향상되었다.
• MSML의 루트 소프트맥스 메커니즘 덕분에 부정적 클래스에 대한 과적합이 효과적으로 억제된다. 이는 부재한 레이블에 대해 출력을 정규화하기 때문이다.