PhAST: Physics-Aware, Scalable, and Task-specific GNNs for Accelerated Catalyst Design

TL;DR

PhAST tackles the bottleneck of applying graph neural networks to catalyst design by introducing task-specific, physics-informed enhancements that improve both accuracy and computational efficiency. It rethinks graph construction, atom embeddings, and output heads (energy and forces) to align with the IS2RE and S2EF tasks on OC20, achieving MAE reductions of up to $42\%$ and inference-time reductions of $3$–$8\times$, with CPU training speedups up to $40\times$. The method demonstrates strong, architecture-agnostic gains across multiple baselines (SchNet, DimeNet++, ForceNet, GemNet, GemNet-OC) and shows practical impact by enabling CPU-based training and deploying scalable catalyst discovery workflows. By combining physics-aware representations with efficient graph-rewiring and energy-conserving force handling, PhAST provides a scalable pathway toward rapid, data-driven electrocatalyst design and broader applicability to related molecular modeling tasks.

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 numerous industrial processes key to this transition, such as renewable energy storage and electrofuel synthesis. To reduce the energy spent on such activities, we must quickly discover more efficient catalysts to drive electrochemical reactions. Machine learning (ML) holds the potential to efficiently model materials properties from large amounts of data, accelerating electrocatalyst design. The Open Catalyst Project OC20 dataset was constructed to that end. However, ML models trained on OC20 are still neither scalable nor accurate enough for practical applications. In this paper, we propose task-specific innovations applicable to most architectures, enhancing both computational efficiency and accuracy. This includes improvements in (1) the graph creation step, (2) atom representations, (3) the energy prediction head, and (4) the force prediction head. We describe these contributions, referred to as PhAST, and evaluate them thoroughly on multiple architectures. Overall, PhAST improves energy MAE by 4 to 42$\%$ while dividing compute time by 3 to 8$\times$ depending on the targeted task/model. PhAST also enables CPU training, leading to 40$\times$ speedups in highly parallelized settings. Python package: \url{https://phast.readthedocs.io}.

Figures (4)

• Figure 1: Common GNN inference pipeline for 3D material modeling. The graph creation step remains unchanged across all methods: it creates ${\pmb A}$ using cutoff distances and periodic boundary conditions. The Embedding and Output blocks slightly differ across models but the underlying idea is the same. The Embedding block learns a representation for each chemical element and the Output block applies a global pooling of each node's representation to obtain the energy prediction. The key distinction between methods typically lies in the Interaction block, where the message passing schemes vary.
• Figure 2: A comparison of the improvements brought by PhAST to the model's MAE (left) and inference times (right) on OC20 IS2RE. The PhAST components are selected in \ref{['sec:ablation']} and summarised in \ref{['subsec:final-model']}. MAE improvements (left) are computed as in Equation \ref{['eq:mae_improvement']} and values to the right of the dashed line at 0.0 denote an improvement of PhAST with respect to the baseline. The results are averaged over 3 runs, with bootstrapped confidence intervals represented by small horizontal bars. PhAST leads to a significant MAE improvement for each validation split, up to $13~\%$, in addition to decreasing inference time (s) by several factors (e.g. 3.4$\times$ or 5.4$\times$). Absolute numerical values are provided in \ref{['tab:best-perf-is2re']}.
• Figure 3: Ablation study results on IS2RE of our PhAST contributions about (1) the graph creation step, (2) atom embeddings and (3) the energy prediction head, all detailed in \ref{['sec:method']}, for SchNet, DimeNet++ and ForceNet. All changes lead to significant improvements in model performance compared to the dash line denoting the baseline (Left). All contributions have little impact on inference time except from the graph rewiring steps which divides it by several factors (Right). The best technique from each family selected for PhAST are highlighted with a darker edge around the shaded background.
• Figure 4: CPU-based training of MegNet chen2019graph using 4th Gen Intel Xeon Scalable Processors known as Sapphire Rapids (SPR). The top of the x-axis specifies the number of CPU nodes with each node including 2 SPR CPUs, meaning that at the largest degree of parallization we run on 32 SPR with 64 CPUs. The bottom x-axis outlines the number of MPI ranks, which specifies the number of parallel MPI processes occuring at a given time. Using 32 SPR nodes, we can scale up to 512 MPI ranks which provide significant speedup in model trianing.