[論文レビュー] GPS++: An Optimised Hybrid MPNN/Transformer for Molecular Property Prediction

Dominic Masters, Josef Dean+8 more|arXiv (Cornell University)|Nov 18, 2022

Machine Learning in Materials Science被引用数 10

ひとこと要約

GPS++ は、PCQM4Mv2 データセットにおける分子性質予測を最適化したハイブリッド MPNN/Transformer モデルであり、3次元原子座標とノイズ除去損失を統合することで性能を向上させた。テストセットの平均絶対誤差は 0.0719 を達成し、OGB-LSC 2022 チャレンジで1位を獲得した。Graphcore IPU を用いた加速により、112 モデルのアンサンブルを用いて 1 時間 32 分で推論が完了した。

ABSTRACT

This technical report presents GPS++, the first-place solution to the Open Graph Benchmark Large-Scale Challenge (OGB-LSC 2022) for the PCQM4Mv2 molecular property prediction task. Our approach implements several key principles from the prior literature. At its core our GPS++ method is a hybrid MPNN/Transformer model that incorporates 3D atom positions and an auxiliary denoising task. The effectiveness of GPS++ is demonstrated by achieving 0.0719 mean absolute error on the independent test-challenge PCQM4Mv2 split. Thanks to Graphcore IPU acceleration, GPS++ scales to deep architectures (16 layers), training at 3 minutes per epoch, and large ensemble (112 models), completing the final predictions in 1 hour 32 minutes, well under the 4 hour inference budget allocated. Our implementation is publicly available at: https://github.com/graphcore/ogb-lsc-pcqm4mv2.

研究の動機と目的

• 大規模な PCQM4Mv2 データセットにおける分子性質予測のための、非常に効率的で高精度なモデルの開発。
• 3次元原子座標とインダクティブバイアスを活用し、グラフニューラルネットワークの一般化性能を向上させ、過剰平滑化を低減すること。
• ハードウェア加速を用いたトレーニングと推論により、厳密な推論時間制約（4時間予算）の下で最先端の性能を達成すること。
• ハイブリッド MPNN/Transformer アーキテクチャが、はるかに少ないパラメータ数で純粋な Transformer モデルと同等またはそれを上回る性能を発揮できることを示すこと。
• 予測の堅牢性と一般化性能を最大化するために、112 モデルの多様なアンサンブルを構築すること。

提案手法

• GPS++ は、局所的なインダクティブバイアスとグローバルな情報伝達の両立を図るため、深層メッセージパッシングニューラルネットワーク（MPNN）とバイアス付き自己注意機構を統合している。
• 距離特徴を注意機構に統合するグループ化された入力マスキング戦略を用いて、3次元原子座標をモデルに組み込んでいる。
• 深層アーキテクチャにおける表現学習の向上と過剰平滑化の緩和を図るため、事前学習段階でノイズ除去損失を適用している。
• MPNN レイヤーが局所的なメッセージパッシングを処理し、注意レイヤーが長距離依存性を扱うハイブリッドアーキテクチャを採用しており、注意は近隣原子にバイアスがかかる。
• トレーニングと推論は Graphcore の IPU ハードウェアで加速されており、1エポックあたり3分、112 モデルのアンサンブルで推論が1時間32分で完了した。
• 入力特徴、ドロップアウト率、マスキング戦略の変更により、多様なモデルのアンサンブルを構築し、一般化性能と耐障害性を向上させた。

実験結果

リサーチクエスチョン

• RQ1 ハイブリッド MPNN/Transformer アーキテクチャは、パラメータ数を少なくした状態で、純粋な Transformer モデルを上回る性能を発揮できるか？
• RQ2 グループ化された入力マスキングを用いた 3次元原子座標の統合は、モデル性能と一般化性能の向上にどの程度効果的か？
• RQ3 ノイズ除去事前学習目的関数は、分子グラフにおける深層グラフニューラルネットワークの過剰平滑化をどの程度低減できるか？
• RQ4 IPU を用いたハードウェア加速により、特に深層モデルや大規模アンサンブルにおいて、厳密な時間制約の下でスケーラブルなトレーニングと推論が可能か？
• RQ5 ハイパーパramータの変更により得られるモデルの多様性——入力特徴、ドロップアウト、マスキング戦略の変更——は、PCQM4Mv2 のテスト・チャレンジスプリットにおける最終的予測性能を顕著に向上させるか？

主な発見

• GPS++ は PCQM4Mv2 データセットにおいて、テスト・チャレンジスプリットの平均絶対誤差（MAE）を 0.0719 に達成し、OGB-LSC 2022 チャレンジで上位3位以内を獲得した。
• 最終的な 112 モデルのアンサンブルは、1つの IPU と AMD EPYC CPU を用いて、4時間の予算を大幅に下回る 1 時間 32 分で推論を完了した。
• パラメータ数が 64% に抑えられた（4430万対 6900万）にもかかわらず、Luo ら（2022）の最先端の Transformer と同等の性能を達成した。
• グループ化マスキングによる 3次元距離特徴の統合は、注意機構の有効性を顕著に向上させた。アブレーション実験では、3次元特徴を省いた場合、性能向上がほとんど見られなかった。
• 入力特徴、ドロップアウト、マスキング戦略を変更した多様な構成のアンサンブルにより、MAE は単一モデルの 0.0755 からプロキシセットの 0.0722 へ、最終テストセットでは 0.0719 へと低下した。
• IPU ハードウェア上で 1 エポックあたり 3 分でトレーニングが可能となり、深層（16層）アーキテクチャの迅速な反復と 24 時間以内のトレーニングが可能になった。