Periodic Graph Transformers for Crystal Material Property Prediction
Anonymous Author(s)
Abstract
We consider representation learning on periodic graphs encoding crystal materials.
Different from regular graphs, periodic graphs consist of a minimum unit cellrepeating itself on a regular lattice in 3D space.
How to effectively encode these periodic structures poses unique challenges not present in regular graph representationlearning.
In addition to being E(3) invariant, periodic graph representations needto be periodic invariant.
That is, the learned representations should be invariant toshifts of cell boundaries as they are artificially imposed.
Furthermore, the periodicrepeating patterns need to be captured explicitly as lattices of different sizes andorientations may correspond to different materials.
In this work, we propose atransformer architecture, known as Matformer, for periodic graph representationlearning.
Our Matformer is designed to be invariant to periodicity and can capturerepeating patterns explicitly.
In particular, Matformer encodes periodic patternsby efficient use of geometric distances between the same atoms in neighboringcells.
Experimental results on multiple common benchmark datasets show thatour Matformer outperforms baseline methods consistently.
In addition, our resultsdemonstrate the importance of periodic invariance and explicit repeating patternencoding for crystal representation learning.
Introduction
Crystal material property prediction is important for the discovery of new materials with desirableproperties [38, 35–37, 49, 51, 41, 4, 34, 6].
Different from molecules and proteins [16, 50, 42, 9,48, 27, 44], which are commonly represented as graphs [45, 30, 12, 13, 26, 28], crystals consist ofa minimum unit cell repeating itself on a regular lattice in 3D space.
Thus, crystals are naturallyrepresented as periodic graphs.
A key challenge of crystal material property prediction lies in howto effectively encode periodic structures that are not present in regular molecular graph representations [10, 21, 41, 24, 23, 32, 25, 39, 18, 14, 1].
E(3) invariance for molecular graphs requires therepresentation for a given molecule to be invariant to translation, rotation and reflection transformations in 3D space.
Beyond that, periodic
[Figure]
[Table]
[Figure]
[Table]
+kTufA5fwHTqQqs/hHz8/wk553gRZ0rb9oeVm5mdm1/ILxaWlldW14rrGw0VxpLQOgl5KFseVpQzQeuaaU5bkaQ48DhtesOzcd68oVKxUFzrJKJugPuC+YxgbayrclLuFkt2xc6E/oLzDaXq7qf9uHO3VesW3zu9kMQBFZpwrFTbsSPtplhqRjgdFTqxohEmQ9ynbYMCB1S5adbqCO0bp4f8UJojNMrcny9SHCiVBJ6pDLAeqOlsbP6XtWPtn7opE1GsqSCTj/yYIx2i8dyoxyQlmicGMJHM9IrIAEtMtNlOwSzBmR75LzQOK85x5ejSKVXPbyFTHrZhDw7AgROowgXUoA4E+nAPT/BscevBerFeJ6U5a3LDJvyS9fYFjnGQsA==</latexit> y
In this work, we propose to tackle periodic graph representation learning by incorporating bothperiodic invariance and periodic pattern encoding.
We propose a periodic graph transformer, known as
Matformer, that is periodic invariant and can capture periodic repeating patterns explicitly for crystalrepresentation learning.
Matformer achieves periodic invariance through two uniquely designed graphconstruction methods.
It further encodes periodic patterns by efficient use of geometric distancesbetween the same atoms in neighboring cells, thereby capturing the lattice size and orientationof a given crystal.
We conduct experiments on two commonly used material benchmark datasets,including the Materials Project [17] and Jarvis [8].
Results show that Matformer outperforms baselinemethods consistently on various tasks.
In addition, our results demonstrate the importance of bothperiodic invariance and periodic pattern encoding for crystal representation learning.
Background
Crystal property prediction and crystal structures.
Given a crystal represented as ( A , P , L ) ,crystal property prediction aims to predict a target property value y , which is either real as y ∈ R orcategorical as y ∈ { 1 , 2 , · · · , C } , for regression or classification task with C classes, respectively.
Specifically, a crystal is represented as a unit cell with periodic patterns.
A unit cell is a minimumrepeatable structure for the given crystal, and it can be described by matrices A and P .
A =[ a 1 , a 2 , · · · , a n ] T ∈ R n × d a is the atom feature matrix, where a i ∈ R d a is the d a -dimensional 52 feature vector for atom i in the unit cell.
P =
[ p 1 , p 2 , · · · , p n ] T ∈ R n × 3 is the position matrix, 53 where p i ∈ R 3 contains the Cartesian coordinates for atom i in 3D space.
To further encode periodicpatterns, an additional lattice matrix L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ R 3 × 3 is used to describe how a unit cellrepeats itself in three directions, including ℓ 1 , ℓ 2 , and ℓ 3 .
We show a 2D case of periodic patternsfor easy illustration in Fig.
1 (a).
Note that crystals usually possess irregular shapes in practice.
Hence, ℓ 1 , ℓ 2 , and ℓ 3 are not always orthogonal in 3D space.
Formally, given a crystal representation( A , P , L ) , the infinite crystal structure can be represented as [Equation]
[Equation]
Here, ˆ P contains all possible positions for each atom i , associated with the same a i in ˆ A .
Multi-edge graph construction for crystals.
The multi-edge graph construction proposed by Xie and
Grossman [51] aims to capture atom interactions across cell boundaries, which are imposed artificially.
In a regular molecular graph, a node corresponds to a single atom.
In contrast, in a multi-edge graph,node i represents atom i and all its duplicates in the infinite 3D space.
Apparently, node i contains theatom features vector a i and all positions in the set { ˆ
[Equation]
Formally, the multi-edge graph construction method builds edges between nodes as follows.
Givena prefixed radius r ∈ R , if there exists any 3-tuple ( k ′ 1 , k ′ 2 , k ′ 3 ) , where k ′ 1 , k ′ 2 , k ′ 3 ∈ R , such that
the
Euclidean distance d ′ ji ∈ R satisfies d ′ ji = ||
p j + k ′ 1 ℓ 1 + k ′ 2 ℓ 2 + k ′ 3 ℓ 3 − p
i || 2 ≤ r , an edge is builtfrom j to i with the initial edge feature d ′ ji .
An example of the edge construction is shown in Fig. 1(b).
Intuitively, if there exist m positions of node j within the radius of the center node i , this methodbuilds m edges from node j to node i .
By considering all possible positions of every node within apredefined radius in 3D space, the multi-edge graph construction method can in essence capture atominteractions across cell boundaries [40, 34, 4, 41, 6].
Periodic invariance and periodic pattern encoding for crystals
Different from molecular graphs, crystal graphs consist of a minimum unit cell repeating itself on aregular lattice in 3D space.
When encoding such periodic structures, unique challenges lie in periodicinvariance and periodic pattern encoding.
In this section, we propose to formally define and analyzethe importance of these two components.
3.1 Periodic invariance for crystals
Periodic invariance is proposed based on E(3) invariance, which is defined as below.
Definition 1 (Unit Cell E(3) Invariance) .
A function f : ( A , P , L ) → X is unit cell E(3) invariantsuch that for all Q ∈ R 3 × 3 , | Q | = ± 1 and b ∈ R 3 , we have f ( A , P , L )
=
f ( A , Q P + b, Q L ) ,where Q is rotation and reflection transformations, and b is translation transformations in 3D space.
Intuitively, the structure of a cell remains the same when either applying rotations and reflections toposition matrix P and lattice matrix L together, or applying translations to P only.
Correspondingly,the output of the unit cell E(
[Table] [Table]
[Figure]
In addition to unit cell E(3) invariance,periodic invariance is also shown necessary for generating valid crystal representations.
Specifically, when thelattice matrix L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ 91 R× 3 is fixed for a crystal, we can stillobtain different position matrices P ∈ 93 R n × 3 and different unit cell structuresby shifting the period boundaries.
Asshown in Fig.
2 (a) and (b), the formedunit cell structures are different forthe same crystal by shifting periodboundaries.
To this end, we furtherintroduce periodic invariance, whichshows that when the periodic boundaries are shifted or scaled up, the periodic invariant representation shouldremain the same.
Formally, based on Sec. 2, we further define a function Φ : ( ˆ A , ˆ P , L , p) → ( A , P )simulating how to form different unit cells from a given infinite crystal structure.
For an infinitecrystal structure represented as ( ˆ A , ˆ P ) , Φ uses a corner point p and shape matrix L to form a unitcell represented as ( A , P ) .
In addition, we use α ∈ N 3+ to indicate the scaling up of a repeating unitcell formed by periodic boundaries.
Then the formal definition of periodic invariance is below.
Definition 2 (Periodic Invariance) .
A unit cell E(3) invariant function f : ( A , P , L ) → X is periodicinvariant if f ( A , P , L ) =
f (Φ( ˆ A , ˆ P , α L , p) , α L ) holds for all p ∈ R 3 and α ∈ N 3+ .
Significance of periodic invariance.
Breaking periodic invariance will result in different crystalgraphs for the same crystal.
In Fig.
2 (a) and (b), we show a graph construction method thatbreaks periodic invariance.
This method is employed by a transformer-based model, known as
Graphormer [53], which first uses radius to include all the atoms of interest in nearby cells and thenbuilds a fully connected graph.
A detailed illustration is shown in Fig.
5 in Appendix.
A.1.
Basedon the formal definition of periodic invariance, in this study, we aim to integrate such an important
[Figure] [Figure]
[Table]
1 + `
[Table]
|
[Table]
[Figure]
3.2 Periodic pattern encoding
[Figure]
As introduced in Sec. 2, L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ 120 R× 3 containing periodic patterns is another keycomponent to describe crystal structures.
Essentially, periodic patterns show how the minimumrepeatable structure ( A , P ) expands itself in infinite 3D space.
Without such periodic patternencoding in L , crystal structures are treated asfinite structures similar to molecules.
As shownin Fig.
3, the widely used multi-edge graph construction method [51, 4, 41, 40, 6] only captureslocal interactions among atoms but ignores theimportant periodic patterns.
However, such periodic repeating patterns need to be capturedexplicitly, as lattices of different sizes and orientations may correspond to different materials.
Hence, To better represent the infinite structuresof crystals, we argue that the periodic patterns
L ∈ R 3 × 3 should be explicitly taken into consideration in crystal learning.
The proposed Matformer
4.1 The proposed graph construction methods
In this section, we introduce the proposed graph construction methods for Matformer.
Our methodseffectively integrate periodic invariance and periodic pattern encoding.
Invariant crystal graph construction .
We consider two crystal graph construction methods, including radius-based graph construction and fully connected graph construction.
Both of them satisfyperiodic invariance, and mathematical proofs can be found in Appendix.
A.2.
In Fig.
2, we show byexample that treating atoms as single nodes breaks periodic invariance.
Instead, our methods followthe fashion introduced in Sec. 2 to treat node i as atom i and all its repeated duplicates.
We use the multi-edge graph construction [51] introduced in Sec. 2 as an alternative crystal graphconstruction method.
Note that although the multi-edge graph construction satisfies periodic invariance as in Fig.
4 (a), several existing works [51, 34] do not follow the settings exactly thus breakingperiodic invariance.
Within the given radius shown in Fig.
4 (a), these studies form the neighborhoodof node i by selecting t nearest neighbors ranked by geometric distances.
If there exist severaldifferent atoms with the same distance to node i , there is no deterministic way to select from them, asshown in Fig.
6 in Appendix.
A.1.
As a result, periodic invariance cannot be guaranteed.
Given the fully-connected fashion employed in Graphormer has achieved impressive performance onmolecular learning, we further propose another graph construction method for Matformer, known asthe fully connected graph construction.
This method uses a different strategy to determine neighborsfor each center node.
Specifically, for node i and j , it builds edges for the entries corresponding tothe t smallest distances in { d ij | d ij =
||
p
i − p j + k ′ 1 ℓ 1 + k ′ 2 ℓ 2 + k ′ 3 ℓ 3 || 2 , k ′ 1 , k ′ 2 , k ′ 3 ∈ Z } .
It can be 158 seen every pair i and j is connected in the constructed graph and there are t edges between them.
Overall, the used multi-edge graph construction and proposed fully-connected graph construction bothsatisfy the important periodic invariance.
Particularly, they possess great flexibility to be used in futurestudies for crystal learning.
In this study, we employ both methods as part of our proposed Matformer,and in main experiments, the radius-based method is used due to better empirical performance.
Periodic pattern encoding with self-connecting edges .
In this study, we propose to encode theimportant L =
[ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ R 3 × 3 into crystal graphs by adding self-connecting edges.
Asmentioned in Sec. 3.2, periodic patterns describe sizes and orientations of lattices of a crystalstructure, eventually determining the properties of this crystal.
A natural step for encoding suchrepeating periodic patterns is to consider the relative positions between an atom and its nearbyrepeated duplicates.
Formally, given an atom i with position p i and L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ R 3 × 3 , weneed to encode the atom’s three nearby duplicates with positions p i + ℓ 1 , p i + ℓ 2 , and p i + ℓ 3 .
Itis widely known that a direction vector ℓ i is determined by both its length || ℓ i || 2 and orientation.
Essentially, || ℓ i || 2 indicates the geometric distance between atom i and its corresponding duplicate.
However, the computing of orientation information, such as angles, usually induces high complexity.
Hence, it is not practical to encode such orientation information into transformer architectures.
To thisend, we propose to use geometric distances solely to implicitly consider the orientation information.
Specifically, we use additional distances to determine angles between any two direction vectors in L .
For example, the angle between ℓ 1 and ℓ 2 can be easily computed by || ℓ 1 || 2 , || ℓ 2 || 2 , and an additional 177 distance || ℓ 1 + ℓ 2 || 2 .
Hence, based on these three distances, we can determine lengths of ℓ 1 and ℓ 2 , 178 and the relative orientation between them.
Extensively, we use six geometric distances, including
[Equation]
ℓ
[Equation]
[Equation] [Equation]
, |
[Equation]
+
[Equation]
|
[Equation]
2 , and || ℓ 1 + ℓ 3 || 2 in our study, as shown in Fig.
4 (b).
By doing this, the length of each direction vector and the angle between any two direction vectors canall be determined.
As a result, the shape formed by lattice matrix L is then fixed.
Overall, we buildthe aforementioned six geometric distances as six self-connecting edges for node i .
By doing this, ourmodel is capable of encoding periodic patterns in L completely, resulting in a more accurate crystalrepresentation learning scheme.
Importantly, our approach also guarantees periodic invariance.
Overall, the graph construction for Matformer consists of two necessary stages, including invariantgraph construction and periodic pattern encoding.
For the first stage, we rigorously prove that themulti-edge graph construction satisfies periodic invariance in Appendix.
A.2 and show that severalprevious works [51, 34] break periodic invariance using a different neighbor selection strategy.
Additionally, we propose a fully-connected crystal graph construction method satisfying periodicinvariance, and the method could be used in future studies for crystal representation learning.
For thesecond stage, we naturally encode periodic patterns into constructed graphs by adding self-connectingedges without breaking periodic invariance.
A constructed graph in 2D case is shown in Fig.
4 (c).
4.2 Message passing scheme
Building on constructed graphs introduced in Sec.
4.1, we propose our message passing scheme for
Matformer.
Formally, we denote a constructed crystal graph as G =
( A , E ) .
Here, each a i ∈ Ais the d a -dimensional feature vector for atom i , as introduced in Sec. 2.
Particularly, e hij ∈
E isd e -dimensional feature vector for the h -th edge between nodes i and j .
We follow the regular attentionmechanism that computes query, key, and value [31, 29].
Our proposed message passing schemeis composed of three steps; those are, edge-wise attention coefficients computing, edge-wise value message computing, and node updating.
Formally, we let f ⋆ℓi denote the input feature vector of nodei for the ℓ -th layer of Matformer.
The message passing scheme of the ℓ -th layer is described as below.
In the first step, q hij , k hij and α hij for the h -th edge between i and j are computed as
[Equation]
where LN Q , LN K , and LN E denote the linear transformations to compute query, key, and edgeembedding in ℓ -th layer, respectively.
e h ′ ij is the intermediate output for e hij .
We use ◦ and | to denote
Hadamard Product and concatenation.
Note that q hij is the concatenation of three q i vectors to matchthe dimension of k hij .
By doing this, when computing α hij , q i attends each of k i , k j , and e h ′ ij forintegrating more information in attention.
Particularly, we omit the softmax to enhance the model’scapability to distinguish nodes with different degrees, and to make the whole network more efficient.
After obtaining coefficients α hij , in the second step, we compute m hij that is the message of e hij as
[Equation]
i
[Equation]
⋆ℓ
[Equation] [Equation]
where LN V and LN update are the linear transformations to compute value and the updated message,and LNorm denotes the layer normalization operation.
Finally, in the third step, we compute node i ’s feature vector f ℓi .
Specifically, we first obtain messagem i by aggregating information from node i ’s neighborhood over multiple edges, then achieve f ℓi as
[Equation]
where σ is the used activation function, and BN indicates batch normalization.
In addition, LN msgand LN fea are linear transformations to update the messages on edges and the old atom features.
Graphormer represents an effective transformer variant for molecular graph learning.
The differencesbetween Graphormer and the proposed Matformer lie in both graph construction and message passingscheme.
Firstly, Graphormer treats every atom as a single node and breaks periodic invariance,as mentioned in Sec. 3.1.
For the message passing, Graphormer uses the node-wise attention andencodes pairwise distances as attention bias.
It cannot work properly on multi-edge graphs for crystals.
While Matformer is specifically designed for multi-edge crystal graphs by performing edge-wiseattention and encoding geometric information into edge-wise messages, as described above.
Thedetailed architecture of our Matformer is shown in Fig.
7 in Appendix.
Related work
Crystal property prediction .
Several existing methods [46, 19, 20, 15] model crystals as chemicalformulas and employ sequence models to process them.
Other studies [51, 40, 34, 4, 41, 6] consider3D structures and formulate crystals as 3D graphs, then apply GNNs to learn from crystal graphs.
As crystals are essentially periodically repeated structures, the graph construction needs to considerperiodic invariance and periodic pattern encoding.
There are limited efforts to identify these twounique components.
As an early work, CGCNN [51] proposes to capture atom interactions acrossartificial cell boundaries by using multi-edge graphs described in Sec. 2.
The multi-edge graphsatisfies periodic invariance as described in Sec. 3.1, but fails to consider the important periodicpatterns, as described in Sec. 3.2.
The multi-edge graph construction method is widely used in thefollowing studies [40, 34, 4, 41, 6, 1].
Based on the constructed crystal graphs, many GNN variantshave been proposed for effective crystal representation learning [51, 40, 34, 4, 41, 6, 1].
Specifically, Nequip [1] considers E(3) equivariance for materials, and satisfies periodic invariance using multiedge graphs, but fails to capture periodic repeating patterns.
We also notice a recent work [47] forperiodic graph generation, which considers periodic graphs as finite graphs and breaks periodicinvariance.
Recently, ALIGNN [6] achieves the best performance on two major material datasets.
It uses angle information in the message passing to generate more informative and discriminativerepresentations.
However, the use of angles introduces excessive time complexity.
[Table]
Geometric GNNs and graph transformer .
Many efforts have been made to incorporate 3D geometric information in molecular learning.
Exemplary studies include SchNet [41], DimeNet [24, 23],
SphereNet [32], GemNet [25] [Table] etc.
However, these methods are designed for molecules without periodic patterns.
Recently, graph transformers [53, 11] using geometric information, e.g.,
Graphormer [53], have shown great potential on real-world graph data.
However, Graphormerconsiders neither periodic invariance nor periodic pattern encoding.
Differences with our method.
To the best of our knowledge, periodic invariance and periodic patternencoding described in Sec.
3 are rarely identified and explored in existing works for crystal propertyprediction.
CGCNN [51] breaks periodic invariance on some corner cases because it uses twelvenearest neighbors determined only by distances as described in Sec. 4.1.
In addition, previous methodsincluding CGCNN [51], SchNet [41], MEGNET [4], CYATT [40], GATGNN [34], NEQUIP [1] and
ALIGNN [6], all fail to consider the important periodic pattern encoding as introduced in Sec. 3.2.
Especially, following GAT [45], GATGNN [34] employs a very limited kind of attention mechanismthat is not conditioned on query, as explained in GATv2 [2]. As a result, the model capacity isreduced compared with the self attention mechanism employed in Matformer. In addition, the usageof softmax limits the capability of GATGNN of distinguishing nodes with different degrees, asmentioned in Sec.
For Graphormer [53], although it achieved remarkable success on the Open Catalyst Challenge [3], the employed graph construction method breaks periodic invariance whenapplied to crystals.
Compared with Graphormer, Matformer is specifically designed for crystalsconsidering both periodic invariance and periodic patterns.
Experimental studies
6.1 Experimental setup
We conduct experiments on two material benchmark datasets, including The Materials Project [17]and JARVIS [8].
The detailed descriptions for The Materials Project and JARVIS datasets are shownin Appendix.
A.3 Baseline methods include CFID [7], CGCNN [51], SchNet [41], MEGNET [4],
GATGNN [34], and ALIGNN [6].
Unless otherwise specified, for all the baseline methods, wereport the results taken from the referred papers or provided by original authors.
All Matformermodels are trained using the Adam optimizer [22] with weight decay [33] and one cycle learning ratescheduler [43].
We only slightly adjust learning rates from 0.001 and training epochs from 500 fordifferent tasks.
Detailed Matformer configurations for different tasks are provided in Appendix.
A.4.
6.2 Experimental results
The Materials Project.
We first use The Materials Project-2018.6.1 dataset [4], which contains69239 crystals, to evaluate Matformer.
We notice that previous works [51, 41, 4, 34, 6] compare witheach other either using datasets of different sizes, or using datasets with the same size but splittingthe datasets with different random seeds.
To make the comparison clear and fair, we retrain allcorresponding models using exactly the same training, validation and test sets across all methods andreport the results in Table.
To avoid confusion, we still put original results from referred papers in
Table.
1 inside parentheses.
For retrained baseline models, we provide the detailed configurations in
Appendix.
A.5.
The used metric is test MAE following previous studies [51, 41, 4, 34, 6].
[Table]
It can be seen from Table.
1 that Matformer achieves the best performance on all tasks consistently bysignificant margins.
Specifically, it reduces the formation energy by 4.5% of the second best model,which is a significant margin.
Furthermore, for Bulk Moduli and Shear Moduli tasks with only 4664training samples, Matformer achieves the best performances, indicating Matformer’s adaptive abilityto tasks of small training scales.
JARVIS dataset
[Table]
The quantitative results for Jarvis are shownin Table.
Matformer outperforms the baseline methods significantly on all of these fivetasks.
Compared with ALIGNN, Matformer hasstronger discriminative ability due to explicitencoding of periodic patterns.
Specifically, Matformer reduces Jarvis Ehull by 0.012, which is15.8% of ALIGNN.
Furthermore, Matformerachieves the best performances for Bulk Moduli and Shear Moduli in the Mateirals Project with4664 training samples, and Bandgap(MBJ) in JARVIS with 14537 training samples, indicating itsadaptive ability to tasks of various data scales.
Overall, the superior performances show the effectiveness of periodic pattern encoding in our Matformer message passing.
In addition, compared with
ALIGNN, our Matformer is more efficient.
We evaluate the efficiency of Matformer by comparingwith ALIGNN using JARVIS formation energy dataset.
The mean time of ten runs for training andinference using best model configurations of ALIGNN and Matformer are reported.
We also reportthe total number of parameters of each model.
In Table.
3, we show that Matformer is three timesfaster than ALIGNN in total training time and near three times faster in inference time, for the wholetest set. Matformer is also much lighter than ALIGNN in terms of model size.
Energy within Threshold .
Following OC20 [3], we use energy within threshold (EwT), whichmeasures the percentage of estimated energies that are likely to be practically useful when the absoluteerror is within a certain threshold, to evaluate Matformer’s capability for periodic graph learning.
Thismetric is new, but is well recognized by the community as it is useful in practice.
Due to significantperformance gaps between ALIGNN and other baseline methods on formation energy and totalenergy for these two datasets in terms of mean absolute error (MAE), we only compare Matformerwith ALIGNN.
Table.
4 shows that Matformer outperforms ALIGNN consistently for all three energyprediction tasks.
Interestingly, the performance gains of our Matformer beyond ALIGNN in terms of
EwT mainly come from more accurate energy predictions within absolute error of 0.01.
Comparedwith JARVIS, the Materials Project has 15422 more traning samples.
As a result, the percentage ofpredicted energies obtained by Matformer within 0.01 increases by 14.69%, which is much betterthan ALIGNN, revealing the huge potential of Matformer when larger crystal dataset is available.
6.3 Ablation studies
In this section, we demonstrate the importance of periodic invariance and explicit repeating patternencoding for crystal representation learning by ablation studies.
We also evaluate the building blocksparticularly designed for our Matformer.
Specifically, we conduct experiments on JARVIS formationenergy, and the test MAE is used as the quantitative evaluation metric.
We also provide ablationstudies on the use of sigmoid and layernorm instead of softmax in Matformer layer in Appendix A.6.
[Table]
[Table]
Periodic invariant graph construction .
We demonstrate the importanceof periodic invariant graph construction by comparing radius multi-edgegraph, denoted as Radius, with thegraph construction method proposedby Graphormer, denoted as OCgraph,on the exactly same Matformer architecture.
Note that the constructed crystal graphs by OCgraph are super large,containing more than n 2 edges, wheren is the atom number in a cell.
We adjust the Matformer configurations to train these large graphs ona single RTX A6000 GPU.
It can be seen from Table.
5 that when using OCgraph to our Matformer,the test MAE drops dramatically of 53% because of breaking periodic invariance, compared withradius-based multi-edge graphs in Matformer.
We also compare two periodic invariant graph construction methods described in Sec. 4.1.
We denote fully connected graph with t smallest pairwisedistances as T-fully, and use t = 3 .
The result shows that Radius is better than T-fully.
Encoding of repeating patterns .
We denote periodic pattern encoding as PE.
In Table.
5, we showthat omitting the periodic pattern encoding results in a significant drop of test MAE from 0.0325 to0.0337, revealing the importance of periodic patterns for crystal representation learning.
We show the training time per epoch, total training time, and test MAE.
[Table]
Complexity of introducing angular information.
Dropping angular information largely improves running efficiency of Matformer compared with ALIGNN.
We show that for the original crystal graph with n nodes and 6 n edges,the corresponding line graph with angles willhave 6 n nodes and 66 n edges, leading to highcomputational cost.
The detailed complexityanalysis of adding angular information are provided in Appendix.
A.7.
Additionally, we provide therunning time and performance analysis of Matformer with angle information in Table.
We use two
Matformer layers to process extra angular information, and use Radial Basis Function kernels [6]and Spherical Bessel Functions with Spherical Harmonics [24, 32, 25] to encode angles, denotedas Matformer + Angle RBF and Matformer + Angle SBF.
Table.
6 shows that introducing angularinformation will increase both the training time per epoch and in total by around 3 times, withoutmuch performance gain.
This may due to the periodic invariant graph construction and periodicpatterns encoding in Matformer already capture sufficient information to identify crystal structures.
Conclusion and limitation
In this work, we first propose to formally define periodic invariance and periodic pattern encodingfor periodic graph learning.
We then propose Matformer for periodic graph representation learning,which is invariant to periodicity and can capture repeating patterns explicitly.
Experimental results oncommon benchmark datasets show that our Matformer outperforms baseline methods consistently.
Inaddition, our results demonstrate the importance of periodic invariance and explicit periodic patternencoding for crystal representation learning.
One potential direction beyond this work is to includeangular information properly to satisfy both periodic invariance and to encode periodic patterns withrelatively low time complexity, and this is one limitation of our work.
Besides, negative societalimpacts of material discovery may apply to our work.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope? [Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] SeeConclusion and limitation section.(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix.
A.2.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main ex-
perimental results (either in the supplemental material or as a URL)?
[Yes] We have
included the detailed settings of our experiments.
We will publish our implementation
upon acceptance.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix.
A.4 and Appendix.
A.5.
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[N/A] Error bars are not necessary for baseline methods and
our method.
(d) Did you include the total amount of compute and the type of resources used (e.g.,
type of GPUs, internal cluster, or cloud provider)?
[Yes] All of these are reported in
Appendix.
A.4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
[Table]
A AppendixA.1 Cases breaking periodic invariance
We show two graph construction methods that break periodic invariance in Fig.
5 and Fig.
One is the graph construction method employed by Graphormer in OC20 [3], shown in Fig. 5.
Itbreaks periodic invariance because it treats every atom as a single node in the constructed graph.
When the periodic boundaries are shifted, the inner structure can change a lot in a unit cell.
Hence,the constructed fully connected graphs can be totally different when periodic boundaries are shifted.
[Table]
[Figure]
corresponding radius for atoms.
It can be seen from the comparison of the left and the right that the constructed graphs are totally different when the periodic boundaries are shifted.
In particular, the constructed graph on the left has three blue atoms, but the constructed graph on the right has two blue atoms.
Thus, the graph construction method employed by Graphormer in OC breaks periodic invariance for crystals.
Another is the graph construction method using nearest neighbors based only on geometric pairwisedistances, employed by CGCNN [51] and GATGNN [34], shown in Fig.
If several different atomshave the same geometric distances to the center atom, there is no deterministic way to determinewhich one to choose.
Thus, for these cases, this method breaks periodic invariance.
A.2 Proofs of periodic invariance
Notations .
Recall that we use matrices A , P , and L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ R 3 × 3 to describe a givencrystal structure.
A =
[ a 1 , a 2 , · · · , a n ] T ∈ R n × d a is the atom feature matrix, where a i ∈ R d a is 577 the d a -dimensional feature vector for atom i in the unit cell.
P =
[ p 1 , p 2 , · · · , p n ] T ∈ R n × 3 is the 578 position matrix, where p i ∈ R 3 contains the Cartesian coordinates for atom i in 3D space.
And the infinite crystal structure can be represented as ˆ P =
{ ˆ p i | ˆ p i = p
i + k 1 ℓ 1 + k 2 ℓ 2 + k 3 ℓ 3 , k 1 , k 2 , k 3 ∈ 580 Z , i ∈ Z , 1 ≤ i ≤ n } , and ˆ A
= { ˆ a i | ˆ a i = a i , i ∈ Z , 1 ≤ i ≤ n } .
Here, set ˆ P contains all possiblepositions for each atom i , and set ˆ A contains the corresponding atom feature vector for each atom i .
For crystal property prediction tasks, L representing the minimum repeating patterns is given.
In thefollowing proofs, we don’t consider the case that the provided L ′ =
α L for α ∈ N 3+ , which meansthe provided periodic patterns are not the minimum repeating patterns for a given crystal.
When Lrepresenting the minimum repeating patterns is given, the periodic invariance lies in the invariance toshifts of periodic boundaries.
We prove that the multi-edge graph construction and fully connectedgraph construction for crystals employed by our Matformer satisfy periodic invariance.
Multi-edge graph .
The Multi-edge graph satisfies the periodic invariance by forming edges betweennode i and j using all items in the set { d ij | d ij =
||
p
i − p j + k ′ 1 ℓ 1 + k ′ 2 ℓ 2 + k ′ 3 ℓ 3 || 2 , k ′ 1 , k ′ 2 , k ′ 3 ∈ 590 Z , d ij ≤ r } , where r is a prefixed threshold.
It can be seen that all pairwise Euclidean distancesbetween node i with positions { ˆ p
i
= p
i
+ k 1 ℓ
+ k 2 ℓ 2 + k 3 ℓ 3 , k 1 , k 2 , k 3 ∈ Z } ,
and node j with 592 positions { ˆ p j = p
j + k 1 ℓ 1 + k 2 ℓ 2 + k 3 ℓ 3 , k 1 , k 2 , k 3 ∈ Z } are considered, and the shifts of periodic 593 boundaries will not influence the pairwise Euclidean distances.
Thus, this method satisfies periodicinvariance.
Additionally, by using pairwise Euclidean distances, unit cell E(3) invariance is naturallysatisfied.
We note that for a given node i , the radius can be computed by a deterministic function that takes allpairwise distances between node i and all other nodes as input, and produces a real value r as output.
This will not break periodic invariance due to the fact that the deterministic function will alwaysproduce the same r for the same input { d ij | d ij = ||
p i − p
[Equation]
+ k
[Equation]
+ k
[Equation]
+ k
[Equation]
[Equation]
3 || 2 , k
[Equation]
1 , k
[Equation]
2 , k
[Equation]
3 ∈ 600 Z , 1 ≤ j ≤ n } .
One example is that we can use the 12-th smallest distance d 12 in { d ij | d ij =|
[Equation]
Z ,
[Equation]
j ≤ n } as the radius for node i .
Fully connected graph for crystals .
To construct fully connected graphs for crystals, in our design,node i represents atom with atom feature a i and all its repeats in the infinite 3D space, with positions{ ˆ p
i = p
i + k 1 ℓ 1 + k 2 ℓ 2 + k 3 ℓ 3 , k 1 , k 2 , k 3 ∈ Z } .
In the t -fully-connected graph for crystals, for nodei and j , t smallest pairwise distances in { d ij | d ij = || p
[Equation]
, k
[Equation]
, k
[Equation]
3 ∈ Z } 606 are considered, where t is a hyperparameter to control the edge number.
Similar to the proof of multiedge graph, all pairwise Euclidean distances between node i with positions { ˆ p
i
= p
i
+ k 1 ℓ
+ k 2 ℓ 2 +k 3 ℓ 3 , k 1 , k 2 , k 3 ∈ Z } ,
and node j with positions { ˆ p j = p j + k 1 ℓ 1 + k 2 ℓ 2 + k 3 ℓ 3 , k 1 , k 2 , k 3 ∈ Z } 609 are considered, and the shifts of periodic boundaries will not influence the pairwise Euclideandistances.
Hence, this method satisfies periodic invariance.
In addition, the usage of pairwise
Euclidean distances makes the unit cell E(3) invariance naturally satisfied.
A.3 Dataset descriptions
We show the detailed dataset Descriptions for two crystal datasets used in experimental studies,including The Materials Project and JARVIS.
The Materials Project dataset .
In particular, for tasks of formation energy and band gap, we directlyfollow ALIGNN [6] and use the same training, validation, and test set, including 60000, 5000, and4239 crystals, respectively.
For tasks of Bulk Moduli and Shear Moduli, we follow GATGNN [34],the recent state-of-the-art method for these two tasks, and use the same training, validation, and testsets, including 4664, 393, and 393 crystals.
In Shear Moduli, one validation sample is removedbecause of the negative GPa value.
We either directly use the publicly available codes from theauthors, or re-implement models based on their official codes and configurations to produce theresults.
Detailed configurations of these retrained models are provided in Appendix.
A.5.
The JARVIS dataset .
JARVIS is a newly released database proposed by Choudhary et al.
For JARVIS dataset, we follow ALIGNN [6] and use the same training, validation, and test set.
We evaluate our Matformer on five important crystal property tasks, including formation energy,bandgap(OPT), bandgap(MBJ), total energy, and Ehull.
The training, validation, and test set contains44578, 5572, and 5572 crystals for tasks of formation energy, total energy, and bandgap(OPT).
Thenumbers are 44296, 5537, 5537 for Ehull, and 14537, 1817, 1817 for bandgap(MBJ).
The used metricis test MAE.
The results for CGCNN and CFID are taken from ALIGNN [6], other baseline resultsare obtained by retrained models.
Detailed configurations of these retrained models are provided in
Appendix.
A.5.
[Table]
[Figure]
[Table]
Update
[Table]
latexit> f
⇤i <latexit
[Table]
Output
[Table]
i
[Table]
<
[Table] [Table]
[Figure]
A.4 Matformer configurations
[Figure]
We show the detailed configurations of our Matformer models in this section.
Notations .
a i ∈ A is the d a -dimensional feature vector for node i , e hij ∈
E is d e -dimensional featurevector for the h -th edge between nodes i and j , and f ⋆i is the input feature vector of node i for agiven layer of Matformer.
LN Q , LN K , and LN E denote the linear transformations to compute query,key, and edge embedding in a given Matformer layer, respectively.
q i , k i and v i are the computedquery, key, and value vectors after these linear transformations for node i .
e h ′ ij is the intermediateoutput in a given Matformer layer for e hij .
Crystal graph construction .
For all Matformer models used in two material datasets, including The
Materials Project and JARVIS, we use the radius-based multi-edge graph construction method.
Weuse the 12-th smallest distance between the given atom and all nearby atoms as radius, and include allnearby atoms within the radius as the neighborhood for the given atom.
It satisfies periodic invarianceas described in Appendix.
A.2.
For the encoding of periodic patterns, we use six geometric distances,including || ℓ 1 || 2 , || ℓ 2 || 2 , || ℓ 3 || 2 , || ℓ 1 + ℓ 2 || 2 , || ℓ 2 + ℓ 3 || 2 , and || ℓ 1 + ℓ 3 || 2 in our study.
For any of 646 these distances, if it is already in the radius of node i computed by the previous graph constructionmethod, it will not be added as a self-connecting edge of node i in the final graph.
Node and edge embeddings .
For each node, we map the atomic number to a 92-dimensionalembedding using CGCNN [51] atomic embedding.
We then use a linear transformation to map it to a128-dimensional vector as the input f ∗ i to the first Matformer message passing layer.
For each edge,we map the Euclidean distance to a 128-dimensional embedding using 128 RBF kernels with centersfrom 0.0 to 8.0.
It is then mapped to a 128-dimensional vector as the edge input e hij , by a nonlinearlayer followed by a linear layer.
Matformer layer .
For message passing, we use five layers of Matformer Message Passing layer, withfour attention heads.
The embedding sizes for q i , k i , v i for a single head are 128, mapped from f ∗ iusing corresponding linear transformations LN Q , LN K , and LN V .
For a given layer, different headsuse different LN Q , LN
Readout layer .
We use the mean pooling to aggregate features from all nodes in a graph and thenuse a nonlinear layer with hidden dimension 128 followed by a linear layer to obtain the scalar outputfor a crystal graph.
Training hyperparameters .
For all tasks in The Materials Project and JARVIS, we use the Adamoptimizer [22] with weight decay [33] of 1e-5 and one cycle learning rate scheduler [43].
We use thebatch size of 64.
We use mean square error as the objective function to train and mean absolute erroras the evaluation metric to validate and test.
We only slightly adjust the learning rates and trainingepochs for different tasks in The Materials Project and JARVIS, as shown in Table.
7 and Table.
8,respectively.
We use Pytorch to implement our models.
For all tasks on two benchmark datasets, weuse one NVIDIA RTX A6000 48GB GPU to train our Matformer models.
[Table] [Table]
We also show the detailed Matformer architecture in Fig.
A.5 Configurations of retrained models
We show configurations of retrained models for The Materials Project and JARVIS in this section.
SchNet [41].
We use six layers of SchNet message passing layer following the original paper, withfeature dimension of 64.
We train SchNet on these four tasks with learning rate of 5e-4 and batchsize of 64 for 500 epochs.
The Adam optimizer is used with 1e-5 weight decay.
One cycle learningrate scheduler is also used.
We use the 12-th smallest distance between the given atom and all nearbyatoms to serve as the radius for this given atom for the tasks of formation energy and bandgap in The
Materials Project and all tasks in JARVIS.
For the tasks of shear moduli and bulk moduli, We use the32-th smallest distance between the given atom and all nearby atoms to serve as the radius for everyatom to better the performance.
MEGNET [4].
Following the original paper, we use three layers of MEGNET message passinglayer with the same feature dimensions as mentioned in the paper, and use Set2Set readout function.
We train MEGNET on these four tasks with learning rate of 1e-3 and batch size of 128 for 1000epochs following the configuration settings mentioned in the original paper.
The Adam optimizeris used with 1e-5 weight decay.
One cycle learning rate scheduler is also used.
We try differenthyperparameters of crystal graph construction and batch size to better the performances of MEGNET.
In particular, for formation energy and bandgap, we use a radius of 4.0 for all atoms with batch sizeof 128.
And for bulk moduli and shear moduli and all tasks in JARVIS, we use the 12-th smallestdistance between the given atom and all nearby atoms to serve as the radius for this given atom tobetter the performance, and the batch size of 64 is chosen because of better empirical results.
CGCNN [51].
For CGCNN, we directly use the publicly available code from Xie and Grossman [51],with 128 hidden dimensions, batch size of 256, and three layers of CGCNN message passing layer.
We train CGCNN with Adam optimizer because of better empirical results.
We train CGCNN forthese four tasks using learning rate of 1e-2 for 100 epochs and 1e-3 for the next 900 epochs followingthe official code.
The radius cutoff of 8.0 is used for all atoms, and the nearest 12 neighbors areselected.
GATGNN [34].
For GATGNN, we directly use the publicly available code from Louis et al.
[34],with 128 hidden dimensions, batch size of 256 and keep other default settings.
We train GATGNN with learning rate of 5e-3 with learning rate decay milestone of 300 epochs and decay parameterof 0.5.
We train GATGNN for 500 epochs for the task of formation energy and bandgap in The
Materials Project and all tasks in JARVIS, with early stop.
The radius cutoff of 4.0 is used for allatoms, and the nearest 16 neighbors are selected according to the original paper.
ALIGNN [6].
For ALIGNN, we directly use the publicly available code from Choudhary and DeCost[6].
We use the official best model configurations of ALIGNN to train ALIGNN models on the tasksof bulk moduli and shear moduli, with learning rate of 1e-3 and batch size of 64.
In particular, weuse ALIGNN with four gcn layers and four alignn layers.
Overall, we either directly use the publicly available codes from corresponding authors [51, 34, 6], orre-implement models based on their official codes and configurations [41, 4] to produce the results inour experiments.
A.6 Building block of Matformer
We provide ablation studies about sigmoid and layernorm instead of softmax in our Matformermessage passing layer in this section.
[Table]
We evaluate the operation of sigmoid and normalization instead of softmax in our message passing, which is designedto capture the node degree information in multi-edge crystalgraphs, described in Sec. 4.2.
Compared with using softmaxwith scalar attention coefficient, denoted as Softmax-scalar, orsoftmax with vector attention coefficient, denoted as Softmaxvector, using sigmoid and normalization leads to significantperformance improvement in terms of test MAE, justifying the effectiveness of this operation in
Matformer message passing design, as shown in Table.
A.7 Complexity analysis of introducing angular information
We provide the complexity analysis of adding angular information to Matformer in this section.
Assume that we have n atoms in a single cell and thus n nodes in the original multi-edge graph.
Also assume that every node has at least 12 neighbors following the graph construction method of
Matformer and assume there are no self-connecting edges.
This will result in a graph G =
( V, E )where | V | = n and | E | = 12 / 2 n = 6 n .
When converting graph G into line graph L ( G ) , every 728 edge is treated as a node in the line graph.
So we have 6 n nodes in the line graph.
Every edge inthe original graph is connecting 2 nodes and every node has (12 − 1) other edges, resulting in 22neighboring edges for each edge in the original graph.
So we have 22 ∗ 6 n/ 2 = 66 n edges in theconverted line graph.
Compared with original graph with | V | = n and | E | = 6 n , the converted line 732 graph is super large with node number of 6 n and edge number of 66 n .