Coherent energy and force uncertainty in deep learning force fields

Peter Bjørn Jørgensen, Jonas Busk, Ole Winther, Mikkel N. Schmidt
2023-12-08
Abstract:In machine learning energy potentials for atomic systems, forces are commonly obtained as the negative derivative of the energy function with respect to atomic positions. To quantify aleatoric uncertainty in the predicted energies, a widely used modeling approach involves predicting both a mean and variance for each energy value. However, this model is not differentiable under the usual white noise assumption, so energy uncertainty does not naturally translate to force uncertainty. In this work we propose a machine learning potential energy model in which energy and force aleatoric uncertainty are linked through a spatially correlated noise process. We demonstrate our approach on an equivariant messages passing neural network potential trained on energies and forces on two out-of-equilibrium molecular datasets. Furthermore, we also show how to obtain epistemic uncertainties in this setting based on a Bayesian interpretation of deep ensemble models.
Machine Learning,Computational Physics
What problem does this paper attempt to address?
The problem that this paper attempts to solve is how to effectively estimate the uncertainties of energy and force in deep - learning force fields. Specifically, the paper focuses on how to correlate the aleatoric uncertainty of energy with the uncertainty of force when predicting the energy of atomic systems. Traditional methods usually assume that the noise is white noise (i.e., uncorrelated noise), which makes it infeasible to directly derive the uncertainty of force from the uncertainty of energy. Therefore, this paper proposes a new machine - learning potential energy model. By introducing a spatially - correlated noise process, the uncertainties of energy and force can be correlated with each other. This method not only solves the theoretical problem but also improves the model's ability to optimize data collection and the interpretability of prediction results. ### Main contributions: 1. **Proposed a new uncertainty model**: This model links the uncertainties of energy and force by introducing a spatially - correlated noise process, overcoming the non - differentiable problem under the traditional white - noise assumption. 2. **Provided closed - form expressions for the uncertainties of energy and force**: Through this new noise model, closed - form expressions for the mean and variance of force can be directly derived from the potential - energy model. 3. **Extended the applicable scope of the model**: This method is not only applicable to deep - learning force fields but can also be applied to other fields where energy and gradients need to be observed simultaneously. 4. **Verified the effectiveness of the model**: Through experiments on two molecular data sets, the effectiveness and superiority of this method in practical applications are proved. ### Formula summary: - **Energy observation model**: \[ E_{\text{obs}}(z, r)=E_{\theta}(z, r)+\rho_{\theta}(z, r)\eta \] where \(E_{\theta}(z, r)\) is the mean prediction of energy, \(\rho_{\theta}(z, r)\) is the variance prediction of energy, and \(\eta\) is a stochastic process with a differentiable autocorrelation function. - **Variance of force**: \[ \text{Var}\left(-\frac{\partial E_{\text{obs}}}{\partial r_{i, d}}(z, r)\right)=\hat{\gamma}\rho^{2}_{\theta}(z, r)+\left(\frac{\partial\rho_{\theta}(z, r)}{\partial r_{i, d}}\right)^{2} \] where \(\hat{\gamma}\) is the square of the inverse length scale and can be estimated from the training data. - **Uncertainty of model integration**: - Mean and variance of energy prediction: \[ E_{\theta}[E_{\text{obs}}(z, r)] = E_{\theta}[E_{\theta}(z, r)] \] \[ \text{Var}_{\theta}(E_{\text{obs}}(z, r))=E_{\theta}[\rho^{2}_{\theta}(z, r)]+\text{Var}_{\theta}(E_{\theta}(z, r)) \] - Variance of force prediction: \[ \text{Var}\left(-\frac{\partial E_{\text{obs}}}{\partial r_{i, d}}(z, r)\right)=E_{\theta}\left[\hat{\gamma}\rho^{2}_{\theta}(z, r)+\left(\frac{\partial\rho_{\theta}(z, r)}{\partial r_{i, d}}\right)^{2}\right]+\text{Var}_{\theta}\left(\frac{\partial E_{\theta}(z, r)}{\partial r_{i, d}}\right)