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

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.

11 pages, 7 figures in main article + supporting information

References in corpus (14)

• Gaussian Approximation Potentials: the accuracy of quantum mechanics, without the electrons
• E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials
• A Universal Graph Deep Learning Interatomic Potential for the Periodic Table
• Machine Learning Molecular Dynamics for the Simulation of Infrared Spectra
• Committee neural network potentials control generalization errors and enable active learning
• Fast and Accurate Uncertainty Estimation in Chemical Machine Learning
• TeaNet: universal neural network interatomic potential inspired by iterative electronic relaxations
• Performance Assessment of Universal Machine Learning Interatomic Potentials: Challenges and Directions for Materials' Surfaces
• Uncertainty estimation for molecular dynamics and sampling
• Fast Uncertainty Estimates in Deep Learning Interatomic Potentials
• Use the force! Reduced variance estimators for densities, radial distribution functions and local mobilities in molecular simulations
• Deep Ensembles vs. Committees for Uncertainty Estimation in Neural-Network Force Fields: Comparison and Application to Active Learning
• Model-free quantification of completeness, uncertainties, and outliers in atomistic machine learning using information theory
• Machine-learning interatomic potentials from a users perspective: A comparison of accuracy, speed and data efficiency