[論文レビュー] Kolmogorov--Arnold networks in molecular dynamics

Yuki Nagai, Masahiko Okumura|arXiv (Cornell University)|Jul 25, 2024

History and advancements in chemistry被引用数 4

ひとこと要約

本稿では、Kolmogorov–Arnold Network（KAN）を分子動力学に統合し、原子間ポテンシャルの効率を向上させる手法を提案する。KANの非線形関数を用いて人工ニューラルネットワーク（ANN）記述子を再定義し、線形または3次スプライン補間を適用することで、ポテンシャルおよび力の評価にかかる計算コストを削減しながらも、精度を維持する。Lennard-Jonesおよび埋め込み原子モデル（EAM）ポテンシャルを用いた実験でその有効性が示された。

ABSTRACT

We explore the integration of Kolmogorov Networks (KANs) into molecular dynamics (MD) simulations to improve interatomic potentials. We propose that widely used potentials, such as the Lennard-Jones (LJ) potential, the embedded atom model (EAM), and artificial neural network (ANN) potentials, can be interpreted within the KAN framework. Specifically, we demonstrate that the descriptors for ANN potentials, typically constructed using polynomials, can be redefined using KAN's non-linear functions. By employing linear or cubic spline interpolations for these KAN functions, we show that the computational cost of evaluating ANN potentials and their derivatives is reduced.

研究の動機と目的

• 機械学習分子動力学（MLMD）シミュレーションの効率を向上させるために、一般的な原子間ポテンシャルをKANフレームワーク内で再解釈すること。
• 分子動力学における人工ニューラルネットワーク（ANN）ポテンシャルおよびその微分の評価にかかる計算コストを低減すること。
• KANの非線形活性化関数を活用し、多項式ベースの記述子とは異なる代替手段としてKANをANNポテンシャルに用いることの可能性を検討すること。
• KAN関数の線形または3次スプライン補間により、モデル性能を損なわせることなく、記述子の高速かつ高精度な評価が可能であることを示すこと。
• ænetなどのオープンソースMLMDパッケージへのKANベース記述子の統合を促進する基盤を構築すること。

提案手法

• Kolmogorov–Arnold表現定理の枠組み内で、広く用いられている原子間ポテンシャル（Lennard-Jones、埋め込み原子モデル（EAM）、ANNポテンシャル）を再定式化する。
• ポテンシャルエネルギー面を、非線形関数（KAN層）の重ね合わせとして表現し、活性化関数は学習中に最適化する。
• ANNポテンシャルにおける従来の多項式ベースの記述子（例：チェビシェフ多項式）を、径方向および角度方向の非線形関数を用いて得たKAN由来の記述子に置き換える。
• 線形または3次スプライン補間を用いてKAN活性化関数を近似し、1回の評価あたりの計算コストをO(𝑁)からO(log 𝑁)に削減する。
• 補間ライブラリ（例：Interpolation.jl）を用いて、補間された記述子を効率的に計算し、誤差は残差ノルムおよび平均絶対誤差（MAE）で測定する。
• テストデータセットを用いて性能を検証し、補間有無のKAN記述子のMAEを、基準となるチェビシェフ多項式記述子と比較する。

実験結果

リサーチクエスチョン

• RQ1 Lennard-JonesやEAMといった標準的な原子間ポテンシャルが、KANフレームワークの特殊ケースとして表現可能かどうか。
• RQ2 KANベースの記述子は、従来の多項式記述子（例：チェビシェフ）と比較して、精度と計算コストの両面でどのように異なるか。
• RQ3 KAN関数の線形または3次スプライン補間によって、ANNポテンシャルの評価コストをどの程度低減できるか、精度の劣化はどの程度か。
• RQ4 補間の解像度（点の数）を増加させた場合、補間されたKAN記述子の精度にどのような影響があるか。
• RQ5 KAN記述子は、ænetのような既存のMLMDソフトウェアスタックに効率的に統合可能か。

主な発見

• Lennard-Jones、埋め込み原子モデル（EAM）、および人工ニューラルネットワーク（ANN）ポテンシャルは、いずれもKANフレームワークの特定の例として解釈可能である。
• チェビシェフ多項式に基づくKAN記述子は、多項式記述子を用いた基準となるANNポテンシャルと同等の平均絶対誤差（MAE）2.05796 meV/atomを達成した。
• 線形または3次スプライン補間を用いることで、KAN記述子の評価コストがO(log 𝑁)にまで低減され、顕著な精度損なわず高速な推論が可能になった。
• 元のKAN記述子と補間されたKAN記述子の間の残差ノルムは、補間点数が増加するにつれて減少し、近似の質が向上していることが示された。
• チェビシェフ多項式と補間されたKAN記述子の間のMAE差は、補間点数が増加するにつれて減少し、高解像度ではゼロに近づく。
• 記述子評価における定性的な高速化が実現され、経過時間はO(log 𝑁)にほぼ比例して増加し、理論的予測と整合的であった。