[論文レビュー] Accelerating Molecular Graph Neural Networks via Knowledge Distillation

Filip Ekström Kelvinius, Dimitar Georgiev+2 more|arXiv (Cornell University)|Jun 26, 2023

Machine Learning in Materials Science被引用数 4

ひとこと要約

本稿では、構造的変更なしに分子グラフニューラルネットワーク（GNN）の推論速度を向上させるために、ノード同士、エッジ同士、エッジからノード、ベクトル同士の知識蒸留（KD）戦略——ノード対ノード（n2n）、エッジ対エッジ（e2e）、エッジ対ノード（e2n）、ベクトル対ベクトル（v2v）——を提案する。大規模で高精度な教師GNNから隠れ表現を蒸留することで、学生モデルは推論速度を維持したままエネルギー予測で最大96.7%、力の予測で62.5%の教師モデル性能を達成する。

ABSTRACT

Recent advances in graph neural networks (GNNs) have enabled more comprehensive modeling of molecules and molecular systems, thereby enhancing the precision of molecular property prediction and molecular simulations. Nonetheless, as the field has been progressing to bigger and more complex architectures, state-of-the-art GNNs have become largely prohibitive for many large-scale applications. In this paper, we explore the utility of knowledge distillation (KD) for accelerating molecular GNNs. To this end, we devise KD strategies that facilitate the distillation of hidden representations in directional and equivariant GNNs, and evaluate their performance on the regression task of energy and force prediction. We validate our protocols across different teacher-student configurations and datasets, and demonstrate that they can consistently boost the predictive accuracy of student models without any modifications to their architecture. Moreover, we conduct comprehensive optimization of various components of our framework, and investigate the potential of data augmentation to further enhance performance. All in all, we manage to close the gap in predictive accuracy between teacher and student models by as much as 96.7% and 62.5% for energy and force prediction respectively, while fully preserving the inference throughput of the more lightweight models.

研究の動機と目的

• 最新の分子GNNの計算コストの増大に対処し、大規模シミュレーションへの応用を制限する要因を軽減すること。
• 構造的変更なしに軽量GNNの推論速度と精度のトレードオフを改善すること。
• 分子シミュレーションで用いられる方向性および等長性を持つGNNに特化したKD戦略の開発。
• OC20-2MおよびCOLLを含む多様な教師-学生設定とデータセットにおいて、KDの有効性を評価すること。
• データ拡張およびトレーニングプロトコルの最適化を検討し、学生モデルの性能をさらに向上させること。

提案手法

• 方向性および等長性を持つGNN向けに、ノード対ノード（n2n）、エッジ対エッジ（e2e）、エッジ対ノード（e2n）、ベクトル対ベクトル（v2v）の4つの新しいKD戦略を提案する。
• 事前学習済みで高精度な教師GNNを用い、エネルギー、力、隠れ状態などのターゲット表現を生成して知識蒸留に用いる。
• 標準的な回帰損失と隠れ表現からのKD損失を組み合わせた多成分損失関数を用いて学生GNNを訓練する。
• スカラーおよびベクトル特徴を含むノードレベルおよびエッジレベルの表現に蒸留を適用し、幾何的・構造的情報を保持する。
• データ混合比や実データと合成データ間の損失重み付けなどのトレーニングハイパーパrameterを最適化する。
• ランダムラッティングや合成データ生成（d1M）によるデータ拡張を検討し、学生モデルの汎化性能を向上させる。

Figure 1: Using knowledge distillation, we manage to significantly boost the predictive accuracy of different student models on the OC20-2M [ 17 ] and COLL [ 6 ] datasets while fully preserving their inference throughput.

実験結果

リサーチクエスチョン

• RQ1 知識蒸留は、推論速度を維持したまま分子GNNを効果的に高速化できるか？
• RQ2 ノード対ノード（n2n）、エッジ対エッジ（e2e）、エッジ対ノード（e2n）、ベクトル対ベクトル（v2v）といった特徴ベースの蒸留戦略が、エネルギーおよび力の予測における学生モデル性能にどのように影響するか？
• RQ3 蒸留によって、大規模な教師モデルと小規模な学生モデルの間の精度格差はどの程度縮まるか？
• RQ4 ランダムラッティングや合成データといったデータ拡張は、学生モデルの性能をさらに向上させるか？
• RQ5 データ混合比や損失重み付けといったハイパーパrameterは、蒸留モデルの最終的な性能にどのように影響するか？

主な発見

• 提案されたKD戦略は、構造的変更なしに多数のデータセットおよび教師-学生設定において一貫して学生モデルの性能を向上させる。
• エネルギー予測では教師モデルの96.7%、力の予測では62.5%の精度にまで到達し、顕著に精度格差を縮小する。
• 学生モデルが軽量かつ推論効率的であるため、推論スループットは完全に保持されている。
• ランダムラッティングによるデータ拡張は顕著な向上をもたらさず、勾配ベースの摂動はランダムノイズより効果が低かった。
• 合成データ生成（d1M）は、実データとの共同学習をベースラインとして用いた場合、他の分野とは異なり、性能向上が著しく見られなかった。
• 包括的なアブレーションスタディにより、個々のKDコンponentsの有効性が確認され、特にベクトル対ベクトル（v2v）およびノード対ノード（n2n）戦略が顕著な向上を示した。

Figure 2: Similarity analysis between the node features of SchNet, PaiNN and GemNet-OC using CKA (averaged over $n=987$ nodes).