S2SNet: A Pretrained Neural Network for Superconductivity Discovery

1 Introduction

Superconductors are crucial for sustainable everyday applications and attractive for the United Nations' sustainable development goals (SDGs) [Miryala, 2021]. Superconductors have zero resistance, which allows perfect current transmission without losses. With superconducting cables, a continuous supply of power can be obtained for a timely availability of irrigation water in remote areas, which is a key to sustainability in performing agriculture. Then hunger can be effectively alleviated since it is partly caused by the lack of sustainability in performing agriculture. Affordable and clean energy can also be obtained because DC energy from places where power is plentiful, like solar plants in Africa (365 sunny days) can be delivered worldwide without losses. Zero resistance means "no heat-generation and no energyloss", which is the key contribution in the climate action. For good health and well-being, superconducting magnets, which can produce magnetic field of high-quality, high-intensity, are well used as high-field magnets in magnetic resonance imaging, nuclear magnetic resonance, water purification, magnetic drug delivery, and so on.

∗Equal contribution.

†Corresponding author: rux@zju.edu.cn

Superconductivity has been a hot topic along many fields for its great significance to human beings and SDGs. However, the theory of superconductivity is still a puzzle till now. Benefiting from massive publicly available datasets in the area of material science, many machine learning methods have been applied to predicting material properties, generating structures, and so on [Schmidt et al., 2019]. It also sheds a light to modeling superconductors with machine learning methods. Much of the recent interest in superconductivity modelling has focused on critical temperature regression using measured properties of the material [Stanev et al., 2018; Hamidieh, 2018], which is computationally expensive since feature engineering is required.

In AI for physical science, a lot of efforts have been made to model molecules with Graph Neural Networks (GNN) [Yang et al., 2021; Mercado et al., 2021] since the molecules are graphs naturally. However, quite few methods have been proposed for crystal modelling [Chen et al., 2019; Park and Wolverton, 2020; Xie and Grossman, 2018]. Due to the different annotations among different datasets, the current supervised learning models can only be trained on small datasets, while other unlabeled data are not fully utilized.

Recently, attention mechanism is widely used in Natural Language Processing (NLP) [Bahdanau et al., 2015] and Computer Vision (CV) [Dosovitskiy et al., 2020] to model the interaction of contexts. The pre-training task, MaskedLanguage Modeling (MLM) task, in both NLP and CV, has shown great potential to get model better and more robust with astounding experimental results [Devlin et al., 2019; Bahdanau et al., 2015]. Besides, with pre-training tasks, the unlabeled data can be utilized properly.

According to BCS theory, superconductivity can be microscopically explained as the interaction between electrons and lattice vibrations (phonons) [Bardeen et al., 1957]. Therefore, we propose a machine learning model for the way from Structure to Superconductivity (S2SNet), in which only crystal structures are required, to model the interaction in order to predict potential superconducting. Attention mechanism is applied to crystal structures. S2SNet is pre-trained on a large amount of unlabeled data to learn the interaction between different atoms in lattice.

Through S2SNet, we try to model the potential interaction between lattice dynamics and cooper pairs to predict the potential superconductor properties of material. It is experimentally verified that S2SNet outperforms other existing models. To the best of our knowledge, S2SNet is the first work showing that it is possible to predict the superconductivity of materials via their crystal structures.

The main contributions of our work can be summarized as follows:

• We propose a novel neural network to discover potential superconductivity of materials given only information of crystal structures for the first time.
• A new dataset S2S containing both crystal structure and superconductivity critical temperature is built for the first time. It provides a benchmark for future machine learning based superconductor discovery.
• To provide insight and guide the superconductivity research for physicists, we present a list of potential superconductors for future experimental validation.

2 RelatedWork

Superconductivity Prediction. Previous machine learning methods focus on predicting the superconductivity through the measured properties of materials, while these properties are obtained through complex and expensive feature engineering. The averaged valence-electron numbers, orbital radii differences, and metallic electronegativity differences were three critical features used in [Villars and Phillips, 1988]. Materials Agnostic Platform for Informatics and Exploration (Magpie) and AFLOW were used for feature engineering in [Stanev et al., 2018; Hamidieh, 2018].

Crystal Modeling. Quite few methods are proposed for crystal modeling. By connecting the atoms within a distance threshold, crystals were modeled as a graph with atoms as nodes in CGCNN [Park and Wolverton, 2020; Xie and Grossman, 2018]. External environmental conditions are taken into consideration by taking it as an additional nodes and connecting it with all the other nodes in MEGNet [Chen et al., 2019].

