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

TL;DR

FeNNol presents a modular, JAX/Flax-based library to design, train, and run force-field-enhanced neural-network potentials, addressing fragmentation in existing ML-FF implementations. It introduces CRATE, a configurable atomic embedding that fuses chemical, radial, angular, equivariant, and long-range resources across interaction layers to emulate leading architectures (ANI-like, AIMNet-like, MACE-like). The framework supports multi-stage training, diverse loss terms, and multiple MD pathways (ASE, Deep-HP, and native MD) and demonstrates near-AMOEBA-level GPU performance for ANI-2x on commodity GPUs and scalable performance for large systems. Overall, FeNNol lowers barriers to building, training, and deploying hybrid ML/MM potentials and enables rapid exploration of novel architectures for molecular simulations.

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.

Figures (4)

• Figure 1: Architecture and data flow of the FeNNol library. Black blocks represent functional modules, green arrows represent model operations and blue arrows represent data flows.
• Figure 2: Flow diagram of the execution of a ANI-like model built with FeNNol
• Figure 3: Total energy of a box of 216 water molecules as a function of time for different simulation precisions. All simulations use a timestep of 0.1fs and the ANI-2x model. The height of the green box represents 0.1% of the initial total energy.
• Figure 4: Performance (in ns simulated per day) for various system sizes. The black dashed line corresponds to the AMOEBA force field simulated using Tinker-HP. The solid lines correspond to ANI-2x with FeNNol's implementation and different MD engines. The dotted lines correspond to ANI-2x with TorchANI's implementation (min. img.=minimum image convention). All simulations use float32 precision, a 0.5fs timestep and a Langevin thermostat set at 300K. Simulations with TorchANI and FeNNol's native MD engine were performed on a RTX 3090 GPU. Simulations with Deep-HP were performed on a A100 GPU.