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 [1, 2].
While the performance greatly exceeds the existingmodels, most of the achievements are confined to individual scalar quantities [3–6].
For conventionalmachine learning models, some other structured properties are often hard to deal with, for examplespectral properties [7].
Two fundamental spectral properties are phonon density of states (phDoS)[8] and electronic density of states (eDoS) [9].
In a given frequency interval, phDoS describes thenumber of phonon modes if the density of wave vectors in the Brillouin zone is homogeneouslydistributed.
The phonon density of states is typically normalized to one [10].
Similarly, eDoS depictsthe number of different states at a particular energy level that electrons are allowed to occupy (namely,the number of electron states per unit volume per unit energy)
Mathematically, the density ofstates can be regarded as a sequence depicting the fundamental characteristics of particles.
Eachcrystalline material is associated with unique density of states (phDoS, eDoS).
The goal of density ofstates prediction 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)
Conventional CNN [13] and RNN [14] are not applicable to graphs while graph neural networks[15] are specifically designed for such structures.
The pioneering work applying GNN networks to materials property prediction was CGCNN [3], where the crystal unit cell was formalized as a graphwith atoms as nodes and bonds as edges.
MEGNet [16] further introduced a global state to add extrafeatures like temperature, pressure, and entropy.
More recently, GATGNN [17] suggested using anattention mechanism to capture the global crystal structure beyond the local atomic environments.
In GATGNN, the local augmented graph-attention layers encode the properties of local environmentsfor each atom to obtain an abstract feature, and a global attention layer aggregates these atomicvectors by weighted averaging to produce the global feature vector for prediction.
This exploitation ofstructure markedly facilitated feature extraction from the crystalline systems, and built an extensiblefoundation 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 [18].
The simplest way to adapt the previous methods to predict sequences is to expand the last layer tooutput multiple values instead of a single value.
The general idea is similar to predicting multiplescalar properties simultaneously.
However, such parallel prediction overlooks the sequential natureof the target DoS.
The smoothness property cannot be enforced onto the predicted targets either.
Mat2Spec [18] observes this flaw and re-formulates the problem as a multi-target regression task.
Mat2Spec keeps the framework of expanding the last layer to output more values, but additionally usescontrastive learning [19] and a probabilistic variational autoencoder [20] to model the correlationsamong 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) [21] to learn both the local and globalattention.
Second, we adopt the attention-based transformer to decode the target sequence [22].
Thetransformer reads the one-hot position embedding together with the global crystal representation,and decodes left-to-right auto-regressively.
Differing from the multi-target regression perspective, Xtal2DoS takes the sequence-to-sequence (seq2seq) learning perspective.
In seq2seq, the decodingprocess incorporates both self-attention and source-attention.
Self-attention retrieves the historicalmemory from the previous steps, and source-attention derives the similarity between the decodingsequence and the input sequence.
In this DoS prediction task, the input is a graph rather than asequence, so the source-attention attends each position with all the atom embeddings.
The resultingmodel, Xtal2DoS, thus improves both the encoding and decoding over prior models, for this graphto-sequence learning problem.
Experimental results show our method can outperform the existingstate-of-the-art models on all four evaluation metrics by over 10% on average for phDoS, and over4% 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, the crystal 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) [23] combines the local graph structureand node-level features to yield impressive node predictions.
However, the feature aggregation of GCN depends on structure and is mostly linear, constraining effectiveness of message passing on thegraph.
In GCN, a graph convolution operation can be written as
[Equation]
[Equation]
σ ( · ) is a non-linear activation function.
W l is the feature transformation matrix in the l -th layer.
h l iis 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 [24], to encode the crystalline graph.
Compared to GAT, UniMP adds support for edge features.
For each layer in UniMP, h li is first linearly transformedto 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 [25], UniMP appends a gated residualconnection before passing ˆ h l [-+
-]+1 i to the next layer
[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.
2.3 Sequential Decoding
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]
[Figure] [Figure] 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 [26].
When we decode the last step ( y 127 ),the memory of the first step ( y 0 ) is probably degraded.
This might hurt the prediction quality for longsequences.
Another drawback of RNN is the speed.
Each step has to wait for the previous 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 [22].
Theinput at each step is the concatenation of the positional embeddings and the global feature vector.
A positional embedding is simply a one-hot vector of length l y indicating the index of the currentstep.
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 initial state forthe first self-attention layer.
As in UniMP, we derive collections of queries, keys and values for thehidden 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 [22].
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]
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) [27] or its generalized version, Bregman Divergence (BD) [28], resulting in the loss term, L = D KL ( y || ˆ y ) .
We chose KL-div and BD because they are more suitable than other measures fordata 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 [29].
There are ingeneral three strategies to address the concern.
First, the same models can be trained using more data [30].
The structured inputs can be randomlytransformed without changing the labels.
Some prior works successfully applied data augmentationto exploring materials design spaces [31].
On the other hand, in materials science, even withaugmentation, it is not easy for machine learning models to comprehend the crystalline system, letalone the data collection and augmentation require enormous domain knowledge and expert effort[2].
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 [32, 33].
One notable modelfor materials science applications is E3NN [34].
E3NN performs training based on radial functionsand spherical harmonics [35].
Intuitively, E3NN represents the atoms with the 3D coordinates, massand velocity, and tells the equivariant convolutional filters how the crystals might transform.
Thus, E3NN can capture the full crystal symmetry, while reducing training set sample complexity.
However,radial functions and spherical harmonics may limit the expressiveness of the model.
In some scenarioswhere the data is abundant, E3NN’s benefits might become marginal.
Third, it is often helpful to pre-process the data for simplicity [36].
Instead of using the coordinates,we can re-format the crystals to graphs by retaining only the atoms and the bonds between them.
Some of the structure information is lost from the input, such as the lattice angles, but this could alsoreduce the risk of overfitting.
Once the inputs are converted to graphs, mature graph neural networktechniques [37, 38] can be applied.
A representative work is GATGNN [17].
GATGNN is one ofthe primary methods to encode the atom embeddings and extract the global feature vector.
It stillleads in the performance of many scalar property prediction tasks including absolute energy, Fermienergy, band gap and Poisson-ratio.
Nevertheless, in sequence property predictions, its performancedegrades significantly [18].
Mat2Spec [18] uses the same graph encoder as GATGNN, but improvesthe decoder.
GATGNN only uses an MLP, which is incapable of sequence decoding.
Mat2Specnot only aligns the inputs and targets in a latent multivariate Gaussian space, but also learns anembedding for each energy level through contrastive learning.
Mat2Spec views the DoS prediction asa multi-target regression problem, yet ignores the sequential structure of DoS. Besides, Mat2Specsimply deepens the MLP decoder and plugs contrastive modules into the decoder, while sequencemodeling is cleaner and simpler.
Xtal2DoS follows the design of GATGNN and Mat2Spec because the graph encoding is shown tooutperform the coordinate encoding [18].
On top of the existing designs, Xtal2DoS further replaces GAT encoder with a more advanced variant, UniMP, and replaces the decoder with a powerfulsequence model, transformer.
Evaluation results and ablation studies prove the effectiveness andefficiency 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[34].
We also include a new version of E3NN, denoted by E3NN-new 1 .
phDoS is normalized to1 since phDoS gives the number of modes per unit frequency per unit volume of real space and isessentially a distribution.
phDoS is critical for calculating properties such as the average phononfrequency and the heat capacity at 300K.
The analogous task of predicting eDoS for nonmagneticmaterials was proposed in [18].
eDoS is the same as phDoS except that it depicts electrons.
Materialsscientists are interested in a small energy range between -4 to 4 eV with respect to band edges with63 meV intervals, including both unoccupied and occupied states.
Fermi energy and band edgesare all set to 0 eV on this energy grid.
Predicting eDoS accurately can be very valuable to bandgapenergy estimation.
Both phDoS and eDoS are continuous functions in nature.
We sample 51 pointson phDoS and 128 points on eDoS spanning the frequency or energy range.
In other words, l y = 51for 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]
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]
[Table]
[Table] [Table] [Table] 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 [39] modules.
The encoder is composed ofthree GAT/UniMP layers, each followed by batch normalization [40] and a non-linear activation [41].
The sequence decoder is a stack of six self-attention layers, where each layer has subsequent layernormalization and feed-forward layers.
The whole model is trained with the Adam optimizer [42], 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]
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.