Graph Contrastive Learning for Materials
Abstract
Recent work has shown the potential of graph neural networks to efﬁciently predict material properties, enabling high-throughput screening of materials.
Training these models, however, often requires large quantities of labelled data, obtained via costly methods such as ab initio calculations or experimental evaluation.
By leveraging a series of material-speciﬁc transformations, we introduce CrystalCLR, a framework for constrastive learning of representations with crystal graph neural networks.
With the addition of a novel loss function, our framework is able to learn representations competitive with engineered ﬁngerprinting methods.
We also demonstrate that via model ﬁnetuning, contrastive pretraining can improve the performance of graph neural networks for prediction of material properties and signiﬁcantly outperform traditional ML models that use engineered ﬁngerprints.
Lastly, we observe that CrystalCLR produces material representations that form clusters by compound class.
Introduction
The discovery of novel materials is a problem of considerable interest.
Machine learning (ML) has demonstrated potential in accelerating material discovery with accurate prediction performance and lower computational cost [1].
Traditional ML models used in material discovery rely on expert knowledge to construct material ﬁngerprints as inputs to ML models [2].
The development of ﬁngerprint-based representations require manual design which is often limited to speciﬁc material families and properties and can be time-consuming to compute.
Recent advances in deep learning provide an alternative way of learning feature representations automatically from data.
In particular, graph neural networks (GNNs) have been developed for modeling crystalline materials to predict material properties [3–8].
In addition, Magar et al.
[9] have introduced several augmentations speciﬁc to crystalline structures, including perturb-structure and supercell.
They ﬁnd that training GNNs on augmented crystal data results in improved performance in prediction of several material properties.
While GNNs are becoming increasingly popular, like in other deep learning domains, adoption is limited by the availability of high quality labelled data [10].
Oftentimes, we only have access to very limited labelled data, such as materials’ physical properties (e.g., melting point, thermal conductivity) that require laborious and time-consuming experimentation or computationally expensive ab initio calculations.
One way to tackle this problem is to transfer the learned model trained on a related base task with large labelled data to the target task [11].
However, this approach is limited by the availability of large labelled data for the base task.
In the absence of labelled data, self-supervised methods provide a way to learn representations from samples alone.
Contrastive learning is one such method of self-supervised learning, which seeks to learn representations such that similar pairs remain close in embedding space, while others remain distant.
In recent years, this has become a powerful method of both learning of representations, as well as self-supervised pre-training for transfer learning, primarily in the image domain [12].
You et al. [13] introduced a framework for applying contrastive learning to graphs, using GNNs and broadly
Fe
Fe Fe Fe
P P
P
P (a) Original
Fe
Fe Fe Fe
P P
P
P (b) Strain
Fe
Fe Fe Fe
P P
P
P (c) Perturb
Fe
Fe Fe Fe
P
N
P
P (d) Column replace (e) Supercell applicable transformations to biochemical molecules, social network, and image super-pixel graphs.
Similar work has shown that contrastive pretraining can improve supervised performance in predicting the chemical properties of molecules [14].
Lastly, Khosla et al.
[15] showed how supervised labels can be incorporated into the contrastive learning framework, allowing these labels to be incorporated into the representations.
Despite the success of contrastive learning in a broad range of domains, very little work has been done to apply contrastive learning to crystalline materials.
In this work, we introduce CrystalCLR, a framework for the constrastive representation learning of crystalline materials.
We highlight the importance of augmentation selection, and show that contrastive learning is capable of learning representations competitive with ﬁngerprinting methods.
We also incorporate a novel loss function into the contrastive learning objective, which further improve representation quality.
Furthermore, we show transfer learning with the pretrained contrastive model can outperform supervised learning for the prediction of material properties.
Method
Following [12, 13], we construct our constrastive learning framework with four major components: augmentation , GNN encoder , projection head and contrastive loss .
Let V and E denote the atomic and bond attributes of a material’s crystal structure, respectively.
Each crystal graph G =
( V, E ) is ﬁrst augmented into a similar pair ˆ G i and ˆ G j where ˆ G i
=
( ˆ V i , ˆ E i ) and the augmentations used to transform G into ˆ G are speciﬁc to the domain of crystalline materials (Sec. 2.1).
A GNN encoder , f ( · ) , maps crystal graphs ˆ G i , ˆ G j into representations
h i , h j .
For our work we use the CGCNN architecture as the encoder, following the same graph representation of crystal structure as inputs [3].
The projection head , a two layer MLP g ( · ) , projects representations h
i , h j into 128-dimensional space to create projections z i , z j .
The addition of a non-linear projection prior to the loss function has shown to improve representation quality [12].
The contrastive loss function seeks to maximize the agreement between representations z i , z j augmented from the same crystalline material while minimizing the agreement among the rest of the pairs augmented from different crystalline materials.
We use the NT-Xent loss [16], written for pair z i , z j in a batch of N pairs as:
[Equation]
where sim ( u , v ) denotes cosine similarity ( u (cid:62) v /
[Equation]
u (cid:107)(cid:107) v (cid:107) ).
The ﬁnal loss is the sum across all positive pairs in the batch.
2.1 Crystal Augmentations
We investigate four transformations speciﬁc to crystalline materials (Fig. 1).
Each transformation is isoelectronic, preserving the valence electron conﬁguration of the material: (1) perturb structure : each site is perturbed by a distance uniformly sampled between 0 and 50% of the minimum pairwise distance within the structure; (2) strain : a random, anisotropic tensile strain deformation is applied to the structure.
Each lattice vector is increased by a factor uniformly sampled from [0 , 0 . 05] ; (3) column replace : with probability 1 , a single, randomly selected site within the structure is replaced with an element from the same column in the periodic table as the original element; and (4) supercell : a supercell of the original crystal is created, scaling each lattice vector by a factor of 3 . The perturb structure strength range is selected based on values used in prior work [9], and the strain strength range is selected based on physically reasonable values.
2.2 Optimizing for composition similarity
Under the notion that materials with similar composition may tend to have similar properties, we hypothesize that explicitly guiding representations of similarly composed materials together may improve the representation quality.
We propose an additional loss term, composition similarity (CS), to promote representation similarity among materials that have one or more common elements.
Each crystal is assigned a binary vector a ∈ { 0 , 1 } 100 , indicating which of the possible 100 atoms (that exist in the dataset used) are contained in the crystal.
The set P ( i ) = { p ∈
I : a i · a p > 0 , i (cid:54) = p } is the set of indices that share one or more atom with the i th crystal.
The loss can then be deﬁned as:
[Equation]
Where | P ( i ) | is the cardinality of the set.
While prior work in [15] has shown that placement of the summation over positives outside the log is more theoretically optimal, we found the above formulation to produce better representations empirically.
In the ﬁnal model, (2) is used in combination with the original NT-Xent loss (1) and weighted equally.
2.3 Learning Task Description
Our self-supervised training dataset consists of 90,160 crystal structures downloaded from Materials
Project in April, 2021 [17].
For evaluation, we use additional melting temperature, thermal conductivity, and bulk moduli (K_VRH) datasets of sizes 3,014, 5,540, and 13,121 respectively.
The melting temperature data were experimental data scraped from MatWeb [18].
Thermal conductivity data were downloaded from AFLOW [19]; all values are at 300K. Bulk moduli data were downloaded from
Materials Project.
Each dataset is split 80/10/10 into train, test, and validation sets.
Experimental Setup and Results 3.1 Augmentation Study
P e r t u r b
S t r u c t u r e
S t r a i n
C o l u m n R e p l a c e
S u p e
r c e ll 1st Transformation
Perturb Structure
Strain
ColumnReplace
Supercell 2 n d T r a n s f o r m a t
i o n
[Figure]
To evaluate the effects of individual and combinations of augmentations, we train models under all of the pairwise combinations of augmentations.
For this study, we apply augmentations to only one sample in each positive pair.
Each model is evaluated by performing a linear regression on the learned representations for each of the studied material properties.
In our work, we use no augmentations for evaluation.
Similar to [12, 13], we observe that the composition of transformations is crucial for representation quality.
As shown in Figure 2 for melting point prediction, the best representations are obtained using the perturb-structure and column-replace transformations, and the addition of supercell often deteriorates performance.
We observe a similar trend across all of the studied material properties.
3.2 Model Evaluation
We selected perturb-structure, strain, and column-replace as the ﬁnal set of data augmentations; each augmentation has a 50% chance of being applied to each training sample.
We trained two self-supervised models; each model was trained for 5000 epochs, with a batch size of 512 .
The ﬁrst model (CrystalCLR) uses L CLR alone, while the second model (CrystalCLR+CS) uses a ﬁnal loss of
L CLR + L CS .
We use the Adam optimizer [20] with a learning rate of 1e-4.
Melting temp
[Table]
Random Forest evaluation:
Site Stats Fingerprint
[Table]
Sine Coulomb Fingerprint
[Table]
Element Properties
[Table]
CrystalCLR (ours)
[Table]
CrystalCLR + CS (ours)
[Table]
CrystalCLR
[Table]
ours
[Table]
+
[Table]
ours) [Table] To quantitatively evaluate the learned representations, we trained random-forest regression models [21] to predict material properties. We compare the predictive performance of the learned embeddings to several other material ﬁngerprinting methods, all obtainable without supervision: (1) Site stats ﬁngerprint that aggregates statistics of the CrystalNNFingerprint [22] using structural order parameters, (2) Sine Coulomb ﬁngerprint which is a variant of the Coloumb matrix for periodic crystals [23], and (3) Element property which is a weighted mean, standard deviation, minimum and maximum of element properties obtained from Magpie [24].
See Table 2 for the list of properties used.
We use the default hyperparameters for the random forest from scikit-learn [25].
Table 1 shows the representations learned from contrastive training perform at or above traditional ﬁngerprinting methods, with the exception of element properties.
Element properties contain additional information, including melting point, which explain its high performance.
Furthermore, we observe that the addition of the composition loss yields an improvement in performance across all properties.
[Figure]
We also evaluate the use of contrastive learning as a pretraining task for material property prediction.
For each of aforementioned properties, we train CGCNN models; one randomly initialized, and one initialized with the CGCNN encoder weights from each of the two contrastive training methods.
For each model, we use the Adam optimizer [20] to ﬁne-tune the models, using the learning rate with the best validation performance.
Table 1 shows by transferring weights from the contrastive-pretrained encoder, we are able to obtain more accurate predictions of material properties than a randomly initialized model for melting temperature and bulk moduli.
The beneﬁt of the composition loss is more evident for melting temperature prediction than the thermal conductivity and bulk moduli properties.
Lastly, in Figure 3 we visualize produced crystal embeddings using UMAP [26].
CrystalCLR learns representations that form clusters by compound class, despite training in a self-supervised manor.
Conclusion
In this work, we introduce a method for the contrastive training of graph neural networks for representation learning of materials.
First, we establish several crystalline material speciﬁc transformations, and study the effects of the composition of transformations on representation quality.
In addition, we introduce a novel loss function to explicitly maximize the similarity among materials that share common elements.
We show that our framework is capable of producing representations competitive with engineered material ﬁngerprinting techniques.
Finally we show that constrastive pretraining improves performance over random initialization for downstream tasks.
Future work will investigate additional contrastive learning methods, the ability to use CrystalCLR representations for material retrieval, and the applicability of equivariant graph neural networks. Acknowledgments and Disclosure of Funding DISTRIBUTION STATEMENT A. Approved for public release. Distribution is unlimited.
This material is based upon work supported by the Under Secretary of Defense for Research and Engineering under Air Force Contract No. FA8702-15-D-0001.
Any opinions, ﬁndings, conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reﬂect the views of the Under Secretary of Defense for Research and Engineering.
{+c (cid:13) 
+}
Massachusetts Institute of
Technology.
Delivered to the U.S. Government with Unlimited Rights, as deﬁned in DFARS Part 252.227-7013 or 7014 (Feb 2014).
Notwithstanding any copyright notice, U.S. Government rights in this work are deﬁned by DFARS 252.227-7013 or DFARS 252.227-7014 as detailed above.
Use of this work other than as speciﬁcally authorized by the U.S. Government may violate any copyrights that exist in this work.
Appendix 5.1 Element Properties List
Table 2 lists the element properties used for the Element Property classiﬁer.
[Table]
(position on the periodic table, counting columnwise from H)
[Table]