[論文レビュー] HydraViT: Adaptive Multi-Branch Transformer for Multi-Label Disease Classification from Chest X-ray Images

Şaban Öztürk, Mehmet Y. Turali+1 more|arXiv (Cornell University)|Oct 9, 2023

COVID-19 diagnosis using AI被引用数 5

ひとこと要約

HydraViTは、多ラベル胸部X線（CXR）疾患分類のための新しい適応的マルチブランチトランスフォーマー・モデルを提案する。畳み込み的空間エンコーダーとトランスフォーマー基盤のコンテキストエンコーダーを組み合わせ、長距離依存関係を捉え、学習された重み付けを備えたマルチブランチ出力モジュールにより、個々の病変および併存病変の両方に対する感受性を向上させる。このモデルは最先端の性能を達成し、既存手法よりも平均1.2%高い性能を示した。

ABSTRACT

Chest X-ray is an essential diagnostic tool in the identification of chest diseases given its high sensitivity to pathological abnormalities in the lungs. However, image-driven diagnosis is still challenging due to heterogeneity in size and location of pathology, as well as visual similarities and co-occurrence of separate pathology. Since disease-related regions often occupy a relatively small portion of diagnostic images, classification models based on traditional convolutional neural networks (CNNs) are adversely affected given their locality bias. While CNNs were previously augmented with attention maps or spatial masks to guide focus on potentially critical regions, learning localization guidance under heterogeneity in the spatial distribution of pathology is challenging. To improve multi-label classification performance, here we propose a novel method, HydraViT, that synergistically combines a transformer backbone with a multi-branch output module with learned weighting. The transformer backbone enhances sensitivity to long-range context in X-ray images, while using the self-attention mechanism to adaptively focus on task-critical regions. The multi-branch output module dedicates an independent branch to each disease label to attain robust learning across separate disease classes, along with an aggregated branch across labels to maintain sensitivity to co-occurrence relationships among pathology. Experiments demonstrate that, on average, HydraViT outperforms competing attention-guided methods by 1.2%, region-guided methods by 1.4%, and semantic-guided methods by 1.0% in multi-label classification performance.

研究の動機と目的

• 胸部X線（CXR）画像における病変のサイズや位置のばらつき、および併存の有無という課題に取り組む。
• 従来のCNNやアテンション拡張モデルが病変間の長距離コンテキストおよび併存パターンを捉えるのを制限する点を克服する。
• 個々の疾患ラベルを別々に学習し、同時にラベルの併存関係を統合的にモデル化するマルチブランチ出力モジュールを統合することで、分類性能を向上させる。
• CXRスキャンからの複数の胸部疾患の自動かつ正確な診断を実現する、計算効率的かつ高感度なディープラーニングフレームワークを開発する。

提案手法

• CXR画像から初期特徴マップを抽出する畳み込み的空間エンコーダーと、画像パッチ間の長距離依存関係をモデル化するトランスフォーマー基盤のコンテキストエンコーダーを組み合わせたハイブリッドアーキテクチャを採用する。
• トランスフォーマー・エンコーダーにおける自己アテンション機構を活用し、タスクに重要な病変関連領域に適応的に注目することで、小さな病変や散在する病変に対する感受性を向上させる。
• 14の疾患ラベルそれぞれに専用のヘッドを設け、病変の併存関係をモデル化する集約ヘッドを備えたマルチブランチ出力モジュールを実装する。
• トレーニング中に個々のラベル性能と併存感度の両立を図るために、マルチブランチ出力の間で学習可能な適応的重み付けを導入する。
• クラスの不均衡を緩和し一般化性能を向上させるために、ラベル固有の損失項と結合損失項を組み合わせ、学習可能な重みを用いてトレーニングする。
• 解釈可能性を高めるためにGradCAMを適用し、代表的なCXR画像においてアテンションが病理的領域に正しく局在していることを検証する。

実験結果

リサーチクエスチョン

• RQ1 標準的なCNNと比較して、小規模または散在する病変を有するCXR画像において、ハイブリッド畳み込み的-トランスフォーマー・バックボーンが長距離コンテキストモデリングをどのように改善するか？
• RQ2 学習可能な重み付けを備えたマルチブランチ出力モジュールは、個々の病変と併存病変の両方を同時に最適化することで、多ラベルCXR分類性能を向上させるか？
• RQ3 平均的および各ラベルごとの分類性能において、HydraViTは最先端のアテンション誘導型、領域誘導型、意味的誘導型手法と比較してどの程度優れているか？
• RQ4 提案されたアーキテクチャは、アテンション可視化を通じて、局在化精度とモデルの解釈性をどの程度向上させるか？

主な発見

• HydraViTは、CXRデータセットにおける多ラベル分類性能において、競合するアテンション誘導型手法を平均1.2%上回った。
• 領域誘導型手法に対しては1.4%、意味的誘導型手法に対しては1.0%の平均F1スコアの向上を達成した。
• 予測ラベルと正解ラベルの整合性が優れており、トップ5予測における真のラベルの平均予測スコアは、非存在ラベルのスコア（例：0.4 vs. 0.1）と比較して顕著に高い値を示した。
• アブレーションスタディの結果、トランスフォーマー基盤のコンテキストエンコーダーとマルチブランチ出力モジュールの両方が性能向上に顕著な貢献をしていることが確認され、特に後者は逐次的な単一ブランチ学習に比べて約14倍の高速化を実現した。
• GradCAMのヒートマップの視覚的点検により、HydraViTが病変関連領域を効果的に強調していることが確認され、アテンションの局在化が堅牢であることが示された。
• 低発症率や複雑な併存パターンを示す疾患クラスに対しても、多様な病変クラスにわたり高い性能を維持しており、一般化性能の向上とバイアスの低減が示唆された。