Source-linked AI summary
TorchMD: A deep learning framework for molecular simulations
Stefan Doerr, Maciej Majewsk, Adrià Pérez, Andreas Krämer, Cecilia Clementi, Frank Noe, Toni Giorgino, Gianni De Fabritiis
TL;DR
Molecular dynamics needs accurate force fields but remains limited by force-field construction and computational cost. TorchMD addresses this gap with a PyTorch-based engine that combines classical and machine-learned potentials, supports differentiable simulations, and is validated across several molecular modeling tasks. The reported experiments include good energy conservation, QM9 training performance, and recovery of optimized charges after one training epoch.
Problem
Classical molecular dynamics requires extensive force-field tables and remains computationally intensive for reaching several important physiological timescales.
Method
TorchMD is a PyTorch-based molecular dynamics engine that combines classical force terms with differentiable neural-network potentials and supports end-to-end training.
Results
TorchMD demonstrated Amber all-atom simulations, QM9 neural-network training, end-to-end charge optimization, and coarse-grained neural-network modeling for protein folding.
Takeaways & Limitations
TorchMD provides a practical tool-set for molecular simulations involving machine-learning potentials, including training, simulation, and differentiable parameter optimization.
Takeaways & Limitations
Coarse-grained protein-folding simulations require prior bond and repulsion potentials to constrain dynamics to regions represented in the training data.
Abstract
from arXiv · showhide
Molecular dynamics simulations provide a mechanistic description of molecules by relying on empirical potentials. The quality and transferability of such potentials can be improved leveraging data-driven models derived with machine learning approaches. Here, we present TorchMD, a framework for molecular simulations with mixed classical and machine learning potentials. All of force computations including bond, angle, dihedral, Lennard-Jones and Coulomb interactions are expressed as PyTorch arrays and operations. Moreover, TorchMD enables learning and simulating neural network potentials. We validate it using standard Amber all-atom simulations, learning an ab-initio potential, performing an end-to-end training and finally learning and simulating a coarse-grained model for protein folding. We believe that TorchMD provides a useful tool-set to support molecular simulations of machine learning potentials. Code and data are freely available at \url{github.com/torchmd}.
1 Introduction
TorchMD addresses the computational and force-field limitations of classical molecular dynamics by integrating machine-learned potentials into a PyTorch-based simulation framework. It supports rapid prototyping, differentiable simulations, and open-source workflows for data-driven potentials.
- 1 Introduction: Classical all-atom MD enables quantitative molecular studies but requires extensive force-field tables and remains too compute-intensive for several physiological timescales.These limitations motivate data-driven potentials and more flexible simulation tooling.
- 1 Introduction: Deep neural networks can learn complex, many-body potential energy functions and their derivatives from large databases generated by expensive approaches.SchNet is cited as an example that predicts molecular energies and forces using graph-based continuous-filter convolutions.
- 1 Introduction: TorchMD is a molecular dynamics code built from scratch around PyTorch primitives.Its design targets integration of machine-learning models with molecular simulation.
- 1 Introduction: TorchMD extends standard bonded and non-bonded force terms with DNN-based potentials of arbitrary complexity and supports end-to-end differentiation of simulation parameters.The framework can integrate ab-initio neural network potentials and machine-learning coarse-grained potentials.
- 1 Introduction: The paper describes TorchMD capabilities, supported functional forms, and a fitting strategy for data-driven DNN potentials.The code, tutorial, and corresponding training data are openly available.
2 Methods
TorchMD combines conventional molecular mechanics terms with PyTorch-based external potentials and neural-network models. Its implementation supports standard simulation workflows, automatic differentiation, configurable force fields, and scalable machine-learning training.
- 2 Methods: The force-field implementation includes configurable units, YAML-based force-field files, and AMBER parameter reading through ParmEd.Figure 1 provides an example YAML force field for a water box.
- 2 Methods: TorchMD implements AMBER functional forms including harmonic bonds, angles, torsions, and non-bonded Van der Waals and electrostatic energies.The framework also supports NVT simulations, Langevin thermostats, velocity Verlet integration, periodic systems, and L-BFGS minimization.
- 2 Methods: TorchMD supports external PyTorch potentials that take atomic coordinates as input and return energies and forces.Automatic differentiation can obtain forces from the total energy, while analytical gradients are also implemented for AMBER terms for performance.
- 2 Methods: TorchMD-Net provides a PyTorch codebase for training neural-network potentials, currently using a SchNet-based model but allowing other force calculators.The network receives positions and box information at each timestep and returns external energies and forces.
- 2 Methods: Training supports multiple GPUs, batched concurrent simulations, energy-based QM9 learning, and force matching for coarse-grained models.PyTorch Lightning is used for multi-GPU training, while coarse-grained training targets thermodynamic behavior from atomistic simulations.
3 Results
TorchMD was evaluated across classical all-atom simulations, QM9 neural-network training, differentiable parameter inference, and coarse-grained chignolin folding. The results show accurate force-field energies and energy conservation, competitive QM9 training, rapid charge recovery, and partially successful coarse-grained folding with model-dependent limitations.
- All-atom simulations: 60 times slower was TorchMD’s performance relative to ACEMD3 across a water box, alanine dipeptide, and trypsin–benzamidine system.The discrepancy was attributed mainly to missing neighbor lists for non-bonded interactions.
- All-atom simulations: TorchMD’s potential energies differed from OpenMM by less than 10^-3 kcal/mol across 14 systems spanning ions, water boxes, small molecules, and proteins.The comparison tested all different force-field terms using identical parameters.
- All-atom simulations: 1.1 10^-5 K per degree of freedom was the mean normalized energy-conservation value for a 1 ns NVE water-box simulation.The simulation used a 1 fs timestep.
- QM9 training validation: 10 meV was the best QM9 performance reported for 100,000 training points, marginally better than SchNet’s reported best performance.The test performance followed a linear trend on a log-log scale across training-set sizes.
- End-to-end differentiable simulations: 3% accuracy was achieved in recovering the original partial atomic charges after one epoch, or 1000 iterations, from a short trajectory.The loss was the mean-squared distance between ground-truth and propagated coordinates, and gradients were obtained by backpropagation.
- Coarse-grained chignolin simulations: The CA coarse-grained chignolin model captured the folded state as the global energy minimum and represented unfolded and misfolded minima, but reproduced neither the connecting energy barriers nor all CACB behavior.The CACB model detected the global minimum but mislocated the misfolded minimum, failed to estimate the unfolded-region free energy, and was less stable.
4 Conclusion
TorchMD is presented as a PyTorch-based molecular dynamics engine supporting machine-learning applications from all-atom simulation through coarse-grained protein folding. The authors identify analytical supplements and model selection as important considerations for neural-network protein-folding potentials.
- TorchMD supports biomolecular simulations with machine-learning capabilities across Amber all-atom simulations, end-to-end parameter learning, and coarse-grained protein-folding potentials.
- Neural-network protein-folding potentials require asymptotic analytical bond and repulsion terms to avoid training-data-excluded conformations where predictions are unreliable.
- The authors coarse-grained proteins using either α-carbon atoms or α-carbon and β-carbon atoms, with the CA model currently appearing to work best.
- TorchMD’s end-to-end parameter differentiability may be useful to initiatives such as the Open Force Field Initiative, while further integrations are planned for OpenMM and ACEMD.