[논문 리뷰] A Comparative Study of CNN, ResNet, and Vision Transformers for Multi-Classification of Chest Diseases

Ananya Jain, Aviral Bhardwaj외 2명|arXiv (Cornell University)|2024. 05. 31.

AI in cancer detection인용 수 4

한 줄 요약

이 연구는 NIH 흉부 X-ray 데이터셋을 사용하여 14종의 흉부 질환에 대한 다중 분류를 위해 CNN, ResNet 및 비전 트랜스포머(ViT) 모델을 비교한다. 사전 훈련된 ViT 모델은 ImageNet-21k에서 미세조정된 경우 테스트 정확도 94%를 기록하며 CNN 및 ResNet을 능가하여, 대규모 데이터에서 사전 훈련된 주의 기반 특징 학습이 의료 영상에서 우수함을 입증한다.

ABSTRACT

Large language models, notably utilizing Transformer architectures, have emerged as powerful tools due to their scalability and ability to process large amounts of data. Dosovitskiy et al. expanded this architecture to introduce Vision Transformers (ViT), extending its applicability to image processing tasks. Motivated by this advancement, we fine-tuned two variants of ViT models, one pre-trained on ImageNet and another trained from scratch, using the NIH Chest X-ray dataset containing over 100,000 frontal-view X-ray images. Our study evaluates the performance of these models in the multi-label classification of 14 distinct diseases, while using Convolutional Neural Networks (CNNs) and ResNet architectures as baseline models for comparison. Through rigorous assessment based on accuracy metrics, we identify that the pre-trained ViT model surpasses CNNs and ResNet in this multilabel classification task, highlighting its potential for accurate diagnosis of various lung conditions from chest X-ray images.

연구 동기 및 목표

• X-ray 영상에서 14종의 흉부 질환에 대한 다중 레이블 분류에서 CNN, ResNet 및 비전 트랜스포머(ViT) 아키텍처의 성능을 평가하고 비교하는 것.
• ImageNet-21k와 같은 대규모 데이터셋에서의 사전 훈련이 의료 영상에서 ViT 성능에 미치는 영향을 평가하는 것.
• 저자원 의료 영상 환경에서 ViT와 전통적인 CNN 기반 모델 간의 일반화 능력 및 특징 학습 효율성에 대해 조사하는 것.
• 임상적 관련성과 의사결정 투명성 확보를 위해 ViT의 주의 맵과 모델 해석 가능성에 대해 분석하는 것.
• 다중 분류에서 데이터 불균형 및 모델 최적화 격차 등의 한계를 특정하는 것.

제안 방법

• ImageNet에서 사전 훈련된 ViT와 NIH 흉부 X-ray 데이터셋(112,120장의 정면 X-ray 영상)을 사용해 미세조정한 ViT 두 가지 변종을 훈련.
• 동일한 데이터셋과 훈련 프rotocol를 사용하여 기준 모델로 표준 CNN과 ResNet-50을 훈련 및 평가.
• ViT 인코더의 마지막 멀티헤드 주의층에서 주의 가중치를 추출하여 주의 맵을 생성하기 위해 사전 훈련된 InceptionResNetV2 모델을 사용한 전이 학습.
• 입력으로 224×224 픽셀 RGB 영상을 사용하고, 데이터 증강을 적용하며, 다중 레이블 분류를 위해 교차 엔트로피 손실과 시그모이드 활성화 함수를 사용해 모델을 훈련.
• 성능 및 일반화 능력을 비교하기 위해 정확도, AUC-ROC, 손실 곡선을 사용해 훈련 및 검증 세트에서 모델을 평가.
• 주의 맵을 시각화하기 위해 헤드 및 패치 간 주의 가중치를 평균화하고, 확대한 후 원본 이미지 위에 오버랩하여 모델 예측의 해석을 도모.

실험 결과

연구 질문

• RQ1 사전 훈련된 비전 트랜스포머가 흉부 X-ray 영상의 다중 레이블 분류에서 CNN 및 ResNet을 능가하는가?
• RQ2 ImageNet-21k와 같은 대규모 데이터셋에서의 사전 훈련이 의료 영상 분류에서 ViT 성능에 어떤 영향을 미치는가?
• RQ3 주의 메커니즘이 복잡한 특징 맵에 비해 흉부 X-ray 영상에서 임상적으로 관련 있는 영역을 식별하는 데 어떤 역할을 하는가?
• RQ4 데이터 불균형과 희귀 질환에 대한 제한된 예시가 다양한 아키텍처의 모델 성능에 어떤 영향을 미치는가?
• RQ5 ViT 모델이 기존 CNN보다 유사하거나 더 나은 해석 가능성을 주의 시각화를 통해 달성할 수 있는가?

주요 결과

• 사전 훈련된 비전 트랜스포머(ViT-v2)는 NIH 흉부 X-ray 데이터셋에서 미세조정된 결과 테스트 정확도 94%를 기록하며 CNN 및 ResNet 모델을 뛰어넘었다.
• ImageNet-21k에서 사전 훈련된 ViT 모델는 사전 훈련 없이 학습한 ViT 및 모든 기준 모델보다 뛰어난 성능을 보였으며, 대규모 사전 훈련의 중요성을 입증했다.
• AUC-ROC 분석 결과, ViT 모델은 'No Finding' 클래스에서 성능이 열악했으며(AUC = 0.50), 정상과 비정상 사례를 구분하는 데 어려움을 겪었고, Atelectasis(AUC = 0.51) 및 Pleural Thickening(AUC = 0.53)와 같은 질환에 대해 중간 수준의 성능를 보였다.
• ViT 모델은 Cardiomegaly(AUC = 0.48) 및 Hernia(AUC = 0.43)와 같은 희귀 질환에 대해 어려움을 겪었으며, 주로 데이터셋 내 데이터 불균형 때문이었다.
• ViT 모델에서 생성된 주의 맵은 심장 및 폐 영역과 같은 임상적으로 관련 있는 영역을 강조하여, 모델이 진단적으로 중요한 영역에 집중할 수 있음을 확인했다.
• 평균 성능는 우수했지만 'No Finding' 클래스에서 성능이 열악하여, 향후 연구에서 더 나은 데이터 샘플링 또는 클래스 균형 기법이 필요하다는 점을 시사했다.