On-the-Fly Fine-Tuning of Foundational Neural Network Potentials: A Bayesian Neural Network Approach

TL;DR

This paper introduces a Bayesian neural network-based method for on-the-fly fine-tuning of neural network potentials, enabling automatic model updates, uncertainty quantification, and enhanced detection of rare events during molecular simulations.

Contribution

It presents a novel Bayesian approach for real-time fine-tuning of foundational neural network potentials, addressing uncertainty estimation and rare event sampling.

Findings

Effective on-the-fly model fine-tuning during simulations.

Automatic detection and sampling of rare events.

Maintains desired accuracy throughout the process.

Abstract

Due to the computational complexity of evaluating interatomic forces from first principles, the creation of interatomic machine learning force fields has become a highly active field of research. However, the generation of training datasets of sufficient size and sample diversity itself comes with a computational burden that can make this approach impractical for modeling rare events or systems with a large configuration space. Fine-tuning foundation models that have been pre-trained on large-scale material or molecular databases offers a promising opportunity to reduce the amount of training data necessary to reach a desired level of accuracy. However, even if this approach requires less training data overall, creating a suitable training dataset can still be a very challenging problem, especially for systems with rare events and for end-users who don't have an extensive background in machine learning. In on-the-fly learning, the creation of a training dataset can be largely automated by using model uncertainty during the simulation to decide if the model is accurate enough or if a structure should be recalculated with classical methods and used to update the model. A key challenge for applying this form of active learning to the fine-tuning of foundation models is how to assess the uncertainty of those models during the fine-tuning process, even though most foundation models lack any form of uncertainty quantification. In this paper, we overcome this challenge by introducing a fine-tuning approach based on Bayesian neural network methods and a subsequent on-the-fly workflow that automatically fine-tunes the model while maintaining a pre-specified accuracy and can detect rare events such as transition states and sample them at an increased rate relative to their occurrence.