[論文レビュー] On Graph Neural Network Ensembles for Large-Scale Molecular Property Prediction

Edward Elson Kosasih, Joaquín Cabezas+6 more|arXiv (Cornell University)|Jun 29, 2021

Machine Learning in Materials Science被引用数 7

ひとこと要約

本稿では、3つのグラフニューラルネットワーク（GNN）モデル—GIN-Virtual、ベイジアンGNN、DiffPool強化GNN—を複数回の異なる初期化で訓練したアンサンブルを提案し、大規模な分子性質予測を改善する。アンサンブルはベースライン比で相対的に7.6%の向上を達成し、PCQM4M-LSCデータセットにおいてMAEを0.1290まで低下させた。また、予測の不確実性と予測誤差との間に強い相関関係（ピアソン相関係数R = 0.5181）が確認され、アクティブラーニングへの応用が可能であることが示された。

ABSTRACT

In order to advance large-scale graph machine learning, the Open Graph Benchmark Large Scale Challenge (OGB-LSC) was proposed at the KDD Cup 2021. The PCQM4M-LSC dataset defines a molecular HOMO-LUMO property prediction task on about 3.8M graphs. In this short paper, we show our current work-in-progress solution which builds an ensemble of three graph neural networks models based on GIN, Bayesian Neural Networks and DiffPool. Our approach outperforms the provided baseline by 7.6%. Moreover, using uncertainty in our ensemble's prediction, we can identify molecules whose HOMO-LUMO gaps are harder to predict (with Pearson's correlation of 0.5181). We anticipate that this will facilitate active learning.

研究の動機と目的

• OGB-LSCのPCQM4M-LSCデータセット（380万個の分子を含む）において、大規模な分子性質予測を改善すること。
• 現実的で大規模なデータセットにおけるグラフ機械学習のスケーラビリティとパフォーマンスの課題に対処すること。
• アンサンブル予測における不確実性が、アクティブラーニングのための高い予測誤差領域を示す可能性があるかを調査すること。
• 多様なGNNアーキテクチャとトレーニングランを組み合わせたスケーラブルなアンサンブルフレームワークを構築し、回帰性能を向上させること。

提案手法

• GIN-Virtual、ベイジアンGNN（BNN）、DiffPool強化GNNの3つのGNN弱学習器を、それぞれ異なる初期化から3回ずつ訓練する。
• 各モデルの複数回の予測平均値を用いて個々のアンサンブル（例：GIN-Virtual-アンサンブル）を形成する。
• 全モデルと全ランの予測を組み合わせ、重み付き平均（9つの異なる試行）を用いて最終アンサンブル（アンサンブル-すべて）を構築する。
• GNNの読み出し層にBayes-by-Backpropagation（BBP）を適用し、ベイジアン線形層を介して不確実性推定を導入する。
• DiffPoolを用いてグラフプーリングを実装し、グラフを5ノードに粗視化することで、構造的情報を保持した表現を向上させる。
• 不確実性を9つのアンサンブル試行における予測の標準偏差として測定し、モデルの信頼性の代理指標として用いる。

実験結果

リサーチクエスチョン

• RQ1 多様なGNNアーキテクチャをアンサンブル化することで、個々のモデルを上回る回帰性能が大規模な分子グラフで達成可能か？
• RQ2 アンサンブルの不一致から推定される不確実性は、HOMO-LUMOギャップ予測における実際の予測誤差と相関するか？
• RQ3 提案されたアンサンブル手法は、PCQM4M-LSCデータセットにおいてOGB-LSCベースラインを著しく上回る性能を発揮できるか？
• RQ4 不確実性推定は、予測が難しい分子を特定するためのアクティブラーニングにどの程度有効に機能するか？

主な発見

• アンサンブル-すべてモデルは、検証セットにおいて平均絶対誤差（MAE）が0.1290を達成し、ベースラインのMAE 0.1396に対して相対的に7.6%の改善を示した。
• 不確実性推定値（9つのアンサンブル試行における予測の標準偏差）は、実際の予測誤差とピアソン相関係数0.5181の強い相関を示した。
• GIN-Virtual-BNN-アンサンブル変種が、個々のアンサンブルタイプの中で最小の検証MAE 0.1329を達成した。
• アンサンブル-すべてモデルは、GNN-Virtual（コンテストで最高性能を示した単一モデル）を含むすべての個別モデルおよびベースラインを上回った。
• 不確実性と誤差の間の相関関係は、不確実性推定値がより予測が難しい分子を特定できることを示しており、将来的なアクティブラーニング応用を支援する。
• モデルは高メモリのクラウドGPUインスタンスを用いて300万個の分子で学習されたことから、大規模な分子データセットへのスケーラビリティが実証された。