The Importance of Being Scalable: Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains

要約

スケーリングは、機械学習におけるモデルのパフォーマンスと一般化を向上させる上で重要です。
これには、モデルのサイズや入力データの増加に伴ってモデルのパフォーマンスがどのように変化するか、またこの成長をサポートするために計算リソースがどのように効率的に利用されるかが関係します。
他の分野での成功にもかかわらず、ニューラルネットワーク原子間ポテンシャル（NNIP）のスケーリングに関する研究は依然として限られています。
NNIP は、非経験的な量子力学計算の代理モデルとして機能します。
ここでの主要なパラダイムは、回転等分散などの多くの物理領域制約をモデルに組み込むことです。
これらの複雑な制約が NNIP のスケーリング能力を阻害し、長期的にはパフォーマンスの停滞につながる可能性が高いと私たちは主張します。
この研究では、別のアプローチを採用し、NNIP スケーリング戦略を体系的に研究することから始めます。
私たちの調査結果は、アテンション メカニズムによるモデルのスケーリングが効率的であり、モデルの表現力を向上させることを示しています。
これらの洞察は、スケーラビリティを考慮して設計された NNIP アーキテクチャである Efficiently Scaled Attendant Interatomic Potential (EScAIP) を開発する動機となります。
EScAIP は、グラフ ニューラル ネットワーク内のマルチヘッド セルフ アテンションの定式化を活用し、近隣レベルの表現に注意を適用します。
高度に最適化されたアテンション GPU カーネルで実装された EScAIP は、既存の NNIP と比較して効率が大幅に向上し、推論が少なくとも 10 倍高速になり、メモリ使用量が 5 分の 1 に削減されます。
EScAIP はまた、触媒 (OC20 および OC22)、分子 (SPICE)、材料 (MPTrj) を含む幅広いデータセットに対して最先端のパフォーマンスを実現します。
私たちのアプローチは特定のモデルではなく哲学として考えられるべきであり、スケーリングを通じてより優れた表現力を実現し、増加した計算リソースとトレーニング データで効率的にスケーリングし続ける汎用 NNIP を開発するための概念実証を表すものであることを強調します。
。

要約(オリジナル)

Scaling has been critical in improving model performance and generalization in machine learning. It involves how a model’s performance changes with increases in model size or input data, as well as how efficiently computational resources are utilized to support this growth. Despite successes in other areas, the study of scaling in Neural Network Interatomic Potentials (NNIPs) remains limited. NNIPs act as surrogate models for ab initio quantum mechanical calculations. The dominant paradigm here is to incorporate many physical domain constraints into the model, such as rotational equivariance. We contend that these complex constraints inhibit the scaling ability of NNIPs, and are likely to lead to performance plateaus in the long run. In this work, we take an alternative approach and start by systematically studying NNIP scaling strategies. Our findings indicate that scaling the model through attention mechanisms is efficient and improves model expressivity. These insights motivate us to develop an NNIP architecture designed for scalability: the Efficiently Scaled Attention Interatomic Potential (EScAIP). EScAIP leverages a multi-head self-attention formulation within graph neural networks, applying attention at the neighbor-level representations. Implemented with highly-optimized attention GPU kernels, EScAIP achieves substantial gains in efficiency–at least 10x faster inference, 5x less memory usage–compared to existing NNIPs. EScAIP also achieves state-of-the-art performance on a wide range of datasets including catalysts (OC20 and OC22), molecules (SPICE), and materials (MPTrj). We emphasize that our approach should be thought of as a philosophy rather than a specific model, representing a proof-of-concept for developing general-purpose NNIPs that achieve better expressivity through scaling, and continue to scale efficiently with increased computational resources and training data.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google