[논문 리뷰] SwinCheX: Multi-label classification on chest X-ray images with transformers

Sina Taslimi, Soroush Taslimi외 3명|arXiv (Cornell University)|2022. 06. 09.

COVID-19 diagnosis using AI인용 수 15

한 줄 요약

이 논문은 스위니 트랜스포머를 백본으로 사용하고 다중 작업 학습을 위한 다층 퍼셉트론(MLP) 헤드를 갖춘 다중 레이블 분류 모델인 SwinCheX를 제안한다. 3층의 MLP 헤드를 사용하여 ChestX-ray14 데이터셋에서 최신 기술 수준(SoA) 성능을 달성하며, 평균 AUC가 0.810에 도달하여 이전 SoA인 0.799를 초월한다.

ABSTRACT

According to the considerable growth in the avail of chest X-ray images in diagnosing various diseases, as well as gathering extensive datasets, having an automated diagnosis procedure using deep neural networks has occupied the minds of experts. Most of the available methods in computer vision use a CNN backbone to acquire high accuracy on the classification problems. Nevertheless, recent researches show that transformers, established as the de facto method in NLP, can also outperform many CNN-based models in vision. This paper proposes a multi-label classification deep model based on the Swin Transformer as the backbone to achieve state-of-the-art diagnosis classification. It leverages Multi-Layer Perceptron, also known as MLP, for the head architecture. We evaluate our model on one of the most widely-used and largest x-ray datasets called "Chest X-ray14," which comprises more than 100,000 frontal/back-view images from over 30,000 patients with 14 famous chest diseases. Our model has been tested with several number of MLP layers for the head setting, each achieves a competitive AUC score on all classes. Comprehensive experiments on Chest X-ray14 have shown that a 3-layer head attains state-of-the-art performance with an average AUC score of 0.810, compared to the former SOTA average AUC of 0.799. We propose an experimental setup for the fair benchmarking of existing methods, which could be used as a basis for the future studies. Finally, we followed up our results by confirming that the proposed method attends to the pathologically relevant areas of the chest.

연구 동기 및 목표

• 시각 트랜스포머를 사용한 다중 레이블 흉부 X-ray 이미지 분류를 위한 딥 러닝 모델을 개발한다.
• 기존의 CNN 기반 및 트랜스포머 기반 모델을 초월하여 ChestX-ray14 데이터셋에서 성능을 향상시킨다.
• 환자 기반 데이터 분할을 사용하여 향후 비교를 위한 공정한 벤치마킹 프로토콜을 수립한다.
• 병리학적으로 관련된 영역에서의 주의 맵 분석을 통해 모델의 해석 가능성(의미 부여 가능성)을 검증한다.
• 스위니와 같은 다중 해상도 트랜스포머가 기존의 CNN보다 다중 레이블 CXR 진단에서 우월함을 입증한다.

제안 방법

• 사전 학습된 스위니 트랜스포머를 흉부 X-ray 이미지의 특징 추출기로 적응시킨다.
• 깊이가 다양하게 설정된(1~3층) 다중 헤드 피드포워드 신경망(MLP)을 분류 헤드로 사용한다.
• 데이터 누출을 방지하기 위해 환자 기반의 훈련/검증/테스트 분할을 사용하여 ChestX-ray14 데이터셋에서 모델을 엔드 투 엔드로 훈련시킨다.
• Grad-CAM을 사용하여 시각화된 열매개 지apap을 생성하고 병변 관련 해부학적 영역에 대한 주의 집중 여부를 검증한다.
• 기존 방법과의 공정한 비교를 보장하기 위해 표준화된 평가 프로토콜을 적용하며, 동일한 분할 조건에서 DNet의 재평가도 수행한다.
• 적응형 학습률을 사용하지 않아 하이퍼파rameter의 복잡성을 줄이고 과적합 위험을 감소시킨다.

실험 결과

연구 질문

• RQ1 스위니 트랜스포머와 같은 시각 트랜스포머가 다중 레이블 흉부 X-ray 분류에서 CNN 기반 모델을 능가할 수 있는가?
• RQ2 다중 레이블 분류를 위한 흉부 X-ray 이미지에서 MLP 헤드의 최적 깊이(층수)는 무엇인가?
• RQ3 표준화된 환자 기반 데이터 분할 방식은 모델 평가의 성능 및 공정성에 어떤 영향을 미치는가?
• RQ4 제안된 모델이 Grad-CAM으로 확인된 병리학적으로 관련된 영역(예: 기도폐쇄 시 폐저부, 종양은 오른쪽 폐, 심장 확대는 심장 가장자리)에 주의를 기울이는가?
• RQ5 적응형 학습률을 사용하지 않는 간단한 훈련 설정이 여전히 최신 기술 수준 성능을 달성할 수 있는가?

주요 결과

• SwinCheX의 3층 MLP 헤드는 ChestX-ray14 데이터셋에서 평균 AUC 0.810을 달성하여 새로운 최신 기술 수준 성능을 확립한다.
• 제안된 동일한 평가 프로토콜 하에서 이전 SoA인 DNet보다 평균 AUC에서 0.011 높으며, 1.01%의 여유를 확보한다.
• Grad-CAM 시각화 결과, 모델이 기도폐쇄의 경우 폐저부, 종양의 경우 오른쪽 폐, 심장 확대의 경우 심장 가장자리 등 병변 관련 영역에 주의를 기울이는 것으로 확인된다.
• 모델은 14개의 병변 모두에서 뛰어난 성능를 유지하며, 심장 확대, 흉막 출혈, 흉막 출혈에서 DNet이 약간 우월한 경우를 제외하고는 거의 모든 병변에서 DNet을 능가한다.
• 제안된 공정한 벤치마킹 프로토콜 하에서 재평가한 결과 DNet의 평균 AUC는 0.799로 하락하여 일관된 평가의 중요성을 입증한다.
• SwinCheX에서 적응형 학습률을 사용하지 않음으로써 훈련을 단순화하고 과적합 위험을 감소시키면서도 뛰어난 성능 유지를 달성한다.