[Paper Review] Kolmogorov--Arnold networks in molecular dynamics

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

History and advancements in chemistry4 citations

TL;DR

This paper proposes integrating Kolmogorov–Arnold Networks (KANs) into molecular dynamics to enhance interatomic potential efficiency. By redefining artificial neural network (ANN) descriptors using KANs' non-linear functions and applying linear or cubic spline interpolation, the method reduces computational cost for potential and force evaluations while preserving accuracy, as demonstrated on Lennard-Jones and embedded atom model potentials.

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.

Motivation & Objective

• To improve the efficiency of machine learning molecular dynamics (MLMD) simulations by reinterpreting common interatomic potentials within the KAN framework.
• To reduce the computational cost of evaluating artificial neural network (ANN) potentials and their derivatives in molecular dynamics.
• To explore the use of KANs as an alternative to polynomial-based descriptors in ANN potentials, leveraging their non-linear activation functions.
• To demonstrate that linear or cubic spline interpolation of KAN functions enables fast, accurate evaluation of descriptors without sacrificing model performance.
• To lay the groundwork for integrating KAN-based descriptors into open-source MLMD packages like ænet for broader adoption.

Proposed method

• Reformulates widely used interatomic potentials—Lennard-Jones, embedded atom model (EAM), and ANN potentials—within the Kolmogorov–Arnold representation theorem framework.
• Represents the potential energy surface as a superposition of univariate non-linear functions (KAN layers), where the activation functions are optimized during training.
• Replaces traditional polynomial-based descriptors (e.g., Chebyshev polynomials) in ANN potentials with KAN-derived descriptors using radial and angular non-linear functions.
• Employs linear or cubic spline interpolation to approximate the KAN activation functions, reducing evaluation cost from O(𝑁) to O(log 𝑁) per evaluation.
• Uses interpolation libraries (e.g., Interpolation.jl) to compute interpolated descriptors efficiently, with error measured via residual norm and mean absolute error (MAE).
• Validates performance using test datasets, comparing MAE of KAN descriptors with and without interpolation against reference Chebyshev polynomial descriptors.

Experimental results

Research questions

• RQ1 Can standard interatomic potentials like Lennard-Jones and EAM be expressed as special cases of the KAN framework?
• RQ2 How do KAN-based descriptors compare to traditional polynomial descriptors (e.g., Chebyshev) in terms of accuracy and computational cost?
• RQ3 To what extent can linear or cubic spline interpolation of KAN functions reduce the evaluation cost of ANN potentials without degrading accuracy?
• RQ4 What is the trade-off between interpolation resolution (number of points) and the accuracy of the interpolated KAN descriptor?
• RQ5 Can the KAN descriptor be efficiently integrated into existing MLMD software stacks such as ænet?

Key findings

• The Lennard-Jones, embedded atom model (EAM), and artificial neural network (ANN) potentials can all be interpreted as specific instances of the KAN framework.
• KAN descriptors based on Chebyshev polynomials achieve a mean absolute error (MAE) of 2.05796 meV/atom, matching the reference ANN potential with polynomial descriptors.
• Using linear or cubic spline interpolation, the computational cost of evaluating KAN descriptors is reduced to O(log 𝑁), enabling faster inference without significant accuracy loss.
• The residual norm between original and interpolated KAN descriptors decreases with increasing interpolation points, indicating improved approximation quality.
• The mean absolute error (MAE) difference between Chebyshev and interpolated KAN descriptors diminishes as the number of interpolation points increases, approaching zero at high resolution.
• The method enables a qualitative speedup in descriptor evaluation, with elapsed time scaling approximately as O(log 𝑁), consistent with theoretical expectations.