[논문 리뷰] Vision Transformer for Efficient Chest X-ray and Gastrointestinal Image Classification

Smriti Regmi, Aliza Subedi외 2명|arXiv (Cornell University)|2023. 04. 23.

COVID-19 diagnosis using AI인용 수 10

한 줄 요약

논문은 Vision Transformer (ViT)와 DeiT 아키텍처가 Chest X-ray, Kvasir, Kvasir-Capsule의 세 가지 의학 영상 데이터셋에서 여러 지표에 걸쳐 여러 CNN 베이스라인을 능가하는 것을 보여주며, ViT를 의료 영상 분류 작업에 대한 강력한 벤치마크로 제시한다.

ABSTRACT

Medical image analysis is a hot research topic because of its usefulness in different clinical applications, such as early disease diagnosis and treatment. Convolutional neural networks (CNNs) have become the de-facto standard in medical image analysis tasks because of their ability to learn complex features from the available datasets, which makes them surpass humans in many image-understanding tasks. In addition to CNNs, transformer architectures also have gained popularity for medical image analysis tasks. However, despite progress in the field, there are still potential areas for improvement. This study uses different CNNs and transformer-based methods with a wide range of data augmentation techniques. We evaluated their performance on three medical image datasets from different modalities. We evaluated and compared the performance of the vision transformer model with other state-of-the-art (SOTA) pre-trained CNN networks. For Chest X-ray, our vision transformer model achieved the highest F1 score of 0.9532, recall of 0.9533, Matthews correlation coefficient (MCC) of 0.9259, and ROC-AUC score of 0.97. Similarly, for the Kvasir dataset, we achieved an F1 score of 0.9436, recall of 0.9437, MCC of 0.9360, and ROC-AUC score of 0.97. For the Kvasir-Capsule (a large-scale VCE dataset), our ViT model achieved a weighted F1-score of 0.7156, recall of 0.7182, MCC of 0.3705, and ROC-AUC score of 0.57. We found that our transformer-based models were better or more effective than various CNN models for classifying different anatomical structures, findings, and abnormalities. Our model showed improvement over the CNN-based approaches and suggests that it could be used as a new benchmarking algorithm for algorithm development.

연구 동기 및 목표

• 강력한 장거리 의존성 모델링으로 CNN 대비 효율적인 의료 영상 분류 대안의 필요성 동기 부여.
• 다중 모달 의료 데이터셋에서 ViT 및 DeiT 모델을 CNN 베이스라인과 비교 평가.
• 변환기 기반 의료 영상 분류를 향상시키기 위한 데이터 증강 및 학습 전략 조사.
• 적절한 지표를 사용하여 데이터셋 전반의 ViT 개선의 통계적 유의성 평가.

제안 방법

• 세 가지 의학 데이터셋에서 사전 학습된 ViT 변형(ViT-B/16, ViT-L/16, ViT-L/32)을 미세 조정.
• ImageNet-21k에서의 전이 학습을 사용하여 ViT/DeiT를 CNN 베이스라인 및 앙상블 모델과 비교.
• 클래스 불균형을 다루기 위해 데이터셋별 증강 및 손실 함수(교차 엔트로피 대 포컬 손실) 적용.
• MCC, ROC-AUC, 정밀도, 재현율, F1, 정확도, ROC 곡선을 포함한 지표로 평가; MCC 비교를 위한 대응 표본 t-검정 수행.

Figure 1 : An original ViT [ 7 ] structure for the classification task. The image is first converted into flattened patches through Patch Embedding and Position Embedding, then processed by the Transformer encoder [ 22 ] . The prediction result is obtained after the MLP Head.

실험 결과

연구 질문

• RQ1 ViT가 흉부 X-선, 내시경, 캡슐 내시경 데이터셋에서 MCC 및 ROC-AUC 측면에서 CNN 기반 모델을 능가하는가?
• RQ2 ViT 변형은 DeiT 및 CNN 앙상블과 다양한 의료 영상 모달리티에서 어떻게 비교되는가?
• RQ3 데이터 증강과 손실 함수가 트랜스포머 기반 의료 영상 분류 성능에 어떤 역할을 하는가?
• RQ4 관찰된 트랜스포머 기반 개선이 데이터셋에 걸쳐 통계적으로 유의한가?
• RQ5 ViT 기반 모델이 향후 의료 영상 분류 연구를 위한 강력한 벤치마크로 기능할 수 있는가?

주요 결과

• ViT-L/16은 Chest X-ray에서 평가된 모델 중 가장 높은 MCC를 달성하여 모든 지표에서 강한 성능을 나타냄.
• ViT 변형은 일반적으로 Chest X-ray와 Kvasir 데이터셋에서 CNN 베이스라인 및 DeiT를 여러 지표에서 능가.
• Kvasir-Capsule 데이터셋에서 ViT-B/16이 최상 MCC를 달성하며, 트랜스포머 모델이 가중치 기반 정밀도와 F1-점수에서 우위를 보임.
• ROC 곡선은 세 데이터셋 모두에서 ViT 모델의 경쟁력 있는 또는 우수한 성능을 시사.
• 대응 표본 t-검정은 Chest X-ray 및 Kvasir-Capsule 데이터셋에서 ViT의 MCC 개선이 다수의 SOTA 베이스라인에 대해 통계적으로 유의함을 시사하지만, 일부 Kvasir 비교에서는 항상 유의하지 않을 수 있음.

Figure 2 : Example samples from Chest X-ray, Kvasir, and Kvasir-Capsule datasets.