[论文解读] SwinCheX: Multi-label classification on chest X-ray images with transformers

Sina Taslimi, Soroush Taslimi+3 more|arXiv (Cornell University)|Jun 9, 2022

COVID-19 diagnosis using AI被引用 15

一句话总结

该论文提出SwinCheX，一种基于Swin Transformer作为主干网络并使用多层感知机（MLP）头进行多任务学习的多标签分类模型，用于胸部X光片图像分析。在ChestX-ray14数据集上，采用3层MLP头时，该模型实现了SOTA性能，平均AUC达到0.810，超越了此前SOTA的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.

研究动机与目标

• 开发一种基于视觉Transformer的深度学习模型，用于胸部X光片图像的多标签分类。
• 在ChestX-ray14数据集上超越现有的基于CNN和基于Transformer的模型性能。
• 通过患者级别的数据划分建立公平的基准测试协议，以供未来比较。
• 通过分析注意力图在病理相关区域的表现，验证模型的可解释性。
• 证明多分辨率Transformer（如Swin）在多标签胸部X光片诊断中优于传统CNN。

提出的方法

• 将预训练的Swin Transformer作为胸部X光片图像的特征提取器。
• 采用多头前馈神经网络（MLP）作为分类头，其深度可调（1至3层）。
• 在ChestX-ray14数据集上使用患者级别的训练/验证/测试划分进行端到端训练，以防止数据泄露。
• 使用Grad-CAM生成显著性热力图，验证模型对疾病相关解剖区域的关注程度。
• 应用标准化的评估协议，确保与先前方法的公平比较，包括在相同划分下重新评估DNet。
• 在不使用自适应学习率的情况下优化模型，从而降低超参数复杂度和过拟合风险。

实验结果

研究问题

• RQ1 视觉Transformer（如Swin Transformer）是否能在多标签胸部X光片分类任务中超越基于CNN的模型？
• RQ2 在胸部X光片图像的多标签分类任务中，MLP头的最佳深度是多少？
• RQ3 标准化的患者级别数据划分如何影响模型评估的性能与公平性？
• RQ4 所提出的模型是否如Grad-CAM所证实的那样，关注到肺部的病理相关区域？
• RQ5 在不使用自适应学习率的简化训练设置下，是否仍能实现SOTA性能？

主要发现

• SwinCheX中采用的3层MLP头在ChestX-ray14数据集上实现了0.810的平均AUC，创下新的SOTA纪录。
• 在相同评估协议下，该模型相比此前SOTA模型DNet，平均AUC高出0.011，提升幅度达1.01%。
• Grad-CAM可视化结果表明，模型能正确关注到疾病相关区域，例如：肺底部（肺不张）、右肺（肿块）、心脏边缘（心脏扩大）。
• 该模型在全部14种病理类型上均保持强劲性能，仅在少数情况下DNet表现略优，包括心脏扩大、疝气和积液。
• 当在所提出的公平基准测试协议下重新评估时，DNet的平均AUC下降至0.799，证实了评估一致性的重要性。
• SwinCheX中未使用自适应学习率，简化了训练流程，降低了过拟合风险，同时保持了卓越的性能。