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

Paul Fuchs,Stephan Thaler,Sebastien Röcken,Julija Zavadlav
2024-08-28
Abstract:Neural Networks (NNs) are promising models for refining the accuracy of molecular dynamics, potentially 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.
Chemical Physics,Machine Learning,Computational Physics
What problem does this paper attempt to address?
### Problems the Paper Aims to Solve This paper aims to address several key issues in the training process of neural network (NN) potential models in molecular dynamics (MD): 1. **Data Efficiency**: Traditional bottom-up training methods require a large amount of high-precision reference data, which becomes impractical when dealing with large-scale systems. Therefore, the paper proposes a framework to efficiently combine data from multiple sources, improving data utilization. 2. **Multi-Scale Consistency**: By integrating top-down macroscopic observations with bottom-up microscopic details, the paper ensures consistency across models at different scales. This helps improve the accuracy of the model in macroscopic predictions. 3. **Algorithm Flexibility**: A highly customizable training framework `chemtrain` has been developed, supporting a combination of various training algorithms such as Force Matching, Relative Entropy Minimization, and Differentiable Trajectory Reweighting. This flexibility allows researchers to design more complex training strategies based on specific needs. 4. **Computational Efficiency**: The use of JAX framework's automatic differentiation capabilities simplifies algorithm implementation and supports efficient utilization of computational resources, thereby enhancing overall computational efficiency. Through these methods, the paper aims to provide a flexible, efficient, and easily extensible neural network potential training solution for molecular dynamics simulations.