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

TL;DR

This paper introduces a multi-head committee approach for atomistic models that enables efficient and reliable uncertainty prediction, improving active learning and reducing training data needs without sacrificing accuracy.

Contribution

It presents a novel multi-head committee mechanism for message-passing neural networks, allowing uncertainty estimation and active learning for foundation models with minimal retraining.

Findings

Uncertainty correlates well with true errors across datasets.

Active learning with the method reduces training set size to 5%.

Foundation models maintain accuracy with reliable uncertainty estimates.

Abstract

Machine learning potentials have become a standard tool for atomistic materials modelling. While models continue to become more generalisable, an open challenge relates to efficient uncertainty predictions for active learning and robust error analysis. In this work, we utilise 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, specifically 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.