TorchMD-Net 2.0: Fast Neural Network Potentials for Molecular Simulations

TL;DR

TorchMD-Net 2.0 delivers a modular, fast framework for neural network potentials, introducing architectures like TensorNet and Equivariant Transformer and enabling integration with OpenMM through CUDA-graph-friendly pipelines. The work emphasizes speed and scalability via advanced neighbor search, static-shape CUDA graphs, and PyTorch compilation, while supporting physical priors and diverse datasets for broad applicability. Empirical results show state-of-the-art accuracy on QM9 with compact models, stable long-timescale MD simulations, and meaningful inference-speed gains, highlighting TorchMD-Net’s practicality for research and deployment. While memory and compute demands remain nontrivial, the paper outlines a viable path toward widespread adoption of NNPs in MD workflows, aided by ongoing hardware progress and programming-tooling improvements.

Abstract

Achieving a balance between computational speed, prediction accuracy, and universal applicability in molecular simulations has been a persistent challenge. This paper presents substantial advancements in the TorchMD-Net software, a pivotal step forward in the shift from conventional force fields to neural network-based potentials. The evolution of TorchMD-Net into a more comprehensive and versatile framework is highlighted, incorporating cutting-edge architectures such as TensorNet. This transformation is achieved through a modular design approach, encouraging customized applications within the scientific community. The most notable enhancement is a significant improvement in computational efficiency, achieving a very remarkable acceleration in the computation of energy and forces for TensorNet models, with performance gains ranging from 2-fold to 10-fold over previous iterations. Other enhancements include highly optimized neighbor search algorithms that support periodic boundary conditions and the smooth integration with existing molecular dynamics frameworks. Additionally, the updated version introduces the capability to integrate physical priors, further enriching its application spectrum and utility in research. The software is available at https://github.com/torchmd/torchmd-net.

Figures (6)

• Figure 1: The main module in TorchMD-Net is called TorchMD_Net from the torchmdnet.models.model module. This class combines a given representation model (such as the Equivariant Transformer), an output model (such as the scalar output module) and a prior model (such as the Atomref prior), producing a module that takes as input a series of atoms features and outputs a scalar value (i.e. energy per molecule) and when derivative = True, its negative gradient with respect to the input positions (i.e. atomic forces).
• Figure 2: Performance comparison of cell (solid line) and brute-force (dashed line) neighbor search strategies across different batch sizes for a random cloud of particles with $64$ neighbors per particle on average. Cell list performance tends to degrade with increasing batch size, while the opposite is true for brute force.
• Figure 3: Performance comparison of cell (solid line) and brute-force (dashed line) neighbor search strategies across different batch sizes for a random cloud of $32$k particles with $64$ neighbors per particle on average. The particles are split into a certain number of batches.
• Figure 4: Training and validation curves for two different training benchmark datasets. Up: TensorNet with 3 interaction layers on the QM9 $U_0$, parameters in table S3.Down: Equivariant Transformer on MD17, parameters in Table S2.
• Figure 5: (Left) RMSD analysis for the trajectories of 4 molecules outside of the training set. Simulations are carried out with TensorNet 2L, using the parameters in Table S4, with the exception of A-0L, in which a 0L TensorNet model is showcased. Presented data is plotted only every 4 ns for visualization clarity. (Right) Representation of the simulated molecules. Labels show the PubChem ID for each molecule.