Self-Refining Training for Amortized Density Functional Theory

TL;DR

This paper introduces a self-refining training method for amortized density functional theory models, reducing reliance on large datasets by jointly training the model and sampling molecular conformations, improving efficiency and accuracy.

Contribution

The paper presents a novel self-refining training strategy that jointly optimizes a DFT predictor and samples training data, reducing dataset dependency and enhancing model performance.

Findings

Outperforms models trained solely on pre-collected datasets.

Efficient training with asynchronous sampling and optimization.

Open-source implementation available for reproducibility.

Abstract

Density Functional Theory (DFT) allows for predicting all the chemical and physical properties of molecular systems from first principles by finding an approximate solution to the many-body Schr\"odinger equation. However, the cost of these predictions becomes infeasible when increasing the scale of the energy evaluations, e.g., when calculating the ground-state energy for simulating molecular dynamics. Recent works have demonstrated that, for substantially large datasets of molecular conformations, Deep Learning-based models can predict the outputs of the classical DFT solvers by amortizing the corresponding optimization problems. In this paper, we propose a novel method that reduces the dependency of amortized DFT solvers on large pre-collected datasets by introducing a self-refining training strategy. Namely, we propose an efficient method that simultaneously trains a deep-learning model to predict the DFT outputs and samples molecular conformations that are used as training data for the model. We derive our method as a minimization of the variational upper bound on the KL-divergence measuring the discrepancy between the generated samples and the target Boltzmann distribution defined by the ground state energy. To demonstrate the utility of the proposed scheme, we perform an extensive empirical study comparing it with the models trained on the pre-collected datasets. Finally, we open-source our implementation of the proposed algorithm, optimized with asynchronous training and sampling stages, which enables simultaneous sampling and training. Code is available at https://github.com/majhas/self-refining-dft.