[论文解读] Dynamic Routing on Deep Neural Network for Thoracic Disease Classification and Sensitive Area Localization

Yan Shen, Mingchen Gao|arXiv (Cornell University)|Aug 17, 2018

COVID-19 diagnosis using AI参考文献 15被引用 4

一句话总结

本文提出了一种新颖的动态路由机制，将其集成到密集神经网络中的1×1卷积层内，用于胸部X光片的胸腔疾病分类与定位。通过将路由一致性机制扩展至卷积特征图，并结合Grad-CAM实现可解释性，该模型在参数量更少的情况下实现了最先进（SOTA）的AUC性能，且定位精度更高，在平均准确率上比标准卷积基线模型高出15%。

ABSTRACT

We present and evaluate a new deep neural network architecture for automatic thoracic disease detection on chest X-rays. Deep neural networks have shown great success in a plethora of visual recognition tasks such as image classification and object detection by stacking multiple layers of convolutional neural networks (CNN) in a feed-forward manner. However, the performance gain by going deeper has reached bottlenecks as a result of the trade-off between model complexity and discrimination power. We address this problem by utilizing the recently developed routing-by agreement mechanism in our architecture. A novel characteristic of our network structure is that it extends routing to two types of layer connections (1) connection between feature maps in dense layers, (2) connection between primary capsules and prediction capsules in final classification layer. We show that our networks achieve comparable results with much fewer layers in the measurement of AUC score. We further show the combined benefits of model interpretability by generating Gradient-weighted Class Activation Mapping (Grad-CAM) for localization. We demonstrate our results on the NIH chestX-ray14 dataset that consists of 112,120 images on 30,805 unique patients including 14 kinds of lung diseases.

研究动机与目标

• 解决深度卷积神经网络中因模型复杂度与判别能力之间的权衡而导致的性能瓶颈问题。
• 通过在卷积层中引入新颖的动态路由机制，提升在NIH ChestX-ray14数据集上对胸腔疾病的分类准确率。
• 通过生成高保真Grad-CAM热力图，实现对敏感区域的精确定位，从而增强模型的可解释性。
• 通过遵循官方的训练/验证/测试集划分协议，避免同一患者图像在不同数据集间的数据泄露，确保公平评估。
• 证明在1×1卷积层中应用动态路由可实现更优的特征选择与分类性能，且所需层数更少。

提出的方法

• 将动态路由扩展至密集模块中的1×1卷积层，通过迭代一致性学习路由系数，以选择相关特征图。
• 采用核技巧高效计算路由，将训练与推理时间降低至接近单层复杂度。
• 模型使用下采样模块，包含7×7与3×3的最大池化操作，随后接1×1卷积和2×2平均池化，再进入密集路由层。
• 在4×4平均池化层之前，从32×32分辨率的初级胶囊生成Grad-CAM，以保留空间细节用于定位。
• 路由机制在密集层的特征图之间以及最终分类层的初级胶囊与预测胶囊之间运行。
• 网络采用端到端训练，使用标准交叉熵损失；推理阶段使用最终的路由迭代，并保留可学习的系数。

实验结果

研究问题

• RQ1 与标准卷积相比，1×1卷积层中的动态路由是否能提升胸部X光图像的分类性能？
• RQ2 将路由机制扩展至密集层连接是否能增强特征表示并提高模型准确率？
• RQ3 从低分辨率初级胶囊生成的Grad-CAM是否仍能准确地定位与疾病相关的区域？
• RQ4 在使用官方数据划分的情况下，该模型与最先进模型相比，在AUC和准确率方面表现如何？
• RQ5 通过Grad-CAM实现的模型可解释性在多大程度上与放射科医生标注的边界框在疾病定位上保持一致？

主要发现

• 所提出的模型在NIH ChestX-ray14数据集上实现了最先进（SOTA）的AUC性能，且参数量显著少于先前模型。
• 将1×1卷积胶囊替换为标准1×1卷积后，平均分类准确率下降15%，证明了动态路由的优越性。
• 从32×32初级胶囊生成的Grad-CAM热力图能准确地定位疾病区域，Atelectasis的重叠率（IoBB ≥ 0.1）达69.77%，Cardiomegaly达83.33%。
• 对于Atelectasis，模型在一位患者中正确定位了左上肺病灶，在另一位患者中正确识别了右上肺病灶，与临床标注一致。
• Atelectasis的定位准确率为45.34%（IoBB ≥ 0.25），Cardiomegaly达82.77%，表明与真实标注边界框高度一致。
• 该模型在不牺牲准确率的前提下保持了高水平的可解释性，因为Grad-CAM在全局平均池化前生成，从而保留了空间信息。