Uncertainty-Informed Training Set Construction for Robust Extrapolation in Materials Property Prediction

Abstract

Machine learning models such as Crystal graph convolutional neural networks (CGCNN) have been widely adopted for the rapid and accurate prediction of crystalline material properties. However, these models suffer from a critical limitation: prediction errors increase sharply during extrapolation to out-of-distribution structures. In this study, we integrate Bayesian neural network (BNN) techniques into the CGCNN framework to quantify predictive uncertainty in formation energy estimation. Using a dataset of materials from the Materials Project, we assess the model’s extrapolation performance across varying lattice strains and diverse crystal systems. Our results reveal that CGCNN predictions are governed primarily by local atomic environments rather than by macroscopic crystal symmetry. Consequently, while the model extrapolates effectively to systems with similar local environments, it can exhibit overconfidence in low-symmetry structures where high errors occur despite low uncertainty. Building on these findings, we introduce a dataset construction strategy that utilizes BNN-derived uncertainty metrics while accounting for these architectural limitations in recognizing global structural changes. By strategically prioritizing the inclusion of structurally vulnerable regions identified through both uncertainty and structural analysis, the proposed methodology enables the development of models with improved extrapolation robustness and reliability in complex chemical spaces.