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

arXiv: https://arxiv.org/abs/2408.15852 · HTML: https://mlsanity.com/paper/2408.15852/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 112 references that are in the map.
- First posted on arXiv 28 Aug 2024.

## Builds on (most cited first, 22 of 22)

- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems](https://mlsanity.com/paper/1603.04467/) (2016)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules](https://mlsanity.com/paper/2011.14115/) (2020)
- [Gradients are Not All You Need](https://mlsanity.com/paper/2111.05803/) (2021)
- [DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics](https://mlsanity.com/paper/1712.03641/) (2017)
- [Machine Learning of coarse-grained Molecular Dynamics Force Fields](https://mlsanity.com/paper/1812.01736/) (2018)
- [TorchMD: A deep learning framework for molecular simulations](https://mlsanity.com/paper/2012.12106/) (2020)
- [Flow-matching -- efficient coarse-graining of molecular dynamics without forces](https://mlsanity.com/paper/2203.11167/) (2022)
- [Coarse-Graining Auto-Encoders for Molecular Dynamics](https://mlsanity.com/paper/1812.02706/) (2018)
- [SchNetPack 2.0: A neural network toolbox for atomistic machine learning](https://mlsanity.com/paper/2212.05517/) (2022)
- [Learning neural network potentials from experimental data via Differentiable Trajectory Reweighting](https://mlsanity.com/paper/2106.01138/) (2021)
- [Machine Learning Implicit Solvation for Molecular Dynamics](https://mlsanity.com/paper/2106.07492/) (2021)
- [Accurate machine learning force fields via experimental and simulation data fusion](https://mlsanity.com/paper/2308.09142/) (2023)
- [Scalable Bayesian Uncertainty Quantification for Neural Network Potentials: Promise and Pitfalls](https://mlsanity.com/paper/2212.07959/) (2022)
- [Predicting solvation free energies with an implicit solvent machine learning potential](https://mlsanity.com/paper/2406.00183/) (2024)
- [FeNNol: an Efficient and Flexible Library for Building Force-field-enhanced Neural Network Potentials](https://mlsanity.com/paper/2405.01491/) (2024)
- [Top-down machine learning of coarse-grained protein force-fields](https://mlsanity.com/paper/2306.11375/) (2023)
- [Uncertainty estimation for molecular dynamics and sampling](https://mlsanity.com/paper/2011.08828/) (2020)

## Built on by (most cited first, 12 of 12)

- [chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations](https://mlsanity.com/paper/2506.04055/) (2025)
- [Learning Non-Local Molecular Interactions via Equivariant Local Representations and Charge Equilibration](https://mlsanity.com/paper/2501.19179/) (2025)
- [Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields](https://mlsanity.com/paper/2503.08674/) (2025)
- [Enhanced Sampling for Efficient Learning of Coarse-Grained Machine Learning Potentials](https://mlsanity.com/paper/2510.11148/) (2025)
- [Enhancing Machine Learning Potentials through Transfer Learning across Chemical Elements](https://mlsanity.com/paper/2502.13522/) (2025)
- [Energy-Based Coarse-Graining in Molecular Dynamics: A Flow-Based Framework without Data](https://mlsanity.com/paper/2504.20940/) (2025)
- [Generalization of Long-Range Machine Learning Potentials in Complex Chemical Spaces](https://mlsanity.com/paper/2512.10989/) (2025)
- [Bridging the Simulation-to-Experiment Gap with Generative Models using Adversarial Distribution Alignment](https://mlsanity.com/paper/2604.01169/) (2026)
- [Refining Coarse-Grained Molecular Topologies: A Bayesian Optimization Approach](https://mlsanity.com/paper/2501.02707/) (2025)
- [Refining Machine Learning Potentials through Thermodynamic Theory of Phase Transitions](https://mlsanity.com/paper/2512.03974/) (2025)
- [ConSolv: Solvent-Conditional Machine Learning Implicit Solvent Potential](https://mlsanity.com/paper/2606.24983/) (2026)
- [Transferable Implicit Solvent Machine Learning Potential for Drugs and Proteins Approaching Ab Initio Accuracy](https://mlsanity.com/paper/2607.10887/) (2026)

## Abstract (arXiv)

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.