[論文レビュー] UniBrain: Universal Brain MRI Diagnosis with Hierarchical Knowledge-enhanced Pre-training

Jiayu Lei, Lisong Dai+10 more|arXiv (Cornell University)|Sep 13, 2023

Radiomics and Machine Learning in Medical Imaging被引用数 5

ひとこと要約

UniBrainは、24,770組の画像レポート対を用いて、包括的脳MRI診断を実現する階層的知識強化事前学習フレームワークを提案する。微細なモodal別レポート分解と連携された視覚言語認識を活用し、クロスモーダル整合性を向上させ、特定疾患において人間の専門家並みの精度を達成する優れた診断性能を実現した。

ABSTRACT

Magnetic resonance imaging~(MRI) have played a crucial role in brain disease diagnosis, with which a range of computer-aided artificial intelligence methods have been proposed. However, the early explorations usually focus on the limited types of brain diseases in one study and train the model on the data in a small scale, yielding the bottleneck of generalization. Towards a more effective and scalable paradigm, we propose a hierarchical knowledge-enhanced pre-training framework for the universal brain MRI diagnosis, termed as UniBrain. Specifically, UniBrain leverages a large-scale dataset of 24,770 imaging-report pairs from routine diagnostics. Different from previous pre-training techniques for the unitary vision or textual feature, or with the brute-force alignment between vision and language information, we leverage the unique characteristic of report information in different granularity to build a hierarchical alignment mechanism, which strengthens the efficiency in feature learning. Our UniBrain is validated on three real world datasets with severe class imbalance and the public BraTS2019 dataset. It not only consistently outperforms all state-of-the-art diagnostic methods by a large margin and provides a superior grounding performance but also shows comparable performance compared to expert radiologists on certain disease types.

研究の動機と目的

• 狭い疾患カテゴリと小規模データに焦点を当てた既存のMRI診断モデルの一般化能力の制限に対処すること。
• マルチモodal脳MRI事前学習におけるデータ不足、ドメイン特化知識の統合、解釈可能性の課題を克服すること。
• 多様な脳疾患を診断可能なスケーラブルで包括的なフレームワークを構築し、ゼロショットおよび少数ショット一般化能力を強化すること。
• 画像特徴と微細なレポート部品を対応付ける視覚的根拠付けにより、モデルの解釈可能性を向上させること。
• 可変長の疾患記述をサポートするクエリ駆動型認識モジュールを備え、柔軟な診断を可能とすること。

提案手法

• 微細なクロスモーダル学習を可能とするために、レポートをモダリティ別およびグローバル成分に分解する階層的ビジョン・ランゲージ整合化機構を提案すること。
• 自由テキストレポートから疾患関連セグメントを抽出する自動レポート分解モジュールを設計し、対応するMRIモダリティと標的的に整合させること。
• 疾患記述をクエリとして用い、グローバル画像特徴に注目し、疾患固有の分類およびグランドイング予測を生成する、連携視覚言語認識（CVP）モジュールを導入すること。
• グローバルなレポート・画像整合性とモダリティ別整合性を組み合わせた多段階コントラスト損失を実装し、異なる粒度で特徴表現を強化すること。
• T1WI、T2WI、FLAIR、DWIを含む多様な疾患とMRIシーケンスをカバーする24,770組の脳MRI画像レポート対から構成される大規模データセットを用いてモデルを学習すること。
• 複数レベルでのビジョンと言語特徴の統合を可能とする、ビジョンおよびテキストエンコーダーの二重ブランチアーキテクチャに、クロスアテンション機構を組み込むこと。

実験結果

リサーチクエスチョン

• RQ1 階層的知識強化事前学習フレームワークは、脳MRI診断における多様な疾患にわたる一般化能力を向上させ得るか？
• RQ2 グローバル整合性と比較して、微細でモダリティ別に特化したレポート整合性は、視覚言語表現学習の効率性と正確性を向上させるか？
• RQ3 CVPモジュールで疾患記述をクエリとして使用することで、診断性能と解釈可能性はどの程度向上するか？
• RQ4 SOTA手法と比較して、オープンボリュームおよびドメインシフト設定下でのUniBrainの性能はいかがなものか？
• RQ5 特定の疾患タイプにおいて、UniBrainは専門放射線科医と同等の診断性能を達成できるか？

主な発見

• UniBrainはBraTS2019データセットでマクロF1スコア63.27%を達成し、すべてのSOTA手法を上回り、複数の指標でF1およびAUCが1.5–2.5%向上した。
• 自社データセット（著しいクラス不均衡を伴う）でもマクロF1が62.27%、AUCが90.71%を達成し、強力なロバストネスと一般化能力を示した。
• CVPモジュールを削除するとF1スコアが10ポイント以上低下（63.27% → 51.73%）し、疾患固有特徴学習におけるその重要性を確認した。
• CVPで疾患記述をクエリとして使用した場合、単に疾患名のみを用いる場合と比較して、F1スコアが約11.5ポイント向上した。
• アブレーションスタディでは、グローバルおよびモダリティ別レポート・画像整合性をすべて削除した場合、F1スコアは63.27%から60.61%に低下し、階層的知識強化の重要性を裏付けた。
• 特定の疾患タイプにおいて、UniBrainの診断性能は専門放射線科医と同等であり、特に膠腫瘍および脳梗塞分類において顕著であった。