Machine Learning Force Fields with Data Cost Aware Training
Alexander Bukharin Tianyi Liu Shengjie Wang Simiao Zuo Weihao Gao Wen Yan Tuo Zhao
Machine learning force fields (MLFF) have been proposed to accelerate
molecular dynamics (MD) simulation, which finds widespread applications in
chemistry and biomedical research. Even for the most data-efficient MLFFs,
reaching chemical accuracy can require hundreds of frames of force and energy
labels generated by expensive quantum mechanical algorithms, which may scale as
$O(n^3)$ to $O(n^7)$, with $n$ proportional to the number of basis functions.
To address this issue, we propose a multi-stage computational framework --
ASTEROID, which lowers the data cost of MLFFs by leveraging a combination of
cheap inaccurate data and expensive accurate data. The motivation behind
ASTEROID is that inaccurate data, though incurring large bias, can help capture
the sophisticated structures of the underlying force field. Therefore, we first
train a MLFF model on a large amount of inaccurate training data, employing a
bias-aware loss function to prevent the model from overfitting tahe potential
bias of this data. We then fine-tune the obtained model using a small amount of
accurate training data, which preserves the knowledge learned from the
inaccurate training data while significantly improving the model's accuracy.
Moreover, we propose a variant of ASTEROID based on score matching for the
setting where the inaccurate training data are unlabeled. Extensive experiments
on MD datasets and downstream tasks validate the efficacy of ASTEROID. Our code
and data are available at https://github.com/abukharin3/asteroid.