PhAST: Physics-Aware, Scalable, and Task-Specific GNNs for Accelerated Catalyst Design
Abstract
Mitigating the climate crisis requires a rapid transition towards lower carbon energy.
Catalyst materials play a crucial role in the electrochemical reactions involved in a great number of industrial processes key to this transition, such as renewable energy storage and electrofuel synthesis.
To reduce the amount of energy spent on such processes, we must quickly discover more efficient catalysts to drive the electrochemical reactions.
Machine learning (ML) holds the potential to efficiently model the properties of materials from large amounts of data, and thus to accelerate electrocatalyst design.
The Open Catalyst Project OC20 data set was constructed to that end.
However, most existing ML models trained on OC20 are still neither scalable nor accurate enough for practical applications.
Here, we propose several task-specific innovations, applicable to most architectures, which increase both computational efficiency and accuracy.
In particular, we propose improvements in (1) the graph creation step, (2) atom representations and (3) the energy prediction head.
We describe these contributions and evaluate them on several architectures, showing up to 5 × reduction in inference time without sacrificing accuracy.
Introduction
To mitigate climate change at a global scale it is imperative to reduce the carbon emissions of ubiquitous industrial processes like cement production or fertiliser synthesis, as well as to develop infrastructures for storing low-carbon energy at scale, in order to re-use it wherever and whenever needed.
Since such processes rely on electrochemical reactions, they require the design of more efficient electrocatalysts [28] to become more environmentally and economically viable.
However, discovering easy-to-exploit low-cost catalysts that drive electrochemical reactions at high rates remains an open challenge. In fact, today’s catalyst discovery mostly relies on expensive quantum mechanical simulations such as the Density Functional Theory (DFT) to approximate the behaviour of the materials involved in the targeted chemical reaction. Unfortunately, the high computational cost of these simulations limits the number of candidates that may be efficiently tested, and consequently stagnates further advances in the field.
Machine learning (ML) holds the potential to approximate these calculations while reducing the time needed to assess each candidate by several orders of magnitude [29].
This capability would transform the search for new catalysts from evaluating O (1 , 000) of manually handpicked candidates to powerful AI-guided design of catalysts spanning millions or even billions of candidates [29].
To enable to use of ML for catalyst discovery, the Open Catalyst Project released OC20 [5], a large data set of adsorbate-catalyst systems and their relaxed energy—a relevant metric to assess how good a catalyst is for a given chemical reaction—computed with DFT from the initial atomic structure.
Despite recent progress [9, 27], major challenges remain.
First, state-of-the-art models have not yet reached high enough performance for practical applications.
Second, they are still too computationally expensive to allow the millions of inferences required to explore the large space of potential catalysts.
Third, the graph neural networks (GNN) typically used were designed for general 3D molecular prediction tasks rather than specifically for catalyst discovery, whose complexity may benefit from task-specific architectures.
To address these challenges, we propose multiple model improvements to increase the accuracy and scalability of generic GNNs applied to catalyst discovery.
In particular, our contributions are (1) a graph construction that is tailored to the task at hand, (2) richer physics-based atom representations, and (3) an energy head that learns a weighted sum of per-atom predictions.
We provide a broad evaluation of these contributions on OC20 and a thorough ablation study. In sum, our proposed improvements increase the performance by 5-8% while dividing compute time by a factor of 3 to 5.
We believe our work provides valuable insights for future research as it leverages domain-specific knowledge to improve parts of the pipeline that were not investigated up to now.
The resulting performance and scalability gains open the door to a practical use of GNNs for new electrocatalyst design, the ultimate end goal of this line of research.
Background
The problem we address is the prediction of the relaxed energy y ∈ R of an adsorbate-catalyst system from its initial configuration in space ( X , Z ) , where X ∈ R N × 3 is the matrix of 3D atom positions and Z ∈ N N contains atom characteristic numbers.
This is commonly represented as a graph regression task, where each sample is represented as a 3D graph G with node set V of dimension
N and adjacency matrix A ∈ R N × N .
H ∈ R N × H represents atom embeddings and T ∈ { 0 , 1 , 2 } N corresponds to tag information (see 3.1).
ML models designed for this task generally adopt graph neural networks as an architecture, as it naturally suits 3D molecular prediction.
Such GNNs typically share a common pipeline for how they are applied, as depicted in Fig.
[Figure]
block learns a representation for each chemical element and the block applies a global pooling of each node’s representation to obtain the energy prediction. The key distinction between methods lies essentially in the Interaction block, where the message passing schemes vary.
In 3D molecular prediction tasks, it is desirable to endow ML models with relevant symmetry properties.
In particular, we want predictions to be invariant to translations, rotations and (often) reflections.
Many models enforce these physical priors within the architecture, making it explicitly invariant or equivariant to the desired transformations.
Formal definitions are included in A.1.
Many GNNs in prior work focus on enforcing equivariance, though it is not strictly required for relaxed energy prediction, which calls for invariance.
Equivariant GNNs [22, 2, 8, 3, 4] are expressive and generalize well, but are very expensive computationally as they are constrained by equivariant filters built on spherical harmonics and the Clebsch-Gordan tensor product [4]. Recent methods [18, 16, 21] model equivariant interactions in Cartesian space using both invariant and vector representations. They are faster but their architectures are still very complex and lack theoretical guarantees.
Alternatively, E(3)-invariant methods [17, 23, 19, 27, 1] do not use atom positions directly in their internal workings.
Instead, these methods extract and use quantities that remain invariant under rotations and reflections.
Dimenet++ [13, 12], for example, includes a directional message passing (MP) mechanism that incorporates bond angles in addition to atom relative distances.
However, distances and bond angles do not suffice to uniquely identify the graph 3D structure.
This is achieved by SphereNet [14] and GemNet [9], which additionally extract torsion information (between quadruplets of nodes).
On the downside, these methods are very computationally expensive as they require considering 3-hop neighbourhoods for each update step.
Importantly, all these MP methods aim at broad applicability and do not leverage the specific constraints of individual tasks.
Proposed Method
In this section, we describe PhAST, a Physics-Aware, Scalable, and Task-specific GNN framework for catalyst design. Notably, the architecture innovations in our proposed framework are applicable to most current GNNs used in catalyst discovery.
These include a new graph creation step, richer atom representations and an advanced energy head for graph-level prediction.
3.1 Graph creation
Although the graph construction step is critical in graph ML tasks, it has received little of no attention by previous work on the OC20 data set.
Most methods reuse the original proposal in [5].
In OC20, each graph’s atoms are tagged as part of the adsorbate (tag 2), the catalyst’s surface (tag 1), or its sub-surface volume (tag 0).
Tag 0 atoms were originally added in DFT simulations to represent more explicitly the repeating pattern of the slab.
They are, by definition, further away from the adsorbate and are fixed by construction, unlike tag 1 and tag 2 atoms, which can move during the relaxation.
As a result, we hypothesise they contain redundant information, making them of lesser importance to predict the final relaxed energy.
Besides, since they account for ∼ 65% of the nodes B.3, and since SOTA GNNs often depend on multiple hops to compute bond and torsion angles, we propose to remove these nodes from the graph.
This should greatly reduce inference time without impacting expressivity.
Additionally, we explore forming super nodes that aggregate tag 0 atoms to avoid a potential information loss caused by their total removal.
We briefly describe these changes below, with more details in A.2. remove-tag-0 removes all atoms with tag
( t i = 0 ) from the graph, adapting correspondingly all graph attributes ( X , A , Z , T , etc.). one-supernode-per-graph aggregates all tag 0 nodes from G into a supernode s with position x s = 1 |S| (cid:80) i ∈S x i , adjacency A is =
max ( A ij : j ∈ S ) and a new z s . one-supernode-per-atom-type replicates the above strategy but creates one super node per chemical element in the catalyst subsurface.
Its attributes are defined based on its components, as previously.
3.2 Atom Embeddings
In all previous GNNs, atom representations are learned from scratch based on atomic number
H = H Z .
We propose to leverage domain information to improve these representations.
First, we hypothesise that whether a given atom belongs to the adsorbate, the catalyst surface, or its subsurface is important information.
We therefore include tag information as an additional embedding H T .
Second, we know from previous studies that some atomic properties (e.g. atomic radius or density) are useful for catalyst discovery [20, 25].
We leverage them as an additional embedding vector
H F (see A.3 for the full list of properties).
Lastly, we include an embedding for the group and period information ( H P , H G ) since atoms belonging to the same group or period often share similar behaviours [26].
As a result, our proposed atom embedding H is a concatenation of all of the above.
3.3 Energy head
In the literature, there is often limited focus on the energy head, which is the part of the output block responsible for the energy computation from final atom representations h Li .
To the best of our knowledge, all GNNs compute the relaxed energy using global pooling y
=
(cid:80) i ∈V h
i , where node embeddings are reduced to a scalar h i by linear layers.
We identify two limitations in this procedure: First, all atoms are assigned the same importance, even though the properties of an atom are normally influenced by what element it is.
Second, the graph topology is neglected by simply summing all atom encodings regardless of their 3D positions.
To overcome these limitations, we explore alternative energy heads.
First, a weighted sum of node representations , which grants adaptive importance to each chemical element, expressed as ˆ y = (cid:80) i ∈V α ( h Li ) · h i or ˆ y
=
(cid:80) i ∈V α ( h 0 i ) · h i , where the learnable importance weights α ( · ) depend either on the embedding block initial encodings
h
i or final ones h Li .
Second, a hierarchical pooling approach endowed with the following energy head pipeline: h Li → [Pooling → GCN] ( × 2 ) → Global Pooling → MLP → ˆ y .
By applying a graph convolutional network (GCN) [11] on a coarsened graph, we propagate information differently, allowing us to capture hierarchical graph information.
We implement two pooling operators: (1) GRACLUS [6]: a deterministic way to regroup topologically close nodes together based on the mincut problem.
HOSCPOOL [7]: an end-to-end operator that learns a cluster assignment matrix using a loss function inspired by motif spectral clustering.
Contrary to GRACLUS , it leverages node features, captures higher-order connectivity patterns and is differentiable.
Evaluation
In this section, we evaluate the performance and scalability of our contributions for three well-known
GNNs on the OC20 dataset [5].
Dataset .
OC20 contains 1,281,040 DFT relaxations of randomly selected catalysts and adsorbates from a set of plausible candidates.
In this paper, we focus on the Initial Structure to Relaxed Energy (IS2RE) task [29], that is the direct prediction of the relaxed adsorption energy from the initial atomic structure.
It comes with a pre-defined train/val/test split, 450,000 training samples and hidden test labels.
Experiments are evaluated on the validation set, which has four splits of ∼ 25 K samples: In Domain (ID), Out of Domain Adsorbates (OOD-ads), Out of Domain catalysts (OOD-cat), and
Out of Domain Adsorbates and catalysts (OOD-both). We measure performance by the energy mean average error (MAE) on each validation split, and scalability by the inference time (seconds) over the ID validation set.
Baselines .
We study the enhancements brought by our components to three key GNN architectures for molecular predictions: SchNet [17], DimeNet++ [12] and ForceNet [10].
We describe their functioning in B.1.
We have selected these three baselines based on their popularity and easeof-implementation but note that PhAST improvements are applicable to all recent 3D molecular prediction GNNs (to the best of our knowledge).
We use the hyperparameters, training settings and model architectures provided in the original papers.
We compare every baseline with their PhAST counterpart, incorporating the best components of each category, that is graph creation (3.1), enriched atom embedding (3.2) and advanced energy-head (3.3), as determined by an ablation study, described below and in the Appendix B.2.
Ablation.
Sub-surface atoms appear to contain redundant information as remove-tag-0 does not cause performance drop and aggregating it into super nodes does not yield better results.
There are two potential explanations: (1) the data generation process of DFT simulations is not optimal and tagdoes contain redundant information (2) ML models do not manage to extract meaningful information from this repeated pattern, in which case our approach could be used by future work to demonstrate a better usage of this long range context info.
remove-tag-0 also dramatically decreases compute time.
Enriched atom embeddings improve performance and generalisation, especially when adding tag information (available in the data set).
Hierarchical pooling approaches are not very successful, either due to the difficulty of the task or the absence of hierarchical structures; but both energy-head weighted sums are beneficial.
For a finer understanding of individual contributions, refer to the full ablation study in B.2.
In conclusion, we obtain the following best components for the PhAST version of our models: remove-tag-0, full concatenation of atom embeddings and predicting the system energy as a learned weighted sum of per-atom predictions.
Results .
From Table 1, we conclude that our set of PhAST enhancements consistently improve both performance and inference time upon the original baselines.
More precisely, we improve average MAE over the four validation splits by 7.75%, 5.25% and 5.80% compared to SchNet, D++ and ForceNet, while reducing model inference time by factors of 2.85, 5.25, and 5.03, respectively.
Moreover, we observe an MAE improvement of 12.3% over SchNet and 9.0% over DimeNet++ on [Table] val OOD-both(compared to a 7.75% and 5.25% decrease in mean MAE).
This points to the fact that PhAST models generalise better than the original baselines.
From the ablation study conducted in B.2, this is due to the combination of our extensions, as they all contribute to significantly better performance on out-of-distribution adsorbate-catalyst systems (OOD-both).
Finally, note that inference time gains are doubled from SchNet, a 1-hop message passing (MP) approach, to DimeNet++, a 2-hop MP approach (from 2.85 to 5.25).
We can therefore expect even better computational scaling gains on newer models like GemNet[9], which builds on 3-hop MP methods.
Conclusion
In this work, we presented several enhancements targeted to catalyst discovery and applicable across existing GNN models.
We showed that (1) enriching atom representations with physics-based properties (2) weighting atoms’ importance when computing the system energy, and (3) tailoring the graph creation to the specific task at hand, all reduce inference time by a factor of 3– with better accuracy
Our results suggest that complex practical applications like catalyst discovery benefit from task specific methods rather than general 3D molecular prediction GNNs and that focusing on all aspects of the pipeline instead of only the message passing block is beneficial. Acknowledgements.
This research is supported in part by ANR (French National Research Agency) under the JCJC project GraphIA (ANR-20-CE23-0009-01) and was made possible thanks to Fragkiskos D. Malliaros, who also provided valuable feedback all along. Alexandre Duval acknowledges support from a Mitacs Globalink Research Award. Alex Hernandez-Garcia acknowledges the support of IVADO and the Canada First Research Excellence Fund.
David Rolnick acknowledges support from the Canada CIFAR AI Chairs Program.
The authors also acknowledge material support from NVIDIA in the form of computational resources, and are grateful for technical support from the
Mila IDT team in maintaining the Mila Compute Cluster.
A Method A.1 Invariance and equivariance to symmetries
Let ϕ : V → R and Φ : V → W be arbitrary functions where W,V are linear spaces.
Let G be a group describing a symmetry which we want to incorporate into ϕ , Φ (e.g. euclidean symmetries
E (3)) .
We use group representations ρ 1 : G → GL ( V ) and ρ 2 : G → GL ( W ) , where GL ( V ) is the space of invertible linear maps V → V to represent how the symmetries g ∈ G are applied to vectors X ∈ V, W .
ϕ is an G-invariant function if it satisfies ϕ ( ρ 1 ( g ) X ) = ϕ
( X ) , ∀ g ∈ G and X ∈ V . Φ is an G-equivariant function if it satisfies Φ( ρ 1 ( g ) X )
=
ρ 2 ( g )Φ( X ) , ∀ g ∈ G and X ∈ V .
In this paper, we focus on accelerated catalysis and thus on adslab relaxed adsorption energy prediction.
Like for most 3D molecular prediction tasks, we want GNNs to predict the same energy for two rotated, translated or reflected versions of the same system, since their energy is equal in real-life.
Hence, we target E (3) -invariant models, where E (3) is the Euclidean group in a 3D space (we have 3D atom positions), that is, the transformations of that 3D space that preserve the Euclidean distance between any two points (i.e. rotations, reflections, translations).
Note that we do desire reflection invariance because we rotate the whole adsorbate-catalyst system and not just the adsorbate, in which case chiral molecules may have a different behaviour and shall be considered distinctly.
A.2 Graph creation A.2.1 OC
Chanussot, Lowik, et al.
[5] create each OC20 sample by choosing a bulk material from the Materials
Project database 2 .
Then, they select a surface from the bulk using Miller indices (at random) and replicate it at depth of at least 7 Å and a width of at least 8 Å.
The final slab is defined by a unit cell that is periodic in all directions with a vacuum layer of at least 20 Å applied in the z direction.
Next, they pick a binding site on this surface to attach the adsorbate onto the catalyst.
The graph is now a set of atoms with their 3D positions.
Last but not least, edges are created between any two nodes within a cutoff distance c = 6 Å of each other (considering periodic boundary conditions).
A.2.2 PhAST graph creation process
Although well grounded, the assumptions of this graph creation process are rarely questioned.
We do, with the objective of making the graph sparser and more informative for subsequent GNNs.
We describe more formally the three proposals evoked in 3.1. remove-tag-0 .
We denote the set of tag 0 atoms by S
= { i ∈ V : t i = 0 } .
The graph we create has attributes X = X S where X S is the position of all atoms except those in S .
Similarly for Z = Z S and T = T S .
A is defined as usual based on cutoff distance (and pbc): A ij = 1 if || x i − x j || < c ,otherwise.
The remaining graph attributes (cell offsets, distances, etc.) are updated correspondingly. one-supernode-per-graph .
The position of the new super node is the mean of its components: x s = 1 |S| (cid:80) i ∈S x i (with S as defined above), and we associate it to a new characteristic number z s (corresponding to a new element in atomic table) and
adjacency A is =
max
( a ij : j ∈ S ) .
We now remove all tag-0 atoms using the above method, and finally add a tag-0 attribute t i = 0 to the supernode.
Note that we also remove self-loop for the supernode. one-supernode-per-atom-type .
This extension is similar to the previous one, except that we create one supernode for each chemical element in the sub-surface catalyst.
This complexify a bit the graph definition.
Let a be the number of distinct elements with tag 0 in the graph ( a = 1 , 2 , 3 by construction) and a k be their characteristic number.
Let S a k
=
{ i ∈ V| t i =
and z i = a k } be the set of atoms corresponding to each distinct tag-0 element a k .
Each supernode s k is defined with x s k =
[Equation]
and A s k ′ s k
= 1 ,
A s k s k
=
For both super-node methods, we encode the number of tag-0 nodes aggregated into each super node with Positional Encodings ([24]) to represent their "cardinal".
A.3 Atom properties for the Embedding block
In atom embeddings, we use the following properties from the mendeleev Python package ([15]):
1. atomic radius, 2.
atomic volume 3.
atomic density 4.
dipole polarizability 5.
electron affinity 6. electronegativity (allen) 7.
Van-Der-Walls radius 8.
metallic radius 9.
covalent radius 10.
ionization energy (first and second order).
B Experiments and ablation study B.1 Baseline description and hyperparameters
We target three well-known GNN baselines to study the impact of our contributions.
As evoked in Fig 1 of the paper, they all follow a similar pipeline, mainly differing in their interaction blocks, which we briefly detail below.
SchNet [17] is a simple message passing architecture that leverages relative distances to update atom representations via a continuous filter: h ( l +1)
i =
(cid:80) j h lj ⊙ W l ( x i − x j ) where W l ( x i − x j ) is a radial basis function to encode distance between atom pairs.
DimeNet++ [12] is an optimised version of DimeNet [12], which proposes a directional message passing.
In different terms, they compute and update edge representations instead of atoms) using interatomic distances e RBF (encoded via bessel functions) and bond angles a SBF (encoded via 2D spherical Fourier-Bessel basis):
[Equation]
ForceNet [10] is a scalable force-centric GNN that does not impose explicit physical constraint (energy conservation, rotational invariance), thus avoiding some memory intensive computations.
It still attempts to enforce invariance by efficient rotation-based data augmentation.
Model-wise, it adopts a node message passing approach that leverages node positions directly via a spherical harmonics basis.
Hyperparameters .
We use each method’s optimal set of parameters, provided in the config folder of the OCP repository for the IS2RE task, for the full dataset: https://github.com/
Open-Catalyst-Project/ocp/tree/main/configs/is2re/all .
Since ForceNet was not applied to IS2RE before, we adapted its S2EF configuration file to fit the IS2RE task.
The only change is the smaller number of epochs used for DimeNet++ (10 instead of 20) and SchNet (20 instead of 30), as these additional epochs only lead to a small performance gain for a large amount of additional compute time.
B.2 Ablation study
Notation .
For the embedding block, tag-embed defines atom embeddings H using atom tag information and characteristic number: H = H Z || H T , where || denotes concatenation.
Similarly, phys-embed defines H
=
H Z
|| H F .
l-phys-embed is a learnable alternative H
= H Z
|| MLP ( H F ) . pg refer to period and group embeddings: H =
H Z
|| H P
|| H G .
All is a concatenation of all five embeddings: H =
H Z
||
H F
|| H G
||
H P
|| H T .
For the graph creation step, we have defined these approaches in A.2.2 (note: sn stands for supernode).
For the energy head part, w-init ( w-final ) denote the weighted sum of initial (final) atom embeddings.
graclus and hoscpool refer to the two hierarchical pooling approaches.
See Tables 2, 3, 4.
The results are reported in a similar fashion as for Table 1.
The symbol (cid:134) indicates that a result is better than the baseline model. Bold font shows the best extension for each PhAST improvement category.
[Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table]
[Table]
[Table]
[Table]
hoscpool
[Table]
[Table]
[Table]
[Table]
[Table]
+/-
[Table]
[Table]
[Table]
401 +/
[Table] [Table] [Table] [Table]