[論文レビュー] Bayesian Graph Neural Networks for Molecular Property Prediction

George S. Lamb, Brooks Paige|arXiv (Cornell University)|Nov 25, 2020

Machine Learning in Materials Science参考文献 30被引用数 9

ひとこと要約

本稿では、QM9データセット上で有向MPNNのメッセージパッシング層およびリードアウト層における不確実性評価を用いて、分子性質予測のためのベイジアングラフニューラルネットワーク（GNN）を提案する。全ベイジアンGNNは、標準的なディープラーニングベースラインと比較して、予測精度、キャリブレーション、および下流の分子探索タスクにおける性能を顕著に向上させることを示している。

ABSTRACT

Graph neural networks for molecular property prediction are frequently underspecified by data and fail to generalise to new scaffolds at test time. A potential solution is Bayesian learning, which can capture our uncertainty in the model parameters. This study benchmarks a set of Bayesian methods applied to a directed MPNN, using the QM9 regression dataset. We find that capturing uncertainty in both readout and message passing parameters yields enhanced predictive accuracy, calibration, and performance on a downstream molecular search task.

研究の動機と目的

• 訓練データにおける過学習による、分布外の分子スケルトンへの一般化性能の低さを是正する。
• 特に回帰設定において、モデルのキャリブレーションと不確実性評価を向上させる。
• GNNのメッセージパッシングおよびリードアウト部の両方におけるベイジアン推論の影響を評価する。
• 実世界の分子探索タスクにおけるベイジアン手法のベンチマークを実施し、実用的有用性を評価する。
• 再現可能性およびさらなる研究を促進するため、オープンソースのPyTorchコードを提供する。

提案手法

• 分子表現学習のためのベースとなるGNNアーキテクチャとして、有向メッセージパッシングニューラルネットワーク（D-MPNN）を採用する。
• 8種類のベイジアン推論手法（MAP、GP、MCドロップアウト（DropR、DropA）、SWAG、SGLD、BBP、DUN）を用いて不確実性推定を実施する。
• 特にアレアトリック不確実性が明示的にモデル化されていない場合に備え、後処理のt分布尤度を用いてキャリブレーションを改善する。
• 事後分布の近似からサンプリングすることでアンサンブルベースの予測分布を実装し、不確実性を推定する。
• 変分推論またはMCMCベースの手法を用いて、モデルパラメータの事後分布を近似する形で、エンドツーエンドでモデルを学習する。
• 信頼性図を用いて不確実性を評価し、12個のQM9回帰タスク全体で平均ランクと誤差の面積（miscalibration area）を主な指標として用いる。

実験結果

リサーチクエスチョン

• RQ1 メッセージパッシング層およびリードアウト層におけるベイジアン学習が、分子性質予測における予測精度とキャリブレーションを向上させるか？
• RQ2 異なるベイジアン推論手法は、QM9回帰タスクにおける不確実性評価と性能の面でどのように比較されるか？
• RQ3 ベイジアンGNNは、下流の分子生成または探索タスクにおける性能を向上させられるか？
• RQ4 アレアトリック不確実性をモデル化した場合、ガウス分布とt分布尤度の使用がキャリブレーションに与える影響は何か？
• RQ5 不確実性評価は、未観測の分子スケルトンへの一般化にどのように影響を与えるか？

主な発見

• 特にメッセージパッシング層およびリードアウト層の両方に不確実性を組み込んだ全ベイジアンGNNは、QM9で平均ランク1.95 ± 0.40を達成し、標準的なMAP学習（4.08 ± 0.16）を顕著に上回った。
• 全GNNにMCドロップアウトを適用するDropAは、77.52 ± 0.77%の探索スコアを達成し、MAP（72.22 ± 0.57%）および他のベイジアン手法を上回った。
• 後処理のt分布尤度の使用により、ガウス尤度で観察された病理的なくどみ（underconfidence）が軽減された。
• ベイジアンリードアウト（例：GP、DropR）を備えたモデルは、ベイジアンメッセージパッシングのみのモデルよりも優れたキャリブレーションを示し、予測信頼性においてリードアウトの不確実性がより重要であることを示唆した。
• SGLDおよびBBPは高い不確実性推定値を示したが、予測精度は低く、MCMCおよび変分推論手法がこの設定では必ずしも一般化性能を発揮しない可能性を示唆した。
• DUN手法（対角＋低ランク共分散）は、平均ランク（4.40 ± 0.25）および探索スコア（76.02 ± 1.09）で優れた性能を示し、不確実性推定の堅牢性を示した。