Out-of-distribution materials property prediction using adversarial learning based fine-tuning

TL;DR

The paper tackles the persistent challenge of predicting material properties under out-of-distribution (OOD) conditions. It introduces Crystal Adversarial Learning (CAL), a SAL-based adversarial training pipeline augmented with partial sampling and a targeting fine-tuning mechanism, implemented on a compact IR-net backbone. CAL demonstrates improved stability and OOD performance across covariate, prior, and relation shifts, with pronounced benefits in low-data settings such as piezoelectric materials. The work provides a practical framework for robust materials property prediction under distribution shifts and data scarcity, with implications for more reliable materials discovery.

Abstract

The accurate prediction of material properties is crucial in a wide range of scientific and engineering disciplines. Machine learning (ML) has advanced the state of the art in this field, enabling scientists to discover novel materials and design materials with specific desired properties. However, one major challenge that persists in material property prediction is the generalization of models to out-of-distribution (OOD) samples,i.e., samples that differ significantly from those encountered during training. In this paper, we explore the application of advancements in OOD learning approaches to enhance the robustness and reliability of material property prediction models. We propose and apply the Crystal Adversarial Learning (CAL) algorithm for OOD materials property prediction,which generates synthetic data during training to bias the training towards those samples with high prediction uncertainty. We further propose an adversarial learning based targeting finetuning approach to make the model adapted to a particular OOD dataset, as an alternative to traditional fine-tuning. Our experiments demonstrate the success of our CAL algorithm with its high effectiveness in ML with limited samples which commonly occurs in materials science. Our work represents a promising direction toward better OOD learning and materials property prediction.

Figures (6)

• Figure 1: Dataset selections for covariate and prior shifts. The covariate shift samples are selected from the periphery of the low-dimensional surface formed by the data. The prior shift dataset consists of outliers in formation energy.
• Figure 2: Supervised UMAP embedding of the data, illustrating variance in both $X$ and $y$ (formation energy). The covariate shift data is located on the periphery, far from the majority of the data. The prior shift data forms its own distinct cluster, while the relation shift data is related but does not clearly fall into either of the previous categories.
• Figure 3: The CAL pipeline for out-of-distribution (OOD) material property prediction. It is based on the Stable Adversarial Learning (SAL) pipeline and integrates individual residual networks for material property prediction. At each step, only one partial sample set and the corresponding target set are used for adversarial sample generation. The training schedule begins with five epochs of training on the entire dataset to initialize the network. Subsequently, the training samples are ranked based on their losses, and the top 20% of high-loss samples are selected to continue training the network. In the base CAL model, these selected samples are used to generate adversarial attack samples. In our targeted training scheme, however, samples from the test domain are used to generate the adversarial attack samples.
• Figure 4: Neural network architecture of the IR-net block. The primary difference from a standard multi-layer perceptron is the addition of the outputs to the inputs. If the input and output dimensions differ, a linear projection of the inputs is used instead.
• Figure 5: Visualization of OOD test set samples: (a) Two-dimensional compositional feature space representation, with colored markers indicating individual samples' relationships to local clusters and global structure. (b) Distribution of target values for the OOD test sets, showing the frequency of the highest optical phonon mode peak. The y-axis is on a log scale.