Uncertainty-aware machine learning surrogates of molecular dynamics simulations

We introduce an approach to use the statistical uncertainties associated with outputs of molecular dynamics simulations of soft materials to improve the training of deep neural networks when creating machine learning surrogates for predicting the relationship between input parameters and simulation outputs. The approach is illustrated by designing a surrogate model for molecular dynamics simulations of confined electrolytes, which aims to predict the intricate relationship between the features of the electrolyte solution (e.g., ion size, electrolyte concentration) and the resulting ionic structure. We propose a surrogate model to consider the direct regression and replace the point estimation with a probabilistic distribution in the output space. By incorporating probabilistic embeddings with Kullback–Leibler divergence in the loss function, we show that the model can significantly reduce prediction errors for samples in the unseen test dataset as well as yield higher generalizability and robustness among different datasets.