Unsupervised Learning. Various approaches have been proposed to make full use of the unlabeled data, most of which pre-train a model with the unlabeled dataset and finetune it on specific labeled dataset [Han et al., 2021]. Contrastive learning is commonly used for pre-training in CV [Le-Khac et al., 2020]. As pre-training tasks in NLP, diverse techniques such as mask language model (MLM), replaced token detection (RTD), and so on are often employed

(a) Schematic illustration of S2SNet

(b) Schematic illustration of attention mechanism

Figure 1: Illustration of S2SNet and attention mechanism. (a) Overall illustration of S2SNet. The Attention Module are performed for 10 times. FC1 and FC2 are two fully connected hidden layers. FC2 outputs the target ˆ yt. (b) Illustration of the Attention mechanism. The computation of representation of Atom a1 is presented. Best view in color.

[Hirschberg and Manning, 2015]. Following the pre-training model Bert, Zhang et. al. proposed "MG-Bert" for molecule modelling with unsupervised learning [Zhang et al., 2021].

3 Methodology

Attention mechanism is proposed in the filed of NLP and widely used in CV recently [Hirschberg and Manning, 2015; Dosovitskiy et al., 2020]. It is applied to the crystals in this study.

3.1 S2SNet

Fig. 1 shows the overall illustration of S2SNet and the attention mechanism. In this paper, it is defined that a crystal c ∈ C consists of atoms A = {a1, a2, · · · , an}, the location of atoms P = {p1, p2, · · · , pn}, and the lattice parameters U = {u1, u2, u3}, i.e. c = {A, P,U}. n, a, p and U denote the number of atoms in a unit cell, the atomic number, the actual 3-D location of an atom in the crystal, and the physical dimension of a unit cell respectively. The overall workflow of S2SNet can be described as followings:

1. Atoms A are embedded into 512-dim vectors through the Embedding layer. Expanded distance with Gaussian basis centered at 512 points linearly placed between 0 and 5.12 ˚A is applied to the location P of atoms in the crystal [Chen et al., 2019]. Specifically, the Euclidean distances P ′ between atoms and the center of crystal i.e. the mean coordinates of all atoms, are calculated. Then Gaussian basis are performed as Eq. 1, where P ′0 are 512 distances from 0 to 5.12 ˚A and σ = 2 [Chen et al., 2019].

1. The embedding of atoms A′ are atom-wise aggregated with the distance ˆ P as Eq. 2, where Wα is a matrix to

get the position ˆ P into the same dimension as A′ and model the interaction between different positions.

1. Attentionmechanism is performed on each atom respectively with its neighborhood atoms. Thus the embedding of atom contains all the information of itself and neighbors. neighbor(a′′ ) is defined as the atoms within the Euclidean distance, ρ, from a′′ . In this paper, ρ equals to 2 ˚A according to [Chen et al., 2019].

1. To aggregate the embedding of all atoms in crystal, the Set2Set [Vinyals et al., 2015] is performed on the set A′′′ .

1. Lattice parameters are aggregated to get the final representation of crystal as Eq. 5, where Wθ is a matrix to get the lattice parameters U into the same dimension as b.

1. Through the FC2 layer, target can be obtained and loss function is binary cross entropy as Eq. 6, where yt is the ground truth of whether it is a superconductor and ˆ yt is the prediction of S2SNet.

In Fig. 1(b), an example of the attention mechanism is illustrated. Suppose that the neighbors of atom a1 are a2 and a3. To get the representation of atom a1, a1 and its neighbors, a2 and a3, are firstly multiplied with WQ, WK and WV respectively. Three different groups of {q, k, v} are obtained. q1 from a1 is multiplied with k1, k2, k3 from all the atoms to get the score as Eq. 7, where dk is the dimension of k. Finally, the representation of atom a1 can be obtained as the weighted sum of V from all the atoms as Eq. 8, where N is the number of a1's neighbors plus one. Representations of other atoms ai ∈ A can also be obtained like this.

3.2 MLMPre-training Task

