Periodic Graph Transformers for Crystal Material Property Prediction
Abstract
We consider representation learning on periodic graphs encoding crystal materials.
Different from regular graphs, periodic graphs consist of a minimum unit cell repeating itself on a regular lattice in 3D space.
How to effectively encode these periodic structures poses unique challenges not present in regular graph representation learning.
In addition to being E(3) invariant, periodic graph representations need to be periodic invariant.
That is, the learned representations should be invariant to shifts of cell boundaries as they are artificially imposed.
Furthermore, the periodic repeating patterns need to be captured explicitly as lattices of different sizes and orientations may correspond to different materials.
In this work, we propose a transformer architecture, known as Matformer, for periodic graph representation learning.
Our Matformer is designed to be invariant to periodicity and can capture repeating patterns explicitly.
In particular, Matformer encodes periodic patterns by efficient use of geometric distances between the same atoms in neighboring cells.
Experimental results on multiple common benchmark datasets show that our Matformer outperforms baseline methods consistently.
In addition, our results demonstrate the importance of periodic invariance and explicit repeating pattern encoding for crystal representation learning.
Our code is publicly available at https://github.com/YKQ98/Matformer .
Introduction
Crystal material property prediction is important for the discovery of new materials with desirable properties [38, 35–37, 50, 52, 41, 4, 33, 6].
Different from molecules and proteins [15, 51, 42, 9, 49, 26, 44], which are commonly represented as graphs [45, 29, 11, 12, 25, 27, 34], crystals consist of a minimum unit cell repeating itself on a regular lattice in 3D space.
Thus, crystals are naturally represented as periodic graphs.
A key challenge of crystal material property prediction lies in how to effectively encode periodic structures that are not present in regular molecular graph representations [10, 20, 41, 23, 22, 31, 24, 39, 17, 13, 1, 47].
E(3) invariance for molecular graphs requires the representation for a given molecule to be invariant to translation, rotation and reflection
[Figure]
[Table]
[Figure]
[Table]
+kTufA5fwHTqQqs/hHz8/wk553gRZ0rb9oeVm5mdm1/ILxaWlldW14rrGw0VxpLQOgl5KFseVpQzQeuaaU5bkaQ48DhtesOzcd68oVKxUFzrJKJugPuC+YxgbayrclLuFkt2xc6E/oLzDaXq7qf9uHO3VesW3zu9kMQBFZpwrFTbsSPtplhqRjgdFTqxohEmQ9ynbYMCB1S5adbqCO0bp4f8UJojNMrcny9SHCiVBJ6pDLAeqOlsbP6XtWPtn7opE1GsqSCTj/yYIx2i8dyoxyQlmicGMJHM9IrIAEtMtNlOwSzBmR75LzQOK85x5ejSKVXPbyFTHrZhDw7AgROowgXUoA4E+nAPT/BscevBerFeJ6U5a3LDJvyS9fYFjnGQsA==</latexit> y transformations in 3D space.
Beyond that, periodic graphs of crystals require unique periodic invariance.
Periodic invariance has two facets; those are, the learned representations should be invariant to both the scaling up of the minimum repeatable unit cells and shifts of periodic boundaries.
Although the former has been considered in Xie et al.
[52, 53], the later is rarely identified explicitly and sometimes not considered by previous studies [52, 4, 46, 41, 6, 53, 33, 54, 1, 5].
Furthermore, the repeating patterns of periodic graphs should be captured explicitly.
Given a fixed minimum unit cell structure, lattices of different sizes and orientations may correspond to different materials.
Without such periodic patterns, the infinite structures of crystals may not be represented accurately.
However, the explicit encoding of periodic patterns is not explored in previous studies [52, 4, 41, 6, 33, 1].
In this work, we propose to tackle periodic graph representation learning by incorporating both periodic 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 crystal representation learning.
Matformer achieves periodic invariance through two uniquely designed graph construction methods.
It further encodes periodic patterns by efficient use of geometric distances between the same atoms in neighboring cells, thereby capturing the lattice size and orientation of a given crystal.
We conduct experiments on two commonly used material benchmark datasets, including the Materials Project [16] and Jarvis [8].
Results show that Matformer outperforms baseline methods consistently on various tasks.
In addition, our results demonstrate the importance of both periodic 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 or categorical 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 minimum repeatable 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 feature vector for atom i in the unit cell.
P =
[ p 1 , p 2 , · · · , p n ] T ∈ R n × 3 is the position matrix, where p i ∈ R 3 contains the Cartesian coordinates for atom i in 3D space.
To further encode periodic patterns, an additional lattice matrix L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ R 3 × 3 is used to describe how a unit cell repeats itself in three directions, including ℓ 1 , ℓ 2 , and ℓ 3 .
We show a 2D case of periodic patterns for 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]
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 [52] 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 the atom features vector a i and all positions in the set { ˆ p i | ˆ p i
= p
i
+ k 1 ℓ
+ k 2 ℓ 2 + k 3 ℓ 3 , k 1 , k 2 , k 3 ∈ Z } .
Formally, the multi-edge graph construction method builds edges between nodes as follows.
Given a 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 built from j to i with the initial edge feature d ′ ji .
An example of the edge construction is shown in Fig.(b).
Intuitively, if there exist m positions of node j within the radius of the center node i , this method builds m edges from node j to node i .
By considering all possible positions of every node within a predefined radius in 3D space, the multi-edge graph construction method can in essence capture atom interactions across cell boundaries [40, 33, 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 a regular lattice in 3D space.
When encoding such periodic structures, unique challenges lie in periodic invariance and periodic pattern encoding.
In this section, we propose to formally define and analyze the 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) invariant such 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 to position matrix P and lattice matrix L together, or applying translations to P only.
Correspondingly, the output of the unit cell E(3) invariant function should remain the same.
[Equation]
[Table]
[Figure]
In addition to unit cell E(3) invariance, periodic invariance is also shown necessary for generating valid crystal representations.
Specifically, when the lattice matrix L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ R 3 × 3 is fixed for a crystal, we can still obtain different position matrices P ∈ R n × 3 and different unit cell structures by shifting the period boundaries.
As shown in
Fig.
2 (a) and (b), the formed unit cell structures are different for the same crystal by shifting period boundaries.
To this end, we further introduce periodic invariance, which shows that when the periodic boundaries are shifted or scaled up, the periodic invariant representation should remain 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
[Figure] [Figure]
[Table]
1 + `
[Table]
|
[Table]
[Figure]
from a given infinite crystal structure.
For an infinite crystal structure represented as ( ˆ A , ˆ P ) , Φ uses a corner point p and shape matrix L to form a unit cell represented as ( A , P ) .
In addition, we use α ∈ N 3+ to indicate the scaling up of a repeating unit cell 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 periodic invariant 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 crystal graphs for the same crystal.
In Fig.
2 (a) and (b), we show a graph construction method that breaks periodic invariance.
This method is employed by a transformer-based model, known as
Graphormer [54], which first uses radius to include all the atoms of interest in nearby cells and then builds a fully connected graph.
A detailed illustration is shown in Fig.
6 in Appendix.
A.1.
Based on the formal definition of periodic invariance, in this study, we aim to integrate such an important component in our Matformer.
By doing this, our model is able to construct a distinct crystal graph for a given crystal structure, resulting in a more informative and discriminative crystal learning scheme.
3.2 Periodic pattern encoding
[Figure]
Illustration that the multi-edge graph method only captures local geometric information but ignores periodic patterns for the infinite structure.
As introduced in Sec. 2, L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ R× 3 containing periodic patterns is another key component to describe crystal structures.
Essentially, periodic patterns show how the minimum repeatable structure ( A , P ) expands itself in infinite 3D space.
Without such periodic pattern encoding in L , crystal structures are treated as finite structures similar to molecules.
As shown in Fig.
3, the widely used multi-edge graph construction method [52, 4, 41, 40, 6] only captures local interactions among atoms but ignores the important periodic patterns.
However, such periodic repeating patterns need to be captured explicitly, as lattices of different sizes and orientations may correspond to different materials.
Hence, To better represent the infinite structures of 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 methods effectively 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 satisfy periodic invariance, and mathematical proofs can be found in Appendix.
A.2.
In Fig.
2, we show by example that treating atoms as single nodes breaks periodic invariance.
Instead, our methods follow the fashion introduced in Sec.
2 to treat node i as atom i and all its repeated duplicates.
We use the multi-edge graph construction [52] introduced in Sec. 2 as an alternative crystal graph construction method.
Note that although the multi-edge graph construction satisfies periodic invariance as in Fig.
4 (a), several existing works [52, 33] do not follow the settings exactly thus breaking periodic invariance.
Within the given radius shown in Fig.
4 (a), these studies form the neighborhood of node i by selecting t nearest neighbors ranked by geometric distances.
If there exist several different atoms with the same distance to node i , there is no deterministic way to select from them, as shown in Fig.
7 in Appendix.
A.1.
As a result, periodic invariance cannot be guaranteed.
Given the fully-connected fashion employed in Graphormer has achieved impressive performance on molecular learning, we further propose another graph construction method for Matformer, known as the fully connected graph construction.
This method uses a different strategy to determine neighbors for each center node.
Specifically, for node i and j , it builds edges for the entries corresponding to the 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 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 both satisfy the important periodic invariance.
Particularly, they possess great flexibility to be used in future studies 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 the important L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ R 3 × 3 into crystal graphs by adding self-connecting edges.
As mentioned in Sec. 3.2, periodic patterns describe sizes and orientations of lattices of a crystal structure, eventually determining the properties of this crystal.
A natural step for encoding such repeating periodic patterns is to consider the relative positions between an atom and its nearby repeated duplicates.
Formally, given an atom i with position p i and L = [ ℓ 1 , ℓ 2 , ℓ 3 ] T ∈ R 3 × 3 , we need to encode the atom’s three nearby duplicates with positions p i + ℓ 1 , p i + ℓ 2 , and p i + ℓ 3 .
It is 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 this end, 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 distance || ℓ 1 + ℓ 2 || 2 .
Hence, based on these three distances, we can determine lengths of ℓ 1 and ℓ 2 , and the relative orientation between them.
Extensively, 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, as shown in Fig.
4 (b).
By doing this, the length of each direction vector and the angle between any two direction vectors can all be determined.
As a result, the shape formed by lattice matrix L is then fixed.
Overall, we build the aforementioned six geometric distances as six self-connecting edges for node i .
By doing this, our model is capable of encoding periodic patterns in L completely, resulting in a more accurate crystal representation learning scheme.
Importantly, our approach also guarantees periodic invariance.
Overall, the graph construction for Matformer consists of two necessary stages, including invariant graph construction and periodic pattern encoding.
For the first stage, we rigorously prove that the multi-edge graph construction satisfies periodic invariance in Appendix.
A.2 and show that several previous works [52, 33] break periodic invariance using a different neighbor selection strategy.
Additionally, we propose a fully-connected crystal graph construction method satisfying periodic
[Table]
[Figure]
[Table]
Update
[Table]
latexit> f
⇤i <latexit
[Table]
Output
[Table]
i
[Table]
<
[Table] [Table]
[Figure]
invariance, and the method could be used in future studies for crystal representation learning.
For the second stage
[Figure]
we naturally encode periodic patterns into constructed graphs by adding self-connecting edges 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.
[Figure]
we propose our message passing scheme for
Matformer.
Formally, we denote a constructed crystal graph as G =
( A , E ) .
Here, each a i ∈ A is the d a -dimensional feature vector for atom i , as introduced in Sec. 2.
Particularly, e hij ∈
E is d e -dimensional feature vector for the h -th edge between nodes i and j .
We follow the regular attention mechanism that computes query, key, and value [30, 28].
Our proposed message passing scheme is 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 node i 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 edge embedding 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 match the dimension of k hij .
By doing this, when computing α hij , q i attends each of k i , k j , and e h ′ ij for integrating more information in attention.
Particularly, we omit the softmax to enhance the model’s capability 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]
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 message m 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 msg and 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 differences between Graphormer and the proposed Matformer lie in both graph construction and message passing scheme.
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 and encodes 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-wise attention and encoding geometric information into edge-wise messages, as described above.
The detailed architecture of our Matformer is shown in Fig.
Related work
Crystal property prediction .
Several existing methods [46, 18, 19, 14] model crystals as chemical formulas and employ sequence models to process them.
Other studies [52, 40, 33, 4, 41, 6] consider 3D 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 consider periodic invariance and periodic pattern encoding.
There are limited efforts to identify these two unique components.
As an early work, CGCNN [52] proposes to capture atom interactions across artificial cell boundaries by using multi-edge graphs described in Sec. 2.
The multi-edge graph satisfies periodic invariance as described in Sec. 3.1, but fails to consider the important periodic patterns, as described in Sec. 3.2.
The multi-edge graph construction method is widely used in the following studies [40, 33, 4, 41, 6, 1].
Based on the constructed crystal graphs, many GNN variants have been proposed for effective crystal representation learning [52, 40, 33, 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 [48] for periodic graph generation, which considers periodic graphs as finite graphs and breaks periodic invariance.
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 discriminative representations.
However, the use of angles introduces excessive time complexity.
Geometric GNNs and graph transformer .
Many efforts have been made to incorporate 3D geometric information in molecular learning.
Exemplary studies include SchNet [41], DimeNet [23, 22],
SphereNet [31], GemNet [24], etc.
However, these methods are designed for molecules without periodic patterns.
Recently, graph transformers [54 ? ] using geometric information, e.g.,
Graphormer [54], have shown great potential on real-world graph data.
However, Graphormer considers neither periodic invariance nor periodic pattern encoding.
Differences with our method.
To the best of our knowledge, periodic invariance and periodic pattern encoding described in Sec.
3 are rarely identified and explored in existing works for crystal property prediction.
CGCNN [52] breaks periodic invariance on some corner cases because it uses twelve nearest neighbors determined only by distances as described in Sec. 4.1.
In addition, previous methods including CGCNN [52], SchNet [41], MEGNET [4], CYATT [40], GATGNN [33], 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 [33] employs a very limited kind of attention mechanism that is not conditioned on query, as explained in GATv2 [2].
As a result, the model capacity is reduced compared with the self attention mechanism employed in Matformer.
In addition, the usage of softmax limits the capability of GATGNN of distinguishing nodes with different degrees, as mentioned in Sec. 4.2.
For Graphormer [54], although it achieved remarkable success on the Open
Catalyst Challenge [3], the employed graph construction method breaks periodic invariance when
[Table]
applied to crystals.
Compared with Graphormer, Matformer is specifically designed for crystals considering both periodic invariance and periodic patterns.
Experimental studies 6.1 Experimental setup
We conduct experiments on two material benchmark datasets, including The Materials Project [16] and JARVIS [8].
The detailed descriptions for The Materials Project and JARVIS datasets are shown in Appendix.
A.3 Baseline methods include CFID [7], CGCNN [52], SchNet [41], MEGNET [4],
GATGNN [33]
[Table]
and ALIGNN [6].
Unless otherwise specified, for all the baseline methods, we report the results taken from the referred papers or provided by original authors.
All Matformer models are trained using the Adam optimizer [21] with weight decay [32] and one cycle learning rate scheduler [43].
We only slightly adjust learning rates from 0.001 and training epochs from 500 for different 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 contains 69239 crystals, to evaluate Matformer.
We notice that previous works [52, 41, 4, 33, 6] compare with each other either using datasets of different sizes, or using datasets with the same size but splitting the datasets with different random seeds.
To make the comparison clear and fair, we retrain all corresponding models using exactly the same training, validation and test sets across all methods and report 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 [52, 41, 4, 33, 6].
It can be seen from Table.
1 that Matformer achieves the best performances on all tasks consistently by significant 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 4664 training samples, Matformer achieves the best performances, indicating Matformer’s adaptive ability to tasks of small training scales.
JARVIS dataset .
The quantitative results for Jarvis are shown in Table.
Matformer outperforms the baseline methods significantly on all of these five tasks.
Compared with ALIGNN, Matformer has stronger discriminative ability due to explicit encoding of periodic patterns.
Specifically, Matformer reduces Jarvis Ehull by 0.012, which is 15.8% of ALIGNN.
Furthermore, Matformer achieves the best performances for Bulk Moduli and Shear Moduli in the Mateirals Project with 4664 training samples, and Bandgap(MBJ) in JARVIS with 14537 training samples, indicating its adaptive 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 comparing with ALIGNN using JARVIS formation energy dataset.
The mean time of ten runs for training and inference using best model configurations of ALIGNN and Matformer are reported.
We also report the total number of parameters of each model.
In Table.
3, we show that Matformer is three times faster than ALIGNN
[Table] [Table]
in total training time and near three times faster in inference time, for the whole test set.
Matformer is also much lighter than ALIGNN in terms of model size.
Energy within Threshold .
Following OC20 [3]
[Table]
we use energy within threshold (EwT)
[Table]
which measures the percentage of estimated energies that are likely to be practically useful when the absolute error is within a certain threshold, to evaluate Matformer’s capability for periodic graph learning.
This metric is new
[Table]
but is well recognized by the community as it is useful in practice.
Due to significant performance gaps between ALIGNN and other baseline methods on formation energy and total energy for these two datasets in terms of mean absolute error (MAE), we only compare Matformer with ALIGNN.
Table.
4 shows that Matformer outperforms ALIGNN consistently for all three energy prediction 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.
Compared with JARVIS, the Materials Project has 15422 more traning samples.
As a result, the percentage of predicted energies obtained by Matformer within 0.01 increases by 14.69%, which is much better than 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 pattern encoding for crystal representation learning by ablation studies.
We also evaluate the building blocks particularly designed for our Matformer.
Specifically, we conduct experiments on JARVIS formation energy, and the test MAE is used as the quantitative evaluation metric.
We also provide ablation studies on the use of sigmoid and layernorm instead of softmax in Matformer layer in Appendix A.6.
Periodic invariant graph construction .
We demonstrate the importance of periodic invariant graph construction by comparing radius multi-edge graph, denoted as Radius, with the graph construction method proposed by 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,
[Table] [Table]
where n is the atom number in a cell.
We adjust the Matformer configurations to train these large graphs on a 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 with radius-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 pairwise distances 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 show that omitting the periodic pattern encoding results in a significant drop of test MAE from 0.0325 to 0.0337, revealing the importance of periodic patterns for crystal representation learning.
We show the training time
[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 will have 6 n nodes and 66 n edges, leading to high computational cost.
The detailed complexity analysis of adding angular information are provided in Appendix.
A.7.
Additionally, we provide the running 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 [23, 31, 24] to encode angles, denoted as Matformer + Angle RBF and Matformer + Angle SBF.
Table.
6 shows that introducing angular information will increase both the training time per epoch and in total by around 3 times, without much performance gain.
This may due to the periodic invariant graph construction and periodic patterns encoding in Matformer already capture sufficient information to identify crystal structures.
Conclusions and discussions
In this work, we first propose to formally define periodic invariance and periodic pattern encoding for 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 on common benchmark datasets show that our Matformer outperforms baseline methods consistently.
In addition, our results demonstrate the importance of periodic invariance and explicit periodic pattern encoding for crystal representation learning.
One potential direction beyond this work is to include angular information properly to satisfy both periodic invariance and to encode periodic patterns with relatively low time complexity, and this is one limitation of our work.
Besides, negative societal impacts of material discovery may apply to our work.
Acknowledgments and Disclosure of Funding
We thank Tian Xie for answering our questions on CGCNN.
This work was supported in part by
National Science Foundation grant IIS-2006861.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions 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] See Conclusion and limitation section.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
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 experimental 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 experiments 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] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Appendix A.1 Cases breaking periodic invariance
We show two graph construction methods that break periodic invariance in Fig.
6 and Fig.
One is the graph construction method employed by Graphormer in OC20 [3], shown in Fig.
It breaks 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]
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 pairwise distances, employed by CGCNN [52] and GATGNN [33], shown in Fig.
If several different atoms have the same geometric distances to the center atom, there is no deterministic way to determine which 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 given crystal 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 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 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 ∈ Z , i ∈ Z , 1 ≤ i ≤ n } , and ˆ A
= { ˆ a i | ˆ a i = a i , i ∈ Z , 1 ≤ i ≤ n } .
Here, set ˆ P contains all possible positions 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 the following proofs, we don’t consider the case that the provided L ′ =
α L for α ∈ N 3+ , which means the provided periodic patterns are not the minimum repeating patterns for a given crystal.
When L representing the minimum repeating patterns is given, the periodic invariance lies in the invariance to shifts of periodic boundaries.
We prove that the multi-edge graph construction and fully connected graph construction for crystals employed by our Matformer satisfy periodic invariance.
Multi-edge graph .
The Multi-edge graph satisfies the periodic invariance by forming edges between node 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 ∈ Z , d ij ≤ r } , where r is a prefixed threshold.
It can be seen that 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 } are considered, and the shifts of periodic boundaries will not influence the pairwise Euclidean distances.
Thus, this method satisfies periodic invariance.
Additionally, by using pairwise Euclidean distances, unit cell E(3) invariance is naturally satisfied.
We note that for a given node i , the radius can be computed by a deterministic function that takes all pairwise 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 always produce the same r for the same input { d ij | d ij = ||
p i − p j + k ′ 1 ℓ 1 + k ′ 2 ℓ 2 + k ′ 3 ℓ
|| 2 , k ′ 1 , k ′ 2 , k ′ 3 ∈ Z , 1 ≤ j ≤ n } .
One example is that we can use the 12-th smallest distance d 12 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 , 1 ≤ 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 node i and j , t smallest pairwise 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 } 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 } are considered, and the shifts of periodic boundaries will not influence the pairwise Euclidean distances.
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 directly follow ALIGNN [6] and use the same training, validation, and test set, including 60000, 5000, and 4239 crystals, respectively.
For tasks of Bulk Moduli and Shear Moduli, we follow GATGNN [33], the recent state-of-the-art method for these two tasks, and use the same training, validation, and test sets, including 4664, 393, and 393 crystals.
In Shear Moduli, one validation sample is removed because of the negative GPa value.
We either directly use the publicly available codes from the authors, or re-implement models based on their official codes and configurations to produce the results.
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 contains 44578, 5572, and 5572 crystals for tasks of formation energy, total energy, and bandgap(OPT).
The numbers are 44296, 5537, 5537 for Ehull, and 14537, 1817, 1817 for bandgap(MBJ).
The used metric is test MAE.
The results for CGCNN and CFID are taken from ALIGNN [6], other baseline results are obtained by retrained models.
Detailed configurations of these retrained models are provided in
Appendix.
A.5.
A.4 Matformer configurations
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 feature vector for the h -th edge between nodes i and j , and f ⋆i is the input feature vector of node
i for a given 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 computed query, key, and value vectors after these linear transformations for node i .
e h ′ ij is the intermediate output 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.
We use the 12-th smallest distance between the given atom and all nearby atoms as radius, and include all nearby atoms within the radius as the neighborhood for the given atom.
It satisfies periodic invariance as 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 these distances, if it is already in the radius of node i computed by the previous graph construction method, 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-dimensional embedding using CGCNN [52] atomic embedding.
We then use a linear transformation to map it to a 128-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 centers from 0.0 to 8.0.
It is then mapped to a 128-dimensional vector as the edge input e hij , by a nonlinear layer followed by a linear layer.
Matformer layer .
For message passing, we use five layers of Matformer Message Passing layer, with four attention heads.
The embedding sizes for q i , k i , v i for a single head are 128, mapped from f ∗ i using corresponding linear transformations LN Q , LN K , and LN V .
For a given layer, different heads use different LN Q , LN K , LN V , and LN E , but share other operations.
To obtain the final message m i , the features from four heads are concatenated and mapped to a 128-dimensional vector by a linear transformation.
Readout layer .
We use the mean pooling to aggregate features from all nodes in a graph and then use a nonlinear layer with hidden dimension 128 followed by a linear layer to obtain the scalar output for a crystal graph.
Training hyperparameters .
For all tasks in The Materials Project and JARVIS, we use the Adam optimizer [21] with weight decay [32] of 1e-5 and one cycle learning rate scheduler [43].
We use the batch size of 64.
We use mean square error as the objective function to train and mean absolute error as the evaluation metric to validate and test.
We only slightly adjust the learning rates and training epochs 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, we use one NVIDIA RTX A6000 48GB GPU to train our Matformer models.
[Table] [Table]
[Figure]
We also show the detailed Matformer architecture in Fig.
[Table]
A.5 Configurations of retrained models
We show configurations of retrained models for The Materials Project and JARVIS in this section.
The results from their original papers are shown in Table.
SchNet [41].
We use six layers of SchNet message passing layer following the original paper, with feature dimension of 64.
We train SchNet on these four tasks with learning rate of 5e-4 and batch size of 64 for 500 epochs.
The Adam optimizer is used with 1e-5 weight decay.
One cycle learning rate scheduler is also used.
We use the 12-th smallest distance between the given atom and all nearby atoms 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 the 32-th smallest distance between the given atom and all nearby atoms to serve as the radius for every atom to better the performance.
MEGNET [4].
Following the original paper, we use three layers of MEGNET message passing layer 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 forepochs following the configuration settings mentioned in the original paper.
The Adam optimizer is used with 1e-5 weight decay.
One cycle learning rate scheduler is also used.
We try different hyperparameters 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 size of 128.
And for bulk moduli and shear moduli and all tasks in JARVIS, we use the 12-th smallest distance between the given atom and all nearby atoms to serve as the radius for this given atom to better the performance, and the batch size of 64 is chosen because of better empirical results.
CGCNN [52].
For CGCNN, we directly use the publicly available code from Xie and Grossman [52], 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 for these four tasks using learning rate of 1e-2 for 100 epochs and 1e-3 for the next 900 epochs following the official code.
The radius cutoff of 8.0 is used for all atoms, and the nearest 12 neighbors are selected.
GATGNN [33].
For GATGNN, we directly use the publicly available code from Louis et al.
[33], 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 parameter of 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 all atoms, 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 tasks of bulk moduli and shear moduli, with learning rate of 1e-3 and batch size of 64.
In particular, we use ALIGNN with four gcn layers and four alignn layers.
Overall, we either directly use the publicly available codes from corresponding authors [52, 33, 6], or re-implement models based on their official codes and configurations [41, 4] to produce the results in our experiments.
A.6 Building block of Matformer
We provide ablation studies about sigmoid and layernorm instead of softmax in our Matformer message passing layer in this section.
[Table]
We evaluate the operation of sigmoid and normalization instead of softmax in our message passing, which is designed to capture the node degree information in multi-edge crystal graphs, described in Sec. 4.2.
Compared with using softmax with scalar attention coefficient, denoted as Softmax-scalar, or softmax with vector attention coefficient, denoted as Softmaxvector, using sigmoid and normalization leads to significant performance 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 edge is treated as a node in the line graph.
So we have 6 n nodes in the line graph.
Every edge in the original graph is connecting 2 nodes and every node has (12 − 1) other edges, resulting inneighboring edges for each edge in the original graph.
So we have 22 ∗ 6 n/ 2 = 66 n edges in the converted line graph.
Compared with original graph with | V | = n and | E | = 6 n , the converted line graph is super large with node number of 6 n and edge number of 66 n .