Xtal2DoS: Attention-based Crystal to Sequence Learning for Density of States Prediction
Anonymous Author(s)
AffiliationAddress email
Abstract
Modern machine learning techniques have been extensively applied to the materialsscience, especially for property prediction tasks.
A majority of these methodsaddress the scalar property predictions, while more challenging spectral propertiesremain less emphasized.
We formulate a crystal-to-sequence learning task andpropose a novel attention-based learning method, Xtal2DoS, which decodes thesequential representation of material density of states (DoS) properties by incorporating the learned atomic embeddings through attention networks.
Experimentsshow Xtal2DoS is faster than the existing models and consistently outperformsother state-of-the-art methods on four metrics for two fundamental spectral properties, phonon and electronic DoS.
Introduction
Machine learning for materials discovery has risen to prominence recently, with impressive applications to materials’ property predictions Liu et al.
[2017], Schmidt et al.
While theperformance greatly exceeds the existing models, most of the achievements are confined to individualscalar quantities Xie and Grossman [2018], Isayev et al.
[2017], De Jong et al.
[2016], Zhuo et al.[2018].
For conventional machine learning models, some other structured properties are often hardto deal with, for example spectral properties Eismann [2012].
Two fundamental spectral propertiesare phonon density of states (phDoS) Osborn et al.
[2001] and electronic density of states (eDoS)
Heremans et al.
In a given frequency interval, phDoS describes the number of phonon modesif the density of wave vectors in the Brillouin zone is homogeneously distributed.
The phonon densityof states is typically normalized to one Parlinski [2010].
Similarly, eDoS depicts the number ofdifferent states at a particular energy level that electrons are allowed to occupy (namely, the numberof electron states per unit volume per unit energy) Omar [1975].
Mathematically, the density of statescan be regarded as a sequence depicting the fundamental characteristics of particles.
Each crystallinematerial is associated with unique density of states (phDoS, eDoS).
The goal of density of statesprediction is to estimate this sequence given an input structure.
There are two challenges in this sequence prediction problem.
First, the input crystals are 3Dstructures, which are very different from typical machine learning formats (i.e., vectors or matrices).
Featurizing the crystal structures is the cornerstone of the whole learning process.
If useful information cannot be extracted effectively, the follow-up predictions would be severely hampered.
Essentially,each crystal structure is a collection of atoms.
Each atomic site contains a specific element, and thesesites can be seen as nodes on a graph.
Atoms also have interactions with each other Martin [2020].
Close atoms often have strong bonds.
These bonds can be represented by the connecting edges on the graph Xie and Grossman [2018].
Conventional CNN Krizhevsky et al.
[2012] and RNN Hochreiterand Schmidhuber [1997] are not applicable to graphs, while graph neural networks Hamilton et al.[2017] are specifically designed for such structures.
The pioneering work applying GNN networks tomaterials property prediction was CGCNN Xie and Grossman [2018], where the crystal unit cell wasformalized as a graph with atoms as nodes and bonds as edges.
MEGNet Chen et al.
[2019] furtherintroduced a global state to add extra features like temperature, pressure, and entropy.
More recently,
GATGNN Louis et al.
[2020] suggested using an attention mechanism to capture the global crystalstructure beyond the local atomic environments.
In GATGNN, the local augmented graph-attentionlayers encode the properties of local environments for each atom to obtain an abstract feature, and aglobal attention layer aggregates these atomic vectors by weighted averaging to produce the globalfeature vector for prediction.
This exploitation of structure markedly facilitated feature extractionfrom the crystalline systems, and built an extensible foundation for various predictions.
The second challenge in sequence prediction is the target format.
Works like GATGNN focused onthe prediction of scalar properties such as formation energies, band gaps and elastic moduli.
Theperformance of these models degrades when applied to spectral properties like eDoS and phDoS
Kong et al.
The simplest way to adapt the previous methods to predict sequences is to expandthe last layer to output multiple values instead of a single value.
The general idea is similar topredicting multiple scalar properties simultaneously.
However, such parallel prediction overlooks thesequential nature of the target DoS.
The smoothness property cannot be enforced onto the predictedtargets either.
Mat2Spec Kong et al.
[2022] observes this flaw and re-formulates the problem as amulti-target regression task.
Mat2Spec keeps the framework of expanding the last layer to outputmore values, but additionally uses contrastive learning Henaff [2020] and a probabilistic variationalautoencoder Doersch [2016] to model the correlations among different locations of the sequence.
In this work, we introduce Crystal-to-DoS (Xtal2DoS) learning, to improve both the graph encodingand the sequence decoding for phDoS and eDoS predictions.
First, we inherit the general frameworkof GATGNN and employ graph attention networks (GAT) Veliˇckovi´c et al.
[2017] to learn both thelocal and global attention.
Second, we adopt the attention-based transformer to decode the targetsequence Vaswani et al.
The transformer reads the one-hot position embedding togetherwith the global crystal representation, and decodes left-to-right auto-regressively.
Differing from themulti-target regression perspective, Xtal2DoS takes the sequence-to-sequence (seq2seq) learningperspective.
In seq2seq, the decoding process incorporates both self-attention and source-attention.
Self-attention retrieves the historical memory from the previous steps, and source-attention derivesthe similarity between the decoding sequence and the input sequence.
In this DoS prediction task, theinput is a graph rather than a sequence, so the source-attention attends each position with all the atomembeddings.
The resulting model, Xtal2DoS, thus improves both the encoding and decoding overprior models, for this graph-to-sequence learning problem.
Experimental results show our methodcan outperform the existing state-of-the-art models on all four evaluation metrics by over 15% onaverage for phDoS, and over 4% on average for eDoS.
Crystal-to-DoS Learning2.1 Problem Formulation
The input crystal material x is denoted as a graph G x = ( V, E ) .
V = { u i } ni
=1 is the collection of natoms where u i is an atom embedding learnt from CGCNN.
E = { e ij } 1 ≤
i,j ≤ n denotes the bondsconnecting atoms.
e ij depicts the distance between two atoms i, j .
For a crystalline system, thecrystal structure repeats in the 3D space and each atom has a 3D coordinate.
Though theoretically,each atom has connections with any other atom, for simplicity and generality, we choose n max _ nbrneighbors to be the upper bound for the number of neighbors of each atom.
In other words, the sumof each row or column of G x ’s adjacency matrix is smaller than or equal to n max _ nbr .
Our goal is topredict y for each x .
y is the desired density of states (phDoS or eDoS), with length l y .
l y = 51 forphDoS and l y =
for eDoS. l is predefined for each type of DoS. The learnt model maps x to y( f : x → y ).
Figure 1 presents an example (caesium fluoride , CsF ) of DoS prediction.
2.2 Graph Attention Encoder
As the prototype of GAT, graph convolutional network (GCN) Kipf and Welling [2016] combinesthe local graph structure and node-level features to yield impressive node predictions.
However, thefeature aggregation of GCN depends on structure and is mostly linear, constraining effectiveness ofmessage passing on the graph.
In GCN, a graph convolution operation can be written as
[Equation]
σ ( · ) is a non-linear activation function.
W l is the feature transformation matrix in the l -th layer.
h liis the latent node embedding for node i in the l -th layer, with h 0 i
= u i .
c ij = (cid:112) | N ( i ) ||
N ( j ) | is anormalization constant and N ( i ) is the set of node i ’s neighbors.
In Xtal2DoS, we use a variant of GAT, UniMP Shi et al.
[2020], to encode the crystalline graph.
Compared to GAT, UniMP adds support for edge features.
For each layer in UniMP, h li is first linearlytransformed to queries, keys and values
[Equation]
Similarly, we also transform the edge features between node i and node j
[Equation]
The edge feature transformation is not layer-dependent and W e , b e are shared for all edges.
Theattention coefficients are then computed as follows
[Equation]
where ⟨ q, k ⟩ = exp( q T k √ d ) and d is the dimensionality of q li , k li or v li
Afterwards, the embeddingsfrom neighbors are aggregated, together with the edge embeddings, weighted by attention scores
[Equation]
From Eq.
1 and Eq.
5, GAT considers similarity scores and edge information when performing theaggregation to gain more expressiveness.
Finally, motivated by the success of residual neural networks He et al.
[2016], UniMP appends agated residual connection before passing ˆ h l +
i to the next layer
[Equation] [Equation]
where ; denotes concatenation, a l is the learnable weight parameter, and f LN represents the layernorm.
Note that for the last layer of UniMP, σ ( · ) and f LN are skipped.
[Figure]
GAT learns embeddings for each node, h Li , if there are L layers in total.
GATGNN and Mat2Specsimply take the average of all the embeddings as the global feature vector
[Equation]
Then the global feature vector is used to decode the target DoS, using MLP or VAE.
We compared various sequence models including RNN and transformer for decoding.
By default, weused the GRU model for RNN.
RNN Figure 2a sketches the RNN decoding of the GAT-encoded crystal.
RNN is naturally suitablefor sequence decoding.
It reads the global feature vector as the cell state and a special ⟨ start ⟩ symbolto indicate the start of the decoding.
The output at each step is also the input for the next step.
However, RNN is criticized for lacking long-term memory Zhao et al.
When we decode thelast step ( y 127 ), the memory of the first step ( y 0 ) is probably degraded.
This might hurt the predictionquality for long sequences.
Another drawback of RNN is the speed.
Each step has to wait for theprevious output, and thus the decoding process could take a very long time.
Chunk RNN Figure 2b gives an overview of a chunk RNN, which is a simple RNN variant.
Chunk
RNN predicts a segment rather than a single value at each step.
The benefit is that the decoding timecould be greatly shortened.
For instance, on eDoS, if the segment length is 32, then sequence lengthcan be reduced from 128 to 4.
On the other hand, each segment still waits for the generation of theprevious segment, and longer segments harm the accuracy gradually.
Chunk RNN + Attention To compensate for the performance degradation brought by chunk RNN,we further introduce the attention mechanism into the sequence model (Figure 3).
The attentionmodule directly attends to the latent sequence states with the atom embeddings.
The latent states cantake specific atoms into consideration and build more fine-grained correlations.
In practice, we alsofind including the attention module boosts the evaluation results.
Transformer Given that we have already introduced the attention mechanism, it is intuitive tocompletely drop the recurrent framework and embrace the fully attentional transformer Vaswani et al.
The input at each step is the concatenation of the positional embeddings and the globalfeature vector.
A positional embedding is simply a one-hot vector of length l y indicating the index ofthe current step.
The global feature vector h x additionally provides the global context for decoding.
Formally, let z i =
[ s i ; h x ] where s i is the one-hot position embedding with length l y .
z i is the initialstate for the first self-attention layer.
As in UniMP, we derive collections of queries, keys and valuesfor the hidden states of each self-attention layer, denoted by Q, K, and V, respectively.
Transformer adopts dot-product attention to calculate the outputs
[Equation]
where d is the latent dimensionality.
Multiple sets of Q, K, and V matrices can be learned, andeach set is known as a head.
Multiple heads (multi-head) can give self-attention greater power ofdiscrimination.
The outputs of all the heads are concatenated, layer-normalized and decoded througha fully connected feed-forward network.
A more detailed depiction can be found in Vaswani et al.[2017].
During the decoding phase, each attention layer also computes the source-attention between theoutput states and the input atom nodes.
If we denote the queries, keys, values in decoding as
Q dec , K dec , V dec , and the ones from the input graph encoding as Q enc , K dec , V dec , we have
[Equation]
Source
[Equation]
for the sequential decoding.
Source-attention usually happens after self-attention.
Furthermore,the transformer model adds masking to prevent the current positions from attending to the futurepositions, to prevent leftward information flow in the decoder, and to preserve the auto-regressiveproperty, ensuring each decoded state only knows the history but not the future.
The final generatedlength- l y sequence ˆ y is compared with the ground-truth y through Kullback–Leibler divergence(KL-div) Kullback [1997] or its generalized version, Bregman Divergence (BD) Bregman [1967],resulting in the loss term, L = D KL ( y || ˆ y ) .
We chose KL-div and BD because they are more suitablethan other measures for data containing peaks, which are ubiquitous in DoS.
Our model is a synergistic combination of a GAT (UniMP) encoder and a transformer decoder.
Thewhole model is entirely based on attentions on graphs and sequences.
Figure 4 illustrates the mainidea of Xtal2DoS.
Related Works
Learning for 3D structures has been a long-standing problem.
In materials discovery, two atomicsystems can be exactly the same even though the coordinates are totally different in the 3D spaceif they can be transformed to one another through translation, rotation or inversion.
For standardmachine learning models, it is very hard to capture such invariance and symmetry Cohen and Welling[2016].
There are in general three strategies to address the concern.
First, the same models can be trained using more data Shorten and Khoshgoftaar [2019].
Thestructured inputs can be randomly transformed without changing the labels.
Some prior workssuccessfully applied data augmentation to exploring materials design spaces Kim et al.
Onthe other hand, in materials science, even with augmentation, it is not easy for machine learningmodels to comprehend the crystalline system, let alone the data collection and augmentation requireenormous domain knowledge and expert effort Schmidt et al.
Second, rather than expanding the scale on the dataset side, some researchers have explored thepossibility of enforcing the invariance or equivariance inside the model.
This kind of model handlesthe coordinates directly and are often called equivariant neural networks Satorras et al.
[2021], Duet al.
One notable model for materials science applications is
E3NN performs training based on radial functions and spherical harmonics Esteves et al.
Intuitively, E3NN represents the atoms with the 3D coordinates, mass and velocity, and tells theequivariant convolutional filters how the crystals might transform.
Thus, E3NN can capture the fullcrystal symmetry, while reducing training set sample complexity.
However, radial functions andspherical harmonics may limit the expressiveness of the model.
In some scenarios where the data isabundant, E3NN’s benefits might become marginal.
Third, it is often helpful to pre-process the data for simplicity Kumar [2018].
Instead of using thecoordinates, we can re-format the crystals to graphs by retaining only the atoms and the bondsbetween them.
Some of the structure information is lost from the input, such as the lattice angles,but this could also reduce the risk of overfitting.
Once the inputs are converted to graphs, maturegraph neural network techniques Wu et al.
[2020] can be applied.
A representative work is GATGNN
Louis et al.
GATGNN is one of the primary methods to encode the atom embeddings andextract the global feature vector.
It still leads in the performance of many scalar property predictiontasks including absolute energy, Fermi energy, band gap and Poisson-ratio.
Nevertheless, in sequenceproperty predictions, its performance degrades significantly Kong et al.
Mat2Spec Kong et al.[2022] uses the same graph encoder as GATGNN, but improves the decoder.
GATGNN only usesan MLP, which is incapable of sequence decoding.
Mat2Spec not only aligns the inputs and targetsin a latent multivariate Gaussian space, but also learns an embedding for each energy level throughcontrastive learning.
Mat2Spec views the DoS prediction as a multi-target regression problem, yetignores the sequential structure of DoS. Besides, Mat2Spec simply deepens the MLP decoder andplugs contrastive modules into the decoder, while sequence modeling is cleaner and simpler.
Xtal2DoS follows the design of GATGNN and Mat2Spec because the graph encoding is shown tooutperform the coordinate encoding Kong et al.
On top of the existing designs, Xtal2DoSfurther replaces GAT encoder with a more advanced variant, UniMP, and replaces the decoder with apowerful sequence model, transformer.
Evaluation results and ablation studies prove the effectivenessand efficiency of our introduced model.
Experiments on DoS Prediction
We validate the performance of Xtal2DoS on two datasets, phDoS and eDoS. The task of directlypredicting phDoS in crystalline solids from atomic species and positions was first proposed in E3NN
Chen et al.
We also include a new version of E3NN, denoted by E3NN-new 1 .
phDoS isnormalized to 1 since phDoS gives the number of modes per unit frequency per unit volume ofreal space and is essentially a distribution.
phDoS is critical for calculating properties such as theaverage phonon frequency and the heat capacity at 300K.
The analogous task of predicting eDoS fornonmagnetic materials was proposed in Kong et al.
eDoS is the same as phDoS except that itdepicts electrons.
Materials scientists are interested in a small energy range between -4 to 4 eV withrespect to band edges with 63 meV intervals, including both unoccupied and occupied states.
Fermienergy and band edges are all set to 0 eV on this energy grid.
Predicting eDoS accurately can be veryvaluable to bandgap energy estimation.
Both phDoS and eDoS are continuous functions in nature.
We sample 51 points on phDoS and 128 points on eDoS spanning the frequency or energy range.
Inother words, l y = 51 for phDoS and l y = 128 for eDoS.
4.1 Evaluation Metrics
Four metrics are adopted for evaluation: coefficient of determination ( R 2 ), mean absolute error(MAE), mean squared error (MSE) and Wasserstein distance (WD).
R 2 measures the proportion of the variation in the ground-truth targets y that is predictable from thepredictions ˆ y ,
[Equation]
[Table] [Table]
where ¯ y = 1 l y (cid:80) i
y i .
In the best case when y and ˆ y match exactly, R 2 would reach the maximumpossible value, 1.
If y i =
¯ y , R 2 = 0 .
Note that R 2 score can be negative if the predictions are worse.
MAE and MSE are standard metrics for regression.
We also adopt them for evaluation:
[Equation]
Wasserstein distance is another metric defined between probability distributions.
If we view the twodistributions as two piles of earth, WD is the minimum amount of dirt to move in order to turn onepile to be the same as the other.
Therefore, WD is also known as earth mover’s distance.
We definethe WD as
[Equation]
for distributions p and q with cumulative distribution functions (CDF) P and Q , respectively.
Thecontinuous form of WD can also be adapted for discrete cases.
4.2 Implementation
The graph encoder is built with PyTorch Geometric Fey and Lenssen [2019] modules.
The encoderis composed of three GAT/UniMP layers, each followed by batch normalization Ioffe and Szegedy[2015] and a non-linear activation He et al.
The sequence decoder is a stack of six selfattention layers, where each layer has subsequent layer normalization and feed-forward layers.
Thewhole model is trained with the Adam optimizer Kingma and Ba [2014], on Nvidia V100 GPU.
4.3 Prediction Results
Xtal2DoS is first evaluated on phDoS, compared with other state-of-the-art baselines, including E3NN,
GATGNN, Mat2Spec and some of their variants.
Among all the compared methods, Xtal2DoS givesthe best numbers.
On R 2 , MAE, MSE and WD, the relative improvements over the best comparedbaseline are as large as 9.81%, 8.81%, 16.58%, 7.21%, respectively.
The average improvement is10.60%.
Table 2 shows the performance gain achieved by Xtal2DoS over other competitive baselines on themore challenging eDoS dataset.
Xtal2DoS outperformed the best baseline on the four metrics by3.96%, 2.71%, 4.98%, 5.22% respectively.
The average improvement is 4.22%, which is a noticeableincrement especially when the training speed is faster (see section 4.6).
4.4 Ablation Study
To validate the neural module selection of Xtal2DoS, we further compare the performance of RNN,
RNN+attention, Chunk RNN+attention, and the transformer (Xtal2DoS).
We use the same fourmetrics to compare these settings.
In Table 3 and 4, transformers always dominate other RNN-basedmodels on all the metrics.
The improvements shown in the last row indicate the comparison between
Xtal2DoS and the best compared model (usually Chunk RNN+Attn).
[Table] [Table]
[Table]
4.5 Qualitative Results
We also show some qualitative results on both phDoS and eDoS in Figure 5 and 6 (in Appendix).
Both figures demonstrate the accuracy of our predictions.
The red and blue lines are always closelycorrelated.
4.6 Training Speed
One major advantage of transformer models is parallelism.
Unlike RNNs where each step hasto wait for the previous output, transformer computes the outputs at all the steps simultaneouslyand enforces auto-regressiveness through masking.
These operations can be performed in matrixform, taking advantage of powerful GPUs, and thus substantially reduce the training cost.
Table5 shows the training time per epoch for Mat2Spec, RNN-based sequence models and Xtal2DoS.
Mat2Spec does not involve any sequential model and generates outputs through MLP, so it is fasterthan any RNN-based sequence decoder model.
Mat2Spec instead integrates probabilistic samplingand contrastive learning, so it is not surprising that it takes more time than Xtal2DoS. Unsurprisingly,
RNN trains more than 50% slower than Mat2Spec.
Adding the attention module further prolongs thetime cost.
RNN+Attention doubles the training cost compared with Mat2Spec.
This also motivatedus to attempt chunk RNN.
As a result, chunk RNN+Attention can complete a training epoch fasterthan RNN and RNN+Attention, but is still slower than Mat2Spec.
Xtal2DoS is the fastest methodamong all these models.
It only takes one third of Mat2Spec’s training time, one fifth of RNN’straining time, and one fourth of chunk RNN+Attention’s training time.
Xtal2DoS excels at speedbecause of its parallel design and its succinct structure to avoid extra neural module.
Faster modelsare often more advantageous in training, tuning and deploying.
Conclusion
In this work, we introduce a novel graph-to-sequence learning framework, Xtal2DoS , for predictingthe spectral properties of materials, phDoS and eDoS , from the input crystalline structures.
Ourmodel adopts an improved graph attention network (to handle edges) as the graph encoder, and atransformer as the decoder for sequential decoding.
Xtal2DoS is fully attentional and captures allthe atom-atom, atom-sequence and step-step correlations.
Therefore, it considerably exceeds otherexisting state-of-the-art baselines, in quantitative evaluation metrics, qualitative matching results andtraining speed.