Multi-head committees enable direct uncertainty prediction for atomistic foundation models

Abstract

Machine learning potentials have become a standard tool for atomistic materials modeling. While models continue to become more generalizable, an open challenge relates to efficient uncertainty predictions for active learning and robust error analysis. In this work, we utilize MACE and its multi-head mechanism to implement a committee neural network potential for message-passing architectures, where the committee comprises multiple output modules attached to the same atomic environment descriptors. As with traditional committees of independent networks, the standard deviation of the predictions functions as an estimate of the model’s uncertainty. We show for a range of datasets in custom-build models that the uncertainty of the force predictions correlates well with the true errors. We subsequently apply this concept to foundation models, in particular MACE-MP-0, where we train only the newly attached output heads while keeping the remaining part of the model fixed. We use this approach in an active learning workflow to condense the training set of the foundation model to just 5% of its original size. The foundation model multi-head committee trained on the condensed training set enables reliable uncertainty estimation without any substantial decrease in prediction accuracy.