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.

• The paper introduces StABlE Training, a novel dual-modal approach that integrates quantum data with system observables to improve NNIP stability.
• The paper demonstrates that iterative simulation and learning phases effectively optimize NNIPs, achieving high accuracy with smaller reference datasets.
• The paper validates the method on diverse systems, showing significant improvements in capturing long-term molecular dynamics and structural observables.

Stability-Aware Training Enhances the Robustness of Neural Network Interatomic Potentials

Introduction to StABlE Training

Neural network interatomic potentials (NNIPs) have emerged as powerful tools for molecular dynamics (MD) simulations, offering a cost-effective alternative to ab-initio methods. Despite their advantages, NNIPs often struggle with simulation stability, particularly over extended time periods, which restricts their utility in capturing long-term physical phenomena. Addressing this critical issue, the Stability-Aware Boltzmann Estimator (StABlE) Training presents a novel training approach that significantly enhances the stability and accuracy of NNIPs by incorporating system observables alongside conventional supervised learning from quantum mechanical data.

Methodological Advancements

StABlE Training introduces a dual-modal training paradigm that converges on stable and accurate NNIPs by iteratively optimizing both quantum-mechanical energy and forces, as well as macroscopic system observables. This approach leverages the Boltzmann Estimator for the gradient-based optimization of NNIPs towards given observables, enabling efficient computation of gradients while detecting global and localized instabilities.

The methodology comprises two main phases: simulation and learning. The simulation phase aims to explore the phase space to identify unstable regions. Upon encountering instabilities, the process transitions to the learning phase, where the NNIP is refined to mitigate the instabilities by aligning with reference system observables. This cycle repeats, driving the NNIP towards greater stability while maintaining fidelity to quantum-mechanical reference data.

Empirical Validation

StABlE Training was validated on diverse systems including organic molecules, tetrapeptides, and condensed-phase systems, employing three modern NNIP architectures. Across these systems, StABlE-trained models demonstrated significant improvements in simulation stability, accurately capturing structural and dynamic observables. Notably, in some instances, StABlE-trained models surpassed the performance of models trained on datasets up to 50 times larger. This underscores the efficiency and effectiveness of the StABlE Training approach in enhancing NNIP stability and accuracy without the need for extensive reference datasets.

Implications and Future Directions

The development of StABlE Training marks a significant advancement in the training of NNIPs, offering a robust framework for producing stable and accurate MD simulations. This method holds considerable promise for extending the applicability of NNIPs to modeling complex phenomena over long timescales, which has been a longstanding challenge in the field.

Looking ahead, there are several exciting avenues for further research. One potential direction involves integrating StABlE Training with experimental observables, offering a pathway to even more generalizable and robust potentials. Additionally, exploring the incorporation of dynamical observables into the training process could further constrain the learning problem, potentially leading to even greater improvements in NNIP performance.

Conclusion

StABlE Training represents a significant step forward in the quest for stable and accurate NNIPs. By adeptly integrating system observables into the training process, this approach opens new horizons for the application of NNIPs across a wider array of complex simulation tasks. As the field of generative AI and LLMs continues to evolve, innovations like StABlE Training will undoubtedly play a pivotal role in unlocking new scientific insights and advancing our understanding of the molecular world.

The remarkable improvements in NNIP stability and accuracy brought forth by StABlE Training not only underscore the potential of integrating multimodal learning approaches but also set a new standard for future research in the development of robust and reliable NNIPs.