Optimal pre-train/fine-tune strategies for accurate material property predictions

Abstract

Overcoming the challenge of limited data availability within materials science is crucial for the broad-based applicability of machine learning within materials science. One pathway to overcome this limited data availability is to use the framework of transfer learning (TL), where a pre-trained (PT) machine learning model (on a larger dataset) can be fine-tuned (FT) on a target (typically smaller) dataset. Our study systematically explores the effectiveness of various PT/FT strategies to learn and predict material properties with limited data. Specifically, we leverage graph neural networks (GNNs) to PT/FT on seven diverse curated materials datasets, encompassing sizes ranging from 941 to 132,752 datapoints. We consider datasets that cover a spectrum of material properties, ranging from band gaps (electronic) to formation energies (thermodynamic) and shear moduli (mechanical). We study the influence of PT and FT dataset sizes, strategies that can be employed for FT, and other hyperparameters on pair-wise TL among the datasets considered. We find our pair-wise PT-FT models to consistently outperform models trained from scratch on the target datasets. Importantly, we develop a GNN framework that is simultaneously PT on multiple properties (MPT), enabling the construction of generalized GNN models. Our MPT models outperform pair-wise PT-FT models on several datasets considered, and more significantly, on a 2D material band gap dataset that is completely out-of-distribution from the PT datasets. Finally, we expect our PT/FT and MPT frameworks to be generalizable to other GNNs and materials properties, which can accelerate materials design and discovery for various applications.