Accelerating Molecular Graph Neural Networks via Knowledge Distillation

要約

グラフ ニューラル ネットワーク (GNN) の最近の進歩により、数分の 1 の計算コストで、従来のゴールドスタンダード手法と同等の精度で分子シミュレーションが可能になりました。
それにもかかわらず、この分野がより大規模でより複雑なアーキテクチャに進歩するにつれて、最先端の GNN は多くの大規模アプリケーションには非常に法外なものになりました。
この論文では、分子 GNN を加速するための知識蒸留 (KD) の有用性を初めて調査します。
この目的のために、方向性と等変性の GNN における隠れた表現の蒸留を容易にする KD 戦略を考案し、エネルギーと力の予測の回帰タスクでのパフォーマンスを評価します。
私たちはさまざまな教師と生徒の構成にわたってプロトコルを検証し、アーキテクチャを変更することなく生徒モデルの予測精度を向上できることを実証しました。
また、フレームワークのさまざまなコンポーネントの包括的な最適化を実行し、パフォーマンスをさらに向上させるためのデータ拡張の可能性を調査します。
全体として、推論時の追加コストをゼロにすることで、GemNet-OC と PaiNN のようなモデル間の予測精度の差を 59% も埋めることに成功しました。

要約(オリジナル)

Recent advances in graph neural networks (GNNs) have allowed molecular simulations with accuracy on par with conventional gold-standard methods at a fraction of the computational cost. 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, for the first time, 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 demonstrate that they can boost the predictive accuracy of student models without altering their architecture. We also 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 as much as 59% of the gap in predictive accuracy between models like GemNet-OC and PaiNN with zero additional cost at inference.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google