[Paper Review] A critical examination of robustness and generalizability of machine learning prediction of materials properties

Kangming Li, Brian DeCost+3 more|arXiv (Cornell University)|Oct 24, 2022

Machine Learning in Materials Science8 citations

TL;DR

This paper demonstrates that machine learning models pre-trained on the Materials Project 2018 (MP18) dataset suffer severe performance degradation when predicting formation energies of new materials in the MP21 dataset due to distribution shift. By using UMAP visualization and query-by-committee active learning, the authors identify out-of-distribution samples and improve prediction accuracy by adding just 1% of the test data, offering a practical framework to enhance robustness and generalizability in materials ML.

ABSTRACT

Recent advances in machine learning (ML) methods have led to substantial improvement in materials property prediction against community benchmarks, but an excellent benchmark score may not imply good generalization of performance. Here we show that ML models trained on the Materials Project 2018 (MP18) dataset can have severely degraded prediction performance on new compounds in the Materials Project 2021 (MP21) dataset. We document performance degradation in graph neural networks and traditional descriptor-based ML models for both quantitative and qualitative predictions. We find the source of the predictive degradation is due to the distribution shift between the MP18 and MP21 versions. This is revealed by the uniform manifold approximation and projection (UMAP) of the feature space. We then show that the performance degradation issue can be foreseen using a few simple tools. Firstly, the UMAP can be used to investigate the connectivity and relative proximity of the training and test data within feature space. Secondly, the disagreement between multiple ML models on the test data can illuminate out-of-distribution samples. We demonstrate that the simple yet efficient UMAP-guided and query-by-committee acquisition strategies can greatly improve prediction accuracy through adding only 1~\% of the test data. We believe this work provides valuable insights for building materials databases and ML models that enable better prediction robustness and generalizability.

Motivation & Objective

• To investigate the generalization performance of ML models trained on historical materials databases when applied to newer, expanded datasets.
• To identify the root cause of performance degradation in ML predictions for new materials, particularly in the context of evolving materials databases.
• To develop practical, data-efficient strategies for diagnosing and mitigating out-of-distribution prediction failures in materials science ML.
• To improve the robustness and reliability of ML models for materials property prediction by leveraging feature space analysis and ensemble disagreement.

Proposed method

• Employed UMAP (Uniform Manifold Approximation and Projection) to reduce high-dimensional Matminer-extracted features (90 after correlation filtering) into a lower-dimensional space for visualizing data distribution shifts between MP18 and MP21.
• Trained a state-of-the-art graph neural network (ALIGNN) and three traditional ML models (XGBoost, Random Forest, Linear Forest) on MP18 data for formation energy prediction.
• Used query-by-committee disagreement across multiple models to detect out-of-distribution test samples that exhibit high prediction uncertainty.
• Applied UMAP-guided active learning by selecting the most outlying samples (based on UMAP clustering) for retraining, improving model accuracy with minimal additional data.
• Evaluated model performance using standard regression metrics (e.g., R², MAE) on MP21 test sets, comparing models trained on full MP18 vs. MP18 plus selected new data.
• Proposed a hybrid prediction strategy incorporating linear models to improve robustness for extrapolative predictions when data scarcity or distribution shift occurs.

Experimental results

Research questions

• RQ1 To what extent do ML models trained on MP18 degrade in performance when predicting formation energies of new materials in the MP21 dataset?
• RQ2 What is the underlying cause of performance degradation in ML models when generalizing to new materials in evolving databases?
• RQ3 Can UMAP-based visualization of feature space reveal distribution shifts that correlate with prediction failure?
• RQ4 Can disagreement among multiple ML models help identify out-of-distribution samples before deployment?
• RQ5 Can active learning guided by UMAP and model disagreement significantly improve prediction accuracy with minimal additional data?

Key findings

• The ALIGNN model, pre-trained on MP18, exhibited strong underestimation of formation energies for a subset of new materials in MP21, with performance degradation exceeding 50% in MAE compared to MP18 test set.
• Traditional models like XGBoost and Linear Forest showed more robust phase stability estimation on MP21 data than the state-of-the-art ALIGNN model.
• UMAP visualization revealed that poorly predicted MP21 materials were located far from the MP18 training data distribution in the feature space, indicating distribution shift.
• Model disagreement (query-by-committee) effectively highlighted out-of-distribution samples, with high disagreement correlating strongly with high prediction error.
• Adding only 1% of the MP21 test data—selected via UMAP clustering and model disagreement—improved prediction accuracy significantly, reducing MAE by up to 30% on the worst-performing subset.
• The study demonstrates that UMAP-guided active learning is a highly efficient strategy for improving model generalizability with minimal data acquisition cost.