MLM pre-training task is a self-supervised learning method to model the interaction between atoms via classification as Fig. 2. In this task, 15% (at least 2) of the atoms in a crystal are masked, and the model are trained to predict the class of masked atoms. Therefore no extra labels except the structure of the crystal are required. Since there are 95 elements in S2S dataset, this pre-training is a 95-classification task. The objective of this pre-training can be defined as Eq. 9, where N and 95 denote the number of masked atoms and the number of element classes respectively. yi,j is the ground truth and ˆ yi,j is the prediction of S2SNet for "atom i is the element j", 0 for False and 1 for True.

After convergence of training, the Embedding, Aggregate, and Attention layers are cloned. Finally, S2SNet is then finetuned on S2S dataset with the pre-trained layers frozen.

4 Experiments

To verify S2SNet, a dataset of superconductors is constructed, on which extensive experiments are conducted. Compared with other models, the experiments demonstrate that S2SNet effectively captures the crystal structural information for superconductivity prediction.

4.1 Setup

S2S Dataset. We collect a total of 69,239 crystal structures from Material Project database [Jain et al., 2013], each of which contains lattice structure, formation energy, and band gap of crystals. 33,810 superconductors and their critical temperatures are obtained from SuperCon∗ , which stores the critical temperature of substances. In addition, this paper also obtains structural information of part of materials from databases such as OQMD [Kirklin et al., 2015], and ICSD [Bergerhoff et al., 1987]. The information in the databases above are used to construct S2S dataset with 1,685 crystals, which includes 421 currently known crystal superconductors and other 1,264 materials without known superconductivity sampled fromMaterial Project database. Each entry includes the name of the substance, its crystal structure(atoms, location of atoms and lattice parameters), and whether it is a superconductor. This statistics of S2S dataset and pre-training dataset are shown in Table 1.

∗ https://supercon.nims.go.jp/en/

Compared Approaches. We mainly compare our proposal with several superconductivity prediction methods, MEGNet [Chen et al., 2019], Statistic [Stanev et al., 2018], and CGCNN [Xie and Grossman, 2018]. MEGNet uses graph networks to model molecules and crystals universally for property prediction. Statistic achieves a pleasant result on predicting superconductor which calculates 145 properties of materials with MAGPIE [Baldow et al., 2017] and put these features into a random forest for training. CGCNN is specifically designed for accurate and interpretable prediction of crystal material properties where convolutional neural networks are performed on crystal graph.

Evaluation Metrics. We mainly employ Accuracy (ACC) and Area Under Curve (AUC) for superconductor prediction. In addition, we also use visualization and interpretability analysis to verify our S2SNet.

4.2 Experimental Results and Discussion

Comparison on S2S Dataset. All these models are trained and tested directly on S2S dataset. To bridge the shortage of data, MEGNet and CGCNN are pre-trained with the task of predicting formation energy following the paper [Chen et al., 2019] and S2SNet is pre-trained with the MLM task both on the Material Project database. 80% of S2S dataset are randomly selected to train the models and the left 20% are used to test the models. All the models are evaluated in the same way. Adam optimizer is used in S2SNet and the settings of other models follow their papers.

The experimental results are shown in Table 2. S2SNet consistently outperforms the other competitors both on accuracy (ACC) and Area Under Curve (AUC). For example, S2SNet gets about 6% improvements on ACC than Statistic and about 7% improvements on AUC than CGCNN without pre-training and almost 11%/14% improvements than MEGNet on ACC/AUC. With the pre-training task, all these three models, MEGNet, CGCNN and S2SNet are improved significantly. Since S2SNet is pre-trained on the unsupervised MLM task, it can be pre-trained on other much bigger dataset without any labels. However, MEGNet and CGCNN must be trained on the dataset with specific labels such as formation energy and so on. This also means that S2SNet can be further improved with larger dataset in the future work.

Fig. 3 shows the receiver operating characteristic (ROC) curve for the superconductor classifier. The overall test accuracy is 91.28%, and the area under curve for the receiver operation conditions is 0.9118.

4.3 Ablation Study

Ablation studies are also conducted to verify our S2SNet. Table 3 shows the experimental results of S2SNet with different number of attention layers on S2S dataset. Experiments demonstrate that more layers lead to better AUC. For both efficiency and accuracy, the number of attention layers in our S2SNet is 10.

Further experiments are conducted to verify that the lattice parameters are helpful in S2SNet. Table 4 shows results of 10 attention layer S2SNet without the input of lattice parameters. With lattice parameters, models perform much better no matter pre-trained or not.

4.4 Interpretability on S2SNet

