chemtrain: Learning Deep Potential Models via Automatic Differentiation and Statistical Physics

TL;DR

This work tackles data-efficient training of neural-network potentials for molecular dynamics by integrating top-down and bottom-up learning within a unified differentiable framework. The authors introduce chemtrain, a modular, JAX-based software that decomposes learning into reusable building blocks (FM, RM, DiffTRe, Delta-Learning) and a high-level API to compose training routines. They demonstrate the approach with two case studies: fusing experimental and simulation data to build a titanium atomistic model and combining FM with RM to train alanine dipeptide in implicit solvent, achieving improved accuracy and data efficiency. The framework emphasizes scalability, uncertainty handling, and practical deployment potential, with future directions including active learning, expanded observables, and tighter MD-code integration for large-scale simulations.

Abstract

Neural Networks (NNs) are effective models for refining the accuracy of molecular dynamics, opening up new fields of application. Typically trained bottom-up, atomistic NN potential models can reach first-principle accuracy, while coarse-grained implicit solvent NN potentials surpass classical continuum solvent models. However, overcoming the limitations of costly generation of accurate reference data and data inefficiency of common bottom-up training demands efficient incorporation of data from many sources. This paper introduces the framework chemtrain to learn sophisticated NN potential models through customizable training routines and advanced training algorithms. These routines can combine multiple top-down and bottom-up algorithms, e.g., to incorporate both experimental and simulation data or pre-train potentials with less costly algorithms. chemtrain provides an object-oriented high-level interface to simplify the creation of custom routines. On the lower level, chemtrain relies on JAX to compute gradients and scale the computations to use available resources. We demonstrate the simplicity and importance of combining multiple algorithms in the examples of parametrizing an all-atomistic model of titanium and a coarse-grained implicit solvent model of alanine dipeptide.

Figures (4)

• Figure 1: Structure of a Maximum Likelihood Estimate (MLE) learning procedure in chemtrain. The trainer executes algorithms to perform gradient-based optimization of a single or multiple objectives. The algorithms assemble from multi-functional building blocks and rely on a modular definition of the objectives.
• Figure 2: Hierarchy of the trainer classes in chemtrain. The figure displays only a subset of all methods and attributes to illustrate similarities between the trainers.
• Figure 3: Density histograms of the backbone dihedral angels of alanine dipeptide from simulations with Force Matching (FM), Relative Entropy Minimization (RM), and FM-then-RM trained models in comparison to an atomistic reference simulation.
• Figure 4: Predicted elastic constants. The black dashed line denotes the experimental reference, the green circles denote the prediction of the bottom-up trained model via Force Matching (FM), and the red stars denote the predictions after sequentially training bottom-up with FM and top-down with Differential Trajectory Reweighting (DiffTRe).