[Paper Review] Vision Transformer for COVID-19 CXR Diagnosis using Chest X-ray Feature Corpus

Sang Joon Park, Gwanghyun Kim+7 more|arXiv (Cornell University)|Mar 12, 2021

COVID-19 diagnosis using AIMedicine29 references26 citations

TL;DR

The paper presents a Vision Transformer that uses a low-level chest X-ray feature corpus, extracted by a pre-trained backbone, to diagnose COVID-19 and other infections with strong generalization across external datasets.

ABSTRACT

Under the global COVID-19 crisis, developing robust diagnosis algorithm for COVID-19 using CXR is hampered by the lack of the well-curated COVID-19 data set, although CXR data with other disease are abundant. This situation is suitable for vision transformer architecture that can exploit the abundant unlabeled data using pre-training. However, the direct use of existing vision transformer that uses the corpus generated by the ResNet is not optimal for correct feature embedding. To mitigate this problem, we propose a novel vision Transformer by using the low-level CXR feature corpus that are obtained to extract the abnormal CXR features. Specifically, the backbone network is trained using large public datasets to obtain the abnormal features in routine diagnosis such as consolidation, glass-grass opacity (GGO), etc. Then, the embedded features from the backbone network are used as corpus for vision transformer training. We examine our model on various external test datasets acquired from totally different institutions to assess the generalization ability. Our experiments demonstrate that our method achieved the state-of-art performance and has better generalization capability, which are crucial for a widespread deployment.

Motivation & Objective

• Motivate robust COVID-19 CXR diagnosis amid limited labeled data by leveraging abundant unlabeled CXRs.
• Propose a Vision Transformer that uses a backbone-derived low-level CXR feature corpus for improved embedding.
• Show that the model generalizes well to external datasets from different institutions and devices.

Proposed method

• Train a backbone network on a large public CXR dataset to extract low-level abnormal features (e.g., consolidation, GGO).
• Construct a feature corpus from intermediate backbone embeddings before PCAM pooling.
• Feed the projected features into a Vision Transformer with a class token to perform image-level diagnosis.
• Use a saliency-map based interpretability method via deep Taylor decomposition for localization.
• Evaluate with AUC, sensitivity, specificity, and accuracy across multiple external datasets.

Experimental results

Research questions

• RQ1 Can a Vision Transformer trained on a backbone-derived low-level CXR feature corpus outperform standard ViT and baselines on COVID-19 CXR diagnosis?
• RQ2 Does using a low-level feature corpus improve generalization to unseen, institutionally diverse data?
• RQ3 Is self-supervised pre-training beneficial for this architecture given the backbone pretraining?
• RQ4 What level of backbone finetuning is advantageous (fixed vs trainable) for generalization?

Key findings

• Achieves state-of-the-art-like performance and strong generalization on three external datasets (AUCs around 0.91–0.95,平均感度约87%,平均特异性约91%).
• Outperforms ResNet-50 baseline and ViT-based SOTA models across external tests.
• Trainable backbone yields better results than fixing backbone weights across external datasets.
• Self-supervised pre-training provides little to no benefit for the proposed model and may slightly hinder performance for some configurations.
• Provides interpretable saliency visualizations illustrating localization of COVID-19 and bacterial infections.