According to [Zhou et al., 2018], the interpretability and reproducibility of known chemical common knowledge for machine learning models in material science is critical. To this end, the embeddings of all the element are extracted from S2SNet. The Pearson correlation coefficient between different elements are calculated as Eq. 10, where m is the dimension of elemetal embedding vector. As is shown in Fig. 4, the Pearson correlation coefficient between elemental embedding vectors follow the trend of the periodic table of elements. For example, the chalcogen, halogen, lanthanoid, transition metals, metalloid, and actinoid show highest similarities within their groups. For better observation on the elemental embedding vectors, t-SNE[Hinton and Roweis, 2002] is used to map the high-dimension embedding vectors into 2-dimension space in Fig. 5. Elements with similar properties are close to each other, such as the Actinoid, the Chalcogen, the Halogen and so on. All the visualizations show that S2SNet learns the common intuition in chemistry.

4.5 Potential SuperconductingMaterials

In order to provide more insight to the field of superconducting materials research, all the materials we have are fed into S2SNet to predict whether they are superconductors or not. Frequency of elements in materials that our model determines to be superconductors is shown in Fig. A1 in Appendix A. Orange and blue ones denote the number of atoms and the number of materials that contain corresponding atom respectively. Si, S, B, La, Fe, Se, Ge, P, Al, N, O, Bi, C, Ca, Te, Ti and Ni appear in more than 300 superconductors which, in our opinion, are supposed to be crucial elements for superconductor. And the known superconducting elements Nb, Pb, and so on are also identified as the crucial elements in S2SNet. To verify that whether our model has learned these knowledge, we make the embedding layer learnable and train the overall model with the superconductivity prediction. Atomic embeddings are visualized by t-sne in Fig. 6. The atoms with high frequency are grouped together.

We get 5468 potential superconductors out of the 69239 materials, and the element combination of part of representative potential superconductors are list in Table A1 in Appendix A. The existing Iron-based, Cuprate, Hydrogen-based and Metal superconductors are identified with high score by our model. Some potential superconductors with similar structures are also predicted. Besides, the other potential superconductors with extremely high score are also given for further experimental synthesis verification. The results are reasonable, since Graphene and SiC, which are not in the training set, are predicted to be superconductors. The prediction accuracis for Iron-based, Cuprate, Hydrogen-based superconductors are 97.64%, 92.00% and 96.89% respectively. It is worth mentioning that all the 21 Iron-based superconductors are correctly predicted.

4.6 Why Are They Superconductors

Class Activation Map (CAM) is used to explore how superconductivity comes [Kwa´sniewska et al., 2017]. CAM is a technique commonly used in CV for interpretability, i. e., highlighting the parts of a picture where the model identifies it as a specific class. Herein, we highlight the atoms in crystals that are critical for superconductivity with CAM. The experimental findings are consistent with the previously proposed theories. We take the BaFe 2 S 3 , H 3 S, and Sr 2 CuClO 2 for example. The contribution of atoms in each crystal to superconductivity is shown in Fig. 7 and the color denotes the contribution. BaFe 2 S 3 in Fig. 7(a) represents the Ironbased superconductors and the Fe is more crucial for superconductivity. H 3 S represents the Hydrogen-based superconductors. Fig. 7(b) shows that both H and S are important for superconductivity and the three hydrogen atoms are of different importance. The hydrogen atoms that are closer to the S is more crucial which is consistent to previous experimental results that with higher pressure, the atoms in a crystal are closer, which can cause the superconductivity. Sr 2 CuClO 2 represents the superconductors based on Cuprate. Fig. 7(c) demonstrate that Cu and O are crucial atoms.

5 Conclusion

A new dataset, S2S, is built as a benchmark for future machine learning based superconductivity discovery. Based on S2S, a novel pre-trained neural network, S2Snet, is built for superconductivity prediction solely from their crystal structures. S2SNet leverages attention mechanism to model the interaction of atoms in the crystal cells. With the pre-training task, MLM, massive unlabeled data can be properly utilized for S2SNet to learn better representation of atoms under interactions. Indeed S2SNet outperforms other methods significantly no matter with or without pre-training. Furthermore, it is straightforward to extend S2SNet to predict other properties of crystals as well. We expect that S2S dataset and S2SNet could inspire the machine learning community to make more impact on AI for science and SDGs. Finally, the current dataset is still relatively small. Enriching S2S dataset will be our future work.