Complete and Efficient Graph Transformers for Crystal Material Property Prediction

TL;DR

A novel approach that utilizes the periodic patterns of unit cells to establish the lattice-based representation for each atom, enabling efficient and expressive graph representations of crystals.

Abstract

Crystal structures are characterized by atomic bases within a primitive unit cell that repeats along a regular lattice throughout 3D space. The periodic and infinite nature of crystals poses unique challenges for geometric graph representation learning. Specifically, constructing graphs that effectively capture the complete geometric information of crystals and handle chiral crystals remains an unsolved and challenging problem. In this paper, we introduce a novel approach that utilizes the periodic patterns of unit cells to establish the lattice-based representation for each atom, enabling efficient and expressive graph representations of crystals. Furthermore, we propose ComFormer, a SE(3) transformer designed specifically for crystalline materials. ComFormer includes two variants; namely, iComFormer that employs invariant geometric descriptors of Euclidean distances and angles, and eComFormer that utilizes equivariant vector representations. Experimental results demonstrate the state-of-the-art predictive accuracy of ComFormer variants on various tasks across three widely-used crystal benchmarks. Our code is publicly available as part of the AIRS library (https://github.com/divelab/AIRS).

Figures (11)

• Figure 1: Illustrations of crystal passive symmetries. We show examples of a real crystal structure in 3D and corresponding simpler demonstrations in 2D. For the 2D demonstrations, we use circles with different colors to represent different kinds of atoms, and use red dotted lines to represent periodic unit cells. We show two cases of periodic transformations, including shifting the periodic unit cells and changing periodic unit cells without changing the volume. Translation, rotation, and periodic transformations will not change the crystal structure and are passive symmetries, while reflection transformation will map the crystal structure to its chiral image if the reflection symmetry is absent.
• Figure 2: Illustration of the proposed ComFormer pipeline. In the left figure, we show different unit cell structures for the same crystal due to passive crystal symmetries, and all of them will map to the same invariant or equivariant crystal graph shown in the middle. In the middle, we demonstrate the information included in our proposed SE(3) invariant and equivariant crystal graphs. Specifically, we include node feature $\boldsymbol{a}_i$ for every node $i$, and for every neighbor $j$ of node $i$, we include edge length $||\mathbf{e_{ji}}||_2$, and three angles $\theta_{ji,ii_*} = \{ \theta_{ji,ii_1}, \theta_{ji,ii_2}, \theta_{ji,ii_3} \}$ in invariant crystal graphs, and edge vector $\mathbf{e_{ji}}$ in equivariant crystal graphs. The proposed iComFormer and eComFormer are shown on the right, with building blocks marked by different colors.
• Figure 3: Demonstration of the multi-edge crystal graph. We use blue arrows to represent periodic patterns $\ell_1$ and $\ell_2$ for the shown crystal structure in 2D case. We use circles with different colors to represent different atoms, and light-green lines to represent the periodic boundaries. We use $\mathcal{N}_i$ to represent the neighborhood of node $i$ within the radius, and use red and green arrows to indicate the captured atomic interactions by the multi-edge crystal graph. (a) A crystal structure with periodic patterns $\ell_1$ and $\ell_2$ in 2D case for clarity. (b) The captured atomic interactions by the multi-edge crystal graph, between the red nodes $j$ and center node $i$. (c) The corresponding multi-edge crystal graph. All periodic duplicates of $j$ are mapped to a single node $j$ in the multi-edge crystal graph.
• Figure 4: Demonstration that multi-edge crystal graphs fail short to capture periodic patterns. We use blue arrows to represent periodic patterns $\ell_1$ and $\ell_2$ for these two different crystal structures. Due to the missing geometric information of periodic patterns, multi-edge crystal graphs map these two different crystals to the same crystal graph $\mathcal{G}$.
• Figure 5: Demonstration that multi-edge crystal graphs fail short to distinguish crystalline materials with different unit cell structures. We use $i, j, k$ to represent three atoms in the unit cell structures of these two crystals, $d_1, d_2$ to represent the encoded bond lengths in the graph, and $\theta_{jik}$ to represent the encoded bond angle between edge $ji$ and $ik$. These two crystals are different but have the identical multi-edge crystal graphs with identical Euclidean distances and bond angles encoded. The number of blue atoms in the middle part is three, and can increase to any arbitrary number to obtain the same conclusion.

Theorems & Definitions (9)

• Definition 1: Geometrically Complete Crystal Graph
• Definition 2: Unit Cell SE(3) Invariance
• Definition 3: Unit Cell SO(3) Equivariance
• Definition 4: Periodic Invariance
• Proposition 1
• proof
• Definition 5: Isometric crystal structures
• proof
• proof