[論文レビュー] Reliable Graph Neural Networks for Drug Discovery Under Distributional Shift

Kehang Han, Balaji Lakshminarayanan+1 more|arXiv (Cornell University)|Nov 25, 2021

Computational Drug Discovery Methods参考文献 29被引用数 15

ひとこと要約

本稿では、スペクトル正規化と距離に敏感なガウス過程分類器を統合することで、分布シフト下での薬物発見における信頼性の高いグラフニューラルネットワーク、GNN-SNGPを提案する。新しいCardioToxベンチマークと3つの既存のデータセットを用いた検証により、特に分布外の分子において過信された誤予測が顕著に減少するが、高い精度を維持するという、優れたロバスト性を示している。

ABSTRACT

The concern of overconfident mis-predictions under distributional shift demands extensive reliability research on Graph Neural Networks used in critical tasks in drug discovery. Here we first introduce CardioTox, a real-world benchmark on drug cardio-toxicity to facilitate such efforts. Our exploratory study shows overconfident mis-predictions are often distant from training data. That leads us to develop distance-aware GNNs: GNN-SNGP. Through evaluation on CardioTox and three established benchmarks, we demonstrate GNN-SNGP's effectiveness in increasing distance-awareness, reducing overconfident mis-predictions and making better calibrated predictions without sacrificing accuracy performance. Our ablation study further reveals the representation learned by GNN-SNGP improves distance-preservation over its base architecture and is one major factor for improvements.

研究の動機と目的

• GNNが薬物発見に用いられる際の過信された誤予測という、信頼性の重要な課題に取り組むこと。特に分布シフト下での発現を想定する。
• 過信された誤予測は、通常、学習データから構造的に離れた分子で発生しており、距離に敏感でないことが示唆される。
• 表現における距離保存性を向上させるとともに、不確実性に敏感な予測を組み込むことで、モデルの信頼性を高める、GNN-SNGPと呼ばれる新規なGNNアーキテクチャを開発すること。
• 実世界のベンチマークとして、分布内と分布シフトされたテストセットを備えたCardioToxを導入し、現実の薬物発見状況下でのモデル信頼性を評価すること。
• 距離に敏感な表現の向上が、より良好なキャリブレーションと過信された偽陰性の低減につながることを示し、安全な薬物開発にとって不可欠である。

提案手法

• スペクトル正規化（SNGP）とランダム特徴量ガウス過程（GP）層を組み合わせた、エンドツーエンドで学習可能なGNN、GNN-SNGPを提案する。不確実性推定に用いる。
• GNNのメッセージパッシング層にスペクトル正規化を適用し、潜在空間における幾何的距離を保存することで、距離に敏感な表現を向上させる。
• 最終層の全結合層を、学習データからの距離に基づいて予測不確実性をモデル化するGP層に置き換える。これにより、分布外の分子における過信を低減する。
• 分子の距離を定量化し、評価用の距離ベースのデータ分割を定義するために、フィンガープrintベースのグラフ距離（RogersとHahn, 2010）を用いる。
• 標準的なGNN損失に不確実性正則化を組み合わせて学習することで、精度を損なわずに信頼性の高い不確実性推定を可能にする。
• アブレーションスタディを実施し、GP層と距離保存型表現が全体のロバスト性に与える寄与を分離して評価する。

実験結果

リサーチクエスチョン

• RQ1 GNNが薬物毒性予測に用いられる際、過信された誤予測は、学習データから構造的に離れた分子でどれほど発生するのか？
• RQ2 GNN表現における距離に敏感な性質を向上させることで、薬物発見における分布シフト下での過信された誤予測を低減できるか？
• RQ3 GNN-SNGPは、ベースラインGNNや他の不確実性推定手法と比較して、分布外データにおけるキャリブレーションとロバスト性において、どのように優れているか？
• RQ4 距離保存型表現とGPベースの分類器の組み合わせは、単体で用いた場合よりも、より良好な不確実性推定を達成できるか？
• RQ5 提案手法は、心毒性を越えて、他のグラフベースの薬物発見タスクにも一般化可能か？

主な発見

• GNN-SNGPは、テスト-OOD1スプリットで最高のAUROC（0.935）と最小のECE（0.002）を達成し、GNNベースラインや他の変種と比較して、キャリブレーションとロバスト性において顕著に優れている。
• 最も挑戦的なテスト-OOD2スプリットでは、GNN-SNGPは過信された偽陰性（OFNs%）を0.00%にまで低減したのに対し、FPs + Denseベースラインでは5.71%であった。
• アブレーションスタディにより、距離保存型GNNエンコーダーとGP層の両方が信頼性の向上に寄与していることが確認され、特にGP層が過信の低減において顕著に有効であることが示された。
• GNN-SNGPは、OODスプリットにおいてECE（0.021まで低下）とブライアースコア（0.208）が改善され、よりキャリブレーションされた不確実性推定であることが示された。
• 過信された予測を低減しながらも、高い精度（AUROC > 0.89）を維持しており、信頼性と性能が相反するものではないことが実証された。
• 本手法は一般化性に優れている：GNN-SNGPは、molHIV、BBBP、BACEデータセットでも一貫した改善を示し、心毒性を越えて広く適用可能であることが確認された。