[论文解读] Kolmogorov--Arnold networks in molecular dynamics

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

History and advancements in chemistry被引用 4

一句话总结

该论文提出将柯尔莫哥洛夫-阿诺德网络（Kolmogorov–Arnold Networks, KANs）集成到分子动力学中，以提升原子间势能的计算效率。通过使用KAN的非线性函数重新定义人工神经网络（ANN）描述符，并应用线性或三次样条插值，该方法在保持精度的同时降低了势能和受力计算的计算成本，已在Lennard-Jones势和嵌入原子模型势上得到验证。

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.

研究动机与目标

• 通过在KAN框架下重新诠释常见原子间势能，提升机器学习分子动力学（MLMD）模拟的效率。
• 降低在分子动力学中评估人工神经网络（ANN）势能及其导数的计算成本。
• 探索使用KAN作为多项式基描述符的替代方案，利用其非线性激活函数。
• 证明对KAN函数进行线性或三次样条插值，可在不损失模型性能的前提下实现快速、准确的描述符评估。
• 为将基于KAN的描述符集成到开源MLMD软件包（如ænet）中奠定基础，以促进更广泛的应用。

提出的方法

• 在柯尔莫哥洛夫-阿诺德表示定理框架下，重新表述广泛使用的原子间势能——Lennard-Jones势、嵌入原子模型（EAM）和ANN势能。
• 将势能面表示为单变量非线性函数（KAN层）的叠加，其中激活函数在训练过程中进行优化。
• 在ANN势能中，用基于径向和角向非线性函数的KAN衍生描述符，替代传统的多项式基描述符（如切比雪夫多项式）。
• 采用线性或三次样条插值来近似KAN激活函数，将每次评估的计算成本从O(𝑁)降低至O(log 𝑁)。
• 使用插值库（如Interpolation.jl）高效计算插值后的描述符，通过残差范数和平均绝对误差（MAE）衡量误差。
• 使用测试数据集验证性能，比较插值前后KAN描述符与参考切比雪夫多项式描述符的MAE。

实验结果

研究问题

• RQ1 Lennard-Jones势和EAM等标准原子间势能是否可作为KAN框架的特例表达？
• RQ2 基于KAN的描述符与传统多项式描述符（如切比雪夫多项式）在精度和计算成本方面相比如何？
• RQ3 对KAN函数进行线性或三次样条插值，能在多大程度上降低ANN势能的评估成本，同时不降低精度？
• RQ4 插值分辨率（插值点数量）与插值后KAN描述符精度之间的权衡如何？
• RQ5 KAN描述符能否高效集成到现有MLMD软件栈（如ænet）中？

主要发现

• Lennard-Jones势、嵌入原子模型（EAM）和人工神经网络（ANN）势能均可被解释为KAN框架的特定实例。
• 基于切比雪夫多项式的KAN描述符实现2.05796 meV/atom的平均绝对误差（MAE），与使用多项式描述符的参考ANN势能精度相当。
• 通过线性或三次样条插值，KAN描述符的评估计算成本降低至O(log 𝑁)，实现更快的推理，且精度损失可忽略。
• 原始KAN描述符与插值后描述符之间的残差范数随插值点数增加而减小，表明近似质量提高。
• 随着插值点数的增加，切比雪夫多项式与插值后KAN描述符之间的MAE差异逐渐减小，高分辨率下趋近于零。
• 该方法实现了描述符评估的定性加速，运行时间近似符合O(log 𝑁)的理论预期。