[論文レビュー] Transfer learning for atomistic simulations using GNNs and kernel mean embeddings

Isak Falk, Leo H. Bonati+3 more|arXiv (Cornell University)|Jun 2, 2023

Machine Learning in Materials Science被引用数 5

ひとこと要約

本論文では、OC20データセットで事前に訓練されたグラフニューラルネットワーク（GNN）表現とカーネル平均埋め込みを組み合わせることで、原子系のポテンシャルエネルギー表面をモデル化するトランスファーラーニングフレームワークであるMean Embedding Kernel Ridge Regression（MEKRR）を提案する。化学種情報を新しいカーネル関数に組み込むことで、触媒系において最先端の性能を達成し、単独のGNNやリッジ回帰よりも顕著に優れており、分布外データではテスト誤差が最大70%まで低減された。

ABSTRACT

Interatomic potentials learned using machine learning methods have been successfully applied to atomistic simulations. However, accurate models require large training datasets, while generating reference calculations is computationally demanding. To bypass this difficulty, we propose a transfer learning algorithm that leverages the ability of graph neural networks (GNNs) to represent chemical environments together with kernel mean embeddings. We extract a feature map from GNNs pre-trained on the OC20 dataset and use it to learn the potential energy surface from system-specific datasets of catalytic processes. Our method is further enhanced by incorporating into the kernel the chemical species information, resulting in improved performance and interpretability. We test our approach on a series of realistic datasets of increasing complexity, showing excellent generalization and transferability performance, and improving on methods that rely on GNNs or ridge regression alone, as well as similar fine-tuning approaches.

研究の動機と目的

• 原子系シミュレーションにおける高精度な原子間ポテンシャルを訓練するための高いデータおよび計算コストの課題に対処すること。
• 事前に訓練されたGNN表現を用いて、複雑な触媒系においてデータ効率的かつ高速なトランスファーラーニングを可能にすること。
• 化学種を記述する物理的インフォームドなカーネル関数とGNN特徴を組み合わせることで、一般化性能とトランスファー可能性を向上させること。
• 物理的対称性（回転・並進および置換対称性）を保持しつつ、解釈可能性と性能向上を実現する手法を開発すること。
• 事前学習分布外の、ますます複雑なデータセットにおいて、本手法の有効性を示すこと。

提案手法

• 触媒系のDFTで最適化された約2.6億個の構造を含む大規模なOC20データセット上でGNN（例：SchNetまたはSCN）を事前学習する。
• 回転・並進不変性を保持する固定された事前学習済みGNN特徴を、原子環境の表現として抽出する。
• 置換対称性を強制し、学習可能な重みを介して化学種情報を組み込むカーネル平均埋め込みに基づくカーネル関数を定義する。
• 提案されたカーネルを用いたカーネルリッジ回帰（KRR）により、小規模なシステム固有のデータセットからポテンシャルエネルギー表面を回帰する。
• 交差検証を用いてカーネル重みを最適化し、特に表面触媒における吸着原子のような低頻度種を強調する。
• カーネル行列を用いたスペクトルクラスタリングを実行し、結合の切断といった反応的イベントを検出することで、回帰を超えた解釈可能性を示す。

実験結果

リサーチクエスチョン

• RQ1 事前学習済みGNN表現とカーネル平均埋め込みを組み合わせることで、原子系におけるポテンシャルエネルギー表面モデリングに有効なトランスファーラーニングが可能になるか？
• RQ2 カーネル関数に化学種情報を組み込むことで、モデルの性能と解釈可能性にどのような影響を与えるか？
• RQ3 提案手法が、事前学習時に見られなかった分布外データセットにどの程度一般化できるか？
• RQ4 ラベルなしで、カーネルベースの類似性測度が結合の切断といった物理的イベントを検出できるか？
• RQ5 データ効率性と精度の観点から、微調整されたGNNや単独のカーネル手法と比較して、MEKRRはどのように差をつけるか？

主な発見

• SchNet特徴を用いたMEKRRは、D1→D2の転送において、教師あり学習済みSchNetと比較してテスト誤差を68%低減し、8.0 meV/atomを達成した。
• 最も複雑なデータセット（D4）では、SCN特徴を用いたMEKRRがテスト誤差55 meV/atomを達成し、微調整済みSchNet-FT（116 meV/atom）を上回った。
• 本手法は、事前学習セットに存在しなかったN-N結合の切断を伴う構造に対しても、強力な一般化性能を示した。
• マルチ重みカーネル（α=1）は、N-N距離といった物理的観測量とより密接に相関しており、反応の主要なイベントに対する感度が向上していることが示された。
• カーネル行列を用いたスペクトルクラスタリングにより、結合切断軌道における反応物および生成物状態が正しく同定され、カーネルの物理的解釈可能性が裏付けられた。
• 化学種に配慮したカーネル重みの導入により、特に吸着原子のような低頻度種において性能が向上した。