[論文レビュー] K-Diag: Knowledge-enhanced Disease Diagnosis in Radiographic Imaging

Chaoyi Wu, Xiaoman Zhang+3 more|arXiv (Cornell University)|Feb 22, 2023

AI in cancer detection被引用数 7

ひとこと要約

K-Diag は、対照学習を用いて医療知識グラフを統合し、学習可能なプロンプトベクトルを用いて視覚的表現を適応させる、知識を強化したフレームワークを提案する。このフレームワークは、長尾分布およびゼロショット疾患認識において最先端の性能を達成し、ゼロショット評価において未観測のレントゲン所見106例のうち79例で AUC >0.600 を達成した。

ABSTRACT

In this paper, we consider the problem of disease diagnosis. Unlike the conventional learning paradigm that treats labels independently, we propose a knowledge-enhanced framework, that enables training visual representation with the guidance of medical domain knowledge. In particular, we make the following contributions: First, to explicitly incorporate experts' knowledge, we propose to learn a neural representation for the medical knowledge graph via contrastive learning, implicitly establishing relations between different medical concepts. Second, while training the visual encoder, we keep the parameters of the knowledge encoder frozen and propose to learn a set of prompt vectors for efficient adaptation. Third, we adopt a Transformer-based disease-query module for cross-model fusion, which naturally enables explainable diagnosis results via cross attention. To validate the effectiveness of our proposed framework, we conduct thorough experiments on three x-ray imaging datasets across different anatomy structures, showing our model is able to exploit the implicit relations between diseases/findings, thus is beneficial to the commonly encountered problem in the medical domain, namely, long-tailed and zero-shot recognition, which conventional methods either struggle or completely fail to realize.

研究の動機と目的

• 従来の教師あり学習が疾患ラベルを独立かつ直交的とみなすという限界を是正するため、疾患ラベル間に内在する意味的関係を無視する問題に対処する。
• 関連する医療概念間の共有知識を活用することで、長尾分布疾患認識における一般化性能を向上させる。
• 視覚的表現学習プロセスにドメイン知識を統合することで、訓練中に観測されなかった疾患のゼロショット診断を可能にする。
• 疾患クエリと視覚的特徴の間のクロスアテンション機構を用いて、説明可能な診断フレームワークを構築する。

提案手法

• 疾患および所見概念を連続的かつ関係に配慮した空間に埋め込むために、医療知識グラフ上で対照学習を用いて知識エンコーダーを訓練する。
• 視覚的エンコーダーは、固定された知識エンコーダーのパラメータを用いて訓練され、学習可能なプロンプトベクトルモジュールが下流タスクへの効率的適応を可能にする。
• Transformerベースの疾患クエリモジュールが、疾患名（クエリ）と視覚的特徴の間でクロスアテンションを実行し、説明可能なアテンション重み付き予測を生成する。
• モデルは、一括ベクトルではなく、知識グラフから導出された連続的ラベル埋め込みを用いるため、病変間の意味的関係を活用できる。
• 視覚的特徴と知識埋め込み済み疾患記述を一致させる対照的目的関数を用いて、エンドツーエンドでモデルを訓練する。
• 推論時に未観測疾患名を知識空間に直接符号化することで、少数ショットおよびゼロショット推論を両方サポートする。

Figure 1: In the conventional training scheme (left), manual annotations are often converted into discrete one-hot vectors, that are orthogonal in the embedding space, thus ignoring the implicit relations between labels. While in our proposed knowledge-enhanced classification framework (right), the

実験結果

リサーチクエスチョン

• RQ1 視覚的表現学習に医療ドメイン知識を統合することで、レントゲン画像における長尾分布疾患認識の性能が向上するか？
• RQ2 知識を強化したフレームワークは、訓練中に観測されなかった疾患のゼロショット診断を可能にするか？
• RQ3 一括ベクトル符号化と比較して、連続的かつ知識埋め込み済みラベル表現は、疾患間の意味的関係をどの程度正確に捉えられるか？
• RQ4 知識を強化したモデルは、解剖的構造の多様性やアノテーションの粒度が異なるデータセットに対し、どの程度一般化できるか？
• RQ5 疾患クエリモジュールのアテンション機構は、モデル予測に対する説明可能なインサイトを提供できるか？

主な発見

• 3つのレントゲン画像データセットにおいて、K-Diag は一括ベクトルを用いたベースラインモデルと比較して、ResNet では AUC 2.53%、ViT では 2.15% の向上を達成した。
• 11個の公開胸部X線データセットを統合して学習した場合、K-Diag はベースラインより平均 AUC 2.14% 高く、より効果的なデータ利用を示した。
• PadChest におけるゼロショット評価で、K-Diag は未観測のレントゲン所見106例のうち79例で AUC が 0.600 以上を達成した。これは従来の教師あり学習では達成不可能な結果であった。
• ゼロショット設定において、14の所見で AUC が 0.800 以上、46の所見で 0.700 以上を達成した。これは、希少または新規病変に対する強力な一般化性能を示している。
• ターゲットクラスのデータ量を増やさずに、クラスの多様性（未観測疾患）を追加することで、すべての11のデータセットで性能が向上し、特に GoogleNIH や SIIM-ACR のような小規模データセットで顕著であった。
• アブレーションスタディにより、知識エンコーダーとプロンプト学習のコンponentsが不可欠であることが確認され、いずれかを削除すると性能が著しく低下した。

Figure 2: Overview of the knowledge-enhanced disease diagnosis workflow. The knowledge encoder (left) is first trained to learn a neural representation of the medical knowledge graph via contrastive learning, and then used to guide the visual representation learning in our knowledge-enhanced classif