[論文レビュー] 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

ひとこと要約

本稿では、前胸部X線における胸部疾患の分類および局在化を目的とした、密度結合ニューラルネットワーク内の1×1畳み込み層に統合された新規な動的ルーティング機構を提案する。ルーティングによる一致を畳み込み特徴マップへ拡張し、Grad-CAMを用いた解釈可能性を実現することで、層数を減らしながらも最先端の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平均プーリングを実施し、密度ルーティング層へと接続する。
• 32×32解像度のプライマリキャプセルから、4×4平均プーリング層の直前でGrad-CAMを生成することで、局在化に必要な空間的詳細を保持する。
• ルーティング機構は、密度層内の特徴マップ間、および最終分類層におけるプライマリキャプセルと予測キャプセル間で動作する。
• ネットワークは標準的な交差エントロピー損失を用いてエンドツーエンドでトレーニングされ、推論時には学習可能な係数を有する最終ルーティングイテレーションが使用される。

実験結果

リサーチクエスチョン

• RQ1 1×1畳み込み層における動的ルーティングは、標準的な畳み込みと比較して胸部X線画像の分類性能を向上させることができるか？
• RQ2 密度結合部の接続にルーティングを拡張することで、特徴表現およびモデルの精度が向上するか？
• RQ3 低解像度のプライマリキャプセルから生成されたGrad-CAMは、依然として疾患関連領域を正確に局在化できるか？
• RQ4 公式のデータ分割を用いた場合、本モデルは最先端モデルと比較してAUCおよび精度において優れているか？
• RQ5 Grad-CAMによる解釈性は、放射線科医がアノテートしたバウンディングボックスとどの程度一致するか？

主な発見

• 提案手法は、先行モデルと比較して顕著に層数を減らしながらも、NIH ChestX-ray14データセットで最先端のAUC性能を達成した。
• 1×1畳み込みキャプセルを標準的な1×1畳み込みに置き換えた場合、平均分類精度が15%低下した。これは、動的ルーティングの優位性を示している。
• 32×32のプライマリキャプセルから生成されたGrad-CAMヒートマップは、疾患関連領域を正確に局在化しており、アテレクトーシスでは69.77%（IoBB ≥ 0.1）、心不全では83.33%の重複率を示した。
• アテレクトーシスの症例では、1例目の患者では左上肺病変を正しく局在化し、2例目では右上肺病変を正しく特定した。臨床的アノテーションと一致した。
• アテレクトーシスでは45.34%（IoBB ≥ 0.25）の局在化精度、心不全では82.77%の精度を達成し、正解バウンディングボックスと強い一致を示した。
• グローバル平均プーリングの直前でGrad-CAMを生成することで、空間情報を保持しつつも、精度を損なわず高い解釈性を維持した。