Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators (2402.13984v2)
Abstract: Machine learning force fields (MLFFs) are an attractive alternative to ab-initio methods for molecular dynamics (MD) simulations. However, they can produce unstable simulations, limiting their ability to model phenomena occurring over longer timescales and compromising the quality of estimated observables. To address these challenges, we present Stability-Aware Boltzmann Estimator (StABlE) Training, a multi-modal training procedure which leverages joint supervision from reference quantum-mechanical calculations and system observables. StABlE Training iteratively runs many MD simulations in parallel to seek out unstable regions, and corrects the instabilities via supervision with a reference observable. We achieve efficient end-to-end automatic differentiation through MD simulations using our Boltzmann Estimator, a generalization of implicit differentiation techniques to a broader class of stochastic algorithms. Unlike existing techniques based on active learning, our approach requires no additional ab-initio energy and forces calculations to correct instabilities. We demonstrate our methodology across organic molecules, tetrapeptides, and condensed phase systems, using three modern MLFF architectures. StABlE-trained models achieve significant improvements in simulation stability, data efficiency, and agreement with reference observables. By incorporating observables into the training process alongside first-principles calculations, StABlE Training can be viewed as a general semi-empirical framework applicable across MLFF architectures and systems. This makes it a powerful tool for training stable and accurate MLFFs, particularly in the absence of large reference datasets.
- “OptNet: Differentiable optimization as a layer in neural networks” In Proceedings of the International Conference on Machine Learning (ICML), 2017, pp. 136–145
- “Gaussian approximation potentials: The accuracy of quantum mechanics, without the electrons” In Phys. Rev. Lett. 104, 2010, pp. 136403
- “A foundation model for atomistic materials chemistry”, 2023 arXiv:2401.00096 [physics.chem-ph]
- “MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields” In Neural Information Processing Systems (NeurIPS), 2022
- “E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials” In Nature Communications 13.1 Springer ScienceBusiness Media LLC, 2022 DOI: 10.1038/s41467-022-29939-5
- J. Behler “Neural network potential-energy surfaces in chemistry: A tool for large-scale simulations” In Phys. Chem. Chem. Phys. 13, 2011, pp. 17930–17955
- “Generalized Neural-Network Representation of High-Dimensional Potential-Energy Surfaces” In Phys. Rev. Lett. 98 American Physical Society, 2007, pp. 146401 DOI: 10.1103/PhysRevLett.98.146401
- “EGraFFBench: Evaluation of Equivariant Graph Neural Network Force Fields for Atomistic Simulations”, 2023 arXiv:2310.02428 [cs.LG]
- “Efficient and Modular Implicit Differentiation” In Neural Information Processing Systems (NeurIPS), 2022 Google Research
- Peter G. Bolhuis, Z.Faidon Brotzakis and Bettina G. Keller “Optimizing molecular potential models by imposing kinetic constraints with path reweighting” In J. Chem. Phys. 159, 2023, pp. 074102 DOI: 10.1063/5.0151166
- “On the Quantum Theory of Molecules” In German In Annalen der Physik 389.20, 1927, pp. 457–484 DOI: 10.1002/andp.19273892002
- “Unified Approach for Molecular Dynamics and Density-Functional Theory” Published 25 November 1985 In Phys. Rev. Lett. 55 American Physical Society, 1985, pp. 2471 DOI: 10.1103/PhysRevLett.55.2471
- Scott P. Carmichael and M.Scott Shell “A New Multiscale Algorithm and Its Application to Coarse-Grained Peptide Models for Self-Assembly” In J. Phys. Chem. B 116.29 American Chemical Society, 2012, pp. 8383–8393 DOI: 10.1021/jp2114994
- “Navigating protein landscapes with a machine-learned transferable coarse-grained model”, 2023 arXiv:2310.18278 [q-bio.BM]
- “A universal graph deep learning interatomic potential for the periodic table” In Nat Comput Sci 2, 2022, pp. 718–728 DOI: 10.1038/s43588-022-00349-3
- “Ab initio thermodynamics of liquid and solid water” Edited by Pablo G. Debenedetti, Princeton University, Princeton, NJ, approved December 3, 2018 (received for review September 4, 2018) In Proc. Natl. Acad. Sci. U.S.A. 116.4, 2019, pp. 1110–1115 DOI: 10.1073/pnas.1815117116
- “Accurate global machine learning force fields for molecules with hundreds of atoms”, 2022 arXiv:2209.14865 [physics.chem-ph]
- “Machine learning of accurate energy-conserving molecular force fields” In Science Advances 3.5, 2017 DOI: 10.1126/sciadv.1603015
- “CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling” In Nature Machine Intelligence 5.9 Nature Publishing Group UK London, 2023, pp. 1031–1041
- “Prediction Errors of Molecular Machine Learning Models Lower than Hybrid DFT Error” Epub 2017 Oct 10 In J Chem Theory Comput 13.11, 2017, pp. 5255–5264 DOI: 10.1021/acs.jctc.7b00577
- “Understanding Molecular Simulation: From Algorithms to Applications” Academic Press, 2001
- Richard A. Friesner “Ab initio quantum chemistry: Methodology and applications” In Proceedings of the National Academy of Sciences 102.19, 2005, pp. 6648–6653 DOI: 10.1073/pnas.0408036102
- “Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations”, 2022 arXiv:2210.07237 [physics.comp-ph]
- Johannes Gasteiger, Florian Becker and Stephan Günnemann In Conference on Neural Information Processing Systems (NeurIPS), 2021
- Johannes Gasteiger, Janek Groß and Stephan Günnemann “Directional message passing for molecular graphs” In 8th International Conference on Learning Representations, ICLR, 2020
- “GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets”, 2022 arXiv:2204.02782 [cs.LG]
- “e3nn: Euclidean Neural Networks”, 2022 arXiv:2207.09453 [cs.LG]
- “On differentiating parameterized argmin and argmax problems with application to bi-level optimization” In arXiv preprint arXiv:1607.05447, 2016
- “ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations”, 2021 arXiv:2103.01436 [cs.LG]
- “Learning Protein Structure with a Differentiable Simulator” In International Conference on Learning Representations, ICLR, 2019
- R.O. Jones “Density functional theory: Its origins, rise to prominence, and future” In Rev. Mod. Phys. 87 American Physical Society, 2015, pp. 897–923 DOI: 10.1103/RevModPhys.87.897
- “Efficient interatomic descriptors for accurate machine learning force fields of extended molecules” Erratum in: Nat Commun. 2023 Jul 11;14(1):4116. PMID: 37322039; PMCID: PMC10272221. In Nat Commun 14.1, 2023, pp. 3562 DOI: 10.1038/s41467-023-39214-w
- Thomas N. Kipf and Max Welling “Semi-Supervised Classification with Graph Convolutional Networks”, 2017 arXiv:1609.02907 [cs.LG]
- “Rescuing off-equilibrium simulation data through dynamic experimental data with dynAMMo” © 2023 The Author(s) In Mach. Learn.: Sci. Technol. 4.4 IOP Publishing Ltd, 2023, pp. 045050 DOI: 10.1088/2632-2153/ac7b4f
- “Linear Atomic Cluster Expansion Force Fields for Organic Molecules: Beyond RMSE” PMID: 34735161 In Journal of Chemical Theory and Computation 17.12, 2021, pp. 7696–7711 DOI: 10.1021/acs.jctc.1c00647
- “MACE-OFF23: Transferable Machine Learning Force Fields for Organic Molecules”, 2023 arXiv:2312.15211 [physics.chem-ph]
- “Scaling the Leading Accuracy of Deep Equivariant Models to Biomolecular Simulations of Realistic Size” In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’23 Denver, CO, USA: Association for Computing Machinery, 2023 DOI: 10.1145/3581784.3627041
- “Learning continuous models for continuous physics” In Communications Physics 6.1 Nature Publishing Group UK London, 2023, pp. 319
- H.G. Kümmel “A biography of the coupled cluster method” In Recent Progress in Many-Body Theories Singapore: World Scientific Publishing, 2002, pp. 334–348
- “The atomic simulation environment — a python library for working with atoms” In Journal of Physics: Condensed Matter, 2017 DOI: 10.1088/1361-648X/aa680e
- “Searching Configurations in Uncertainty Space: Active Learning of High-Dimensional Neural Network Reactive Potentials” In J. Chem. Theory Comput. 17.5, 2021, pp. 2691–2701 DOI: 10.1021/acs.jctc.0c01193
- “Spherical Message Passing for 3D Molecular Graphs” In International Conference on Learning Representations (ICLR), 2022
- M. Majewski, A. Pérez and P. Thölke “Machine learning coarse-grained potentials of protein thermodynamics” In Nat Commun 14 Nature Publishing Group, 2023, pp. 5739 DOI: 10.1038/s41467-023-41343-1
- “Nuclear quantum effects enter the mainstream” In Nat Rev Chem 2, 2018, pp. 0109 DOI: 10.1038/s41570-017-0109
- A. Merchant, S. Batzner and S.S. Schoenholz “Scaling deep learning for materials discovery” In Nature 624 Nature Publishing Group, 2023, pp. 80–85 DOI: 10.1038/s41586-023-06735-9
- Joe D. Morrow, John L.A. Gardner and Volker L. Deringer “How to validate machine-learned interatomic potentials” In J. Chem. Phys. 158, 2023, pp. 121501
- “Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics”, 2022 arXiv:2204.05249 [physics.comp-ph]
- Carles Navarro, Maciej Majewski and Gianni De Fabritiis “Top-Down Machine Learning of Coarse-Grained Protein Force Fields” In J. Chem. Theory Comput., 2023
- Geoffrey Négiar, Michael W. Mahoney and Aditi Krishnapriyan “Learning differentiable solvers for systems with hard constraints” In The Eleventh International Conference on Learning Representations, 2023
- W.G. Noid “Perspective: coarse-grained models for biomolecular systems” In J. Chem. Phys. 139 American Institute of Physics, 2013, pp. 090901
- John P. Perdew, Kieron Burke and Matthias Ernzerhof “Generalized Gradient Approximation Made Simple” In Phys. Rev. Lett. 77 American Physical Society, 1996, pp. 3865–3868 DOI: 10.1103/PhysRevLett.77.3865
- Z Raza “Proton ordering in cubic ice and hexagonal ice; a potential new ice phase–XIc” In Phys. Chem. Chem. Phys. 13, 2011, pp. 19788–19795 DOI: 10.1039/c1cp22167a
- “TorchOpt: An Efficient Library for Differentiable Optimization”, 2022 arXiv:2211.06934 [cs.MS]
- “Accurate machine learning force fields via experimental and simulation data fusion”, 2023 arXiv:2308.09142 [physics.chem-ph]
- “Fast and accurate modeling of molecular atomization energies with machine learning” In Phys. Rev. Lett. 108, 2012, pp. 058301
- “The Graph Neural Network Model” In IEEE Transactions on Neural Networks 20.1, 2009, pp. 61–80 DOI: 10.1109/TNN.2008.2005605
- “Learned Force Fields Are Ready For Ground State Catalyst Discovery”, 2022 arXiv:2209.12466 [cond-mat.mtrl-sci]
- “SchNet – A deep learning architecture for molecules and materials” In J. Chem. Phys. 148.24, 2018, pp. 241722 DOI: 10.1063/1.5019779
- Kristof T. Schütt, Oliver T. Unke and Michael Gastegger “Equivariant message passing for the prediction of tensorial properties and molecular spectra” In International Conference on Machine Learning (ICML), 2021
- “From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction”, 2023 arXiv:2310.16802 [cs.LG]
- “Differentiable Simulations for Enhanced Sampling of Rare Events”, 2023 arXiv:2301.03480 [physics.chem-ph]
- “Less is more: Sampling chemical space with active learning” In J. Chem. Phys. 148, 2018, pp. 241733 DOI: 10.1063/1.5023802
- “How robust are modern graph neural network potentials in long and hot molecular dynamics simulations?” In Machine Learning: Science and Technology 3.4 IOP Publishing, 2022, pp. 045010 DOI: 10.1088/2632-2153/ac9955
- S. Takamoto, C. Shinagawa and D. Motoki “Towards universal neural network potential for material discovery applicable to arbitrary combination of 45 elements” In Nat Commun 13, 2022, pp. 2991 DOI: 10.1038/s41467-022-30687-9
- “Learning neural network potentials from experimental data via Differentiable Trajectory Reweighting” In Nature Communications 12, 2021, pp. 6884 DOI: 10.1038/s41467-021-27241-4
- “Accurate Molecular Van Der Waals Interactions from Ground-State Electron Density and Free-Atom Reference Data” In Phys. Rev. Lett. 102 American Physical Society, 2009, pp. 073005 DOI: 10.1103/PhysRevLett.102.073005
- “Accurate and Efficient Method for Many-Body van der Waals Interactions” In Phys. Rev. Lett. 108 American Physical Society, 2012, pp. 236402 DOI: 10.1103/PhysRevLett.108.236402
- J. Vandermause, S.B. Torrisi and S. Batzner “On-the-fly active learning of interpretable Bayesian force fields for atomistic rare events” In npj Comput Mater 6, 2020, pp. 20 DOI: 10.1038/s41524-020-0283-z
- Joshua A. Vita and Daniel Schwalbe-Koda “Data efficiency and extrapolation trends in neural network interatomic potentials” In Mach. Learn.: Sci. Technol. 4 IOP Publishing, 2023, pp. 035031
- Wujie Wang, Simon Axelrod and Rafael Gómez-Bombarelli “Differentiable Molecular Simulations for Control and Learning”, 2020 arXiv:2003.00868 [physics.comp-ph]
- “Learning pair potentials using differentiable simulations” In The Journal of Chemical Physics 158.4 AIP Publishing, 2023, pp. 044113 DOI: 10.1063/5.0126475
- “Improving machine learning force fields for molecular dynamics simulations with fine-grained force metrics” In J. Chem. Phys. 159, 2023, pp. 035101 DOI: 10.1063/5.0147023
- Yujie Wu, Harald L. Tepper and Gregory A. Voth “Flexible simple point-charge water model with improved liquid-state properties” In The Journal of Chemical Physics 124.2, 2006, pp. 024503 DOI: 10.1063/1.2136877
- Robert W. Zwanzig “High-Temperature Equation of State by a Perturbation Method. I. Nonpolar Gases” In J. Chem. Phys. 22, 1954, pp. 1420–1426