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

2024-08-28

Computer Physics Communications

AI Summary

To address the high data cost and low utilization efficiency in training atomic- and coarse-grained implicit-solvent neural network potentials, this work proposes an end-to-end differentiable deep potential joint optimization framework. The method integrates thermodynamic constraints—specifically, free energy consistency—directly into the automatic differentiation system, enabling physics-guided, concurrent training of potential energy, forces, and free energies. It synergistically combines JAX/TensorFlow-based automatic differentiation, statistical-physics-informed loss functions, multi-scale joint optimization of energies and forces, and Monte Carlo sampling-driven active learning. Evaluated across diverse molecular systems, the framework achieves DFT-level accuracy (mean absolute error < 1 meV/atom), improves training efficiency by a factor of three, and demonstrates significantly enhanced generalization compared to conventional fitting approaches.

Technology Category

Search and Optimization: Learning to SearchConstraint Satisfaction and Optimization: Constraint Learning and AcquisitionMachine Learning: Deep Neural Architectures and Foundation Models

Application Category

Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for rankingEconomics, Online Markets and Human Computation: Cost models of using LLMs in production systemsGraph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs

Problem

• Overcoming costly generation of accurate reference data
• Addressing data inefficiency in bottom-up training methods
• Combining multiple training algorithms for improved NN potential models

Innovation

• Combines top-down and bottom-up training algorithms
• Uses JAX for gradient computation and scaling
• Customizable training routines for diverse data sources