FeNNol: an Efficient and Flexible Library for Building Force-field-enhanced Neural Network Potentials

TL;DR

FeNNol is a versatile library that simplifies the creation and deployment of hybrid neural network potentials enhanced with force-field interactions, achieving high computational efficiency for molecular simulations.

Contribution

It introduces a modular, user-friendly library that combines ML and physical models using Jax, enabling fast evaluation and easy hybrid model construction without programming.

Findings

FeNNol enables ANI-2x to reach speeds comparable to AMOEBA on GPUs.

The library simplifies hybrid model development for molecular simulations.

FeNNol demonstrates significant speed improvements over traditional NNP evaluation.

Abstract

Neural network interatomic potentials (NNPs) have recently proven to be powerful tools to accurately model complex molecular systems while bypassing the high numerical cost of ab-initio molecular dynamics simulations. In recent years, numerous advances in model architectures as well as the development of hybrid models combining machine-learning (ML) with more traditional, physically-motivated, force-field interactions have considerably increased the design space of ML potentials. In this paper, we present FeNNol, a new library for building, training and running force-field-enhanced neural network potentials. It provides a flexible and modular system for building hybrid models, allowing to easily combine state-of-the-art embeddings with ML-parameterized physical interaction terms without the need for explicit programming. Furthermore, FeNNol leverages the automatic differentiation and just-in-time compilation features of the Jax Python library to enable fast evaluation of NNPs, shrinking the performance gap between ML potentials and standard force-fields. This is demonstrated with the popular ANI-2x model reaching simulation speeds nearly on par with the AMOEBA polarizable force-field on commodity GPUs (GPU=Graphics processing unit). We hope that FeNNol will facilitate the development and application of new hybrid NNP architectures for a wide range of molecular simulation problems.