Equivariant Electronic Hamiltonian Prediction with Many-Body Message Passing

Abstract

Machine learning surrogate models of Kohn-Sham Density Functional Theory Hamiltonians provide a powerful tool for accelerating the prediction of electronic properties of materials, such as electronic band structures and density of states. For large-scale applications, an ideal model would exhibit high generalization ability and computational efficiency. Here, we introduce the MACE-H graph neural network, which combines high body-order message passing with a node-order expansion to efficiently obtain all relevant $O(3)$ irreducible representations. The model achieves high accuracy and computational efficiency and captures the full local chemical environment features of, currently, up to $f$ orbital matrix interaction blocks. We demonstrate the model's accuracy and transferability on several open materials benchmark datasets of two-dimensional materials and a new dataset for bulk gold, achieving sub-meV prediction errors on matrix elements and high accuracy on eigenvalues across all systems. We further analyze the interplay of high-body-order message passing and locality that makes this model a good candidate for high-throughput material screening.

Figures (8)

• Figure 1: Illustration of the Hamiltonian Matrix Prediction. a. The composition of Hamiltonian matrices for a twisted Bi2Te3 bilayer with spin-orbit coupling, the matrix can be decomposed into atom-pair-wise blocks, which consist of spin-polarized orbital-pair subblocks, b. The comparison of Hamiltonian matrix calculations via neural network and self-consistent field density functional theory. The training of the neural network is supervised by the labels generated from the density functional theory.
• Figure 2: Graph Neural Network Architecture of MACE-H.a. The overall workflow of module blocks. The node-wise message passing operations (as denoted in the yellow rectangle) consist of $L$ message neural network layers. $\mathrm{Encode(}Z_i\mathrm{)}$ and $\mathrm{Encode(}Z_i, Z_j\mathrm{)}$ are one-hot-encoding of the element $Z$ of atom $i$ and the element pair $Z_i-Z_j$ between atom $i$ and $j$, respectively. The initial node feature in MACE block and edge feature in the edge-update block in the first layer are the element embedding $\mathrm{Embed}\left(Z_{i}\right)$ and distance embedding $\mathrm{Embed}\left(\left|\bm{r}_{ij}\right|\right)$. $e_{\mathrm{N}}\left(\left|\bm{r}_{ij}\right|\right)$ and $e_{\mathrm{E}}\left(\left|\bm{r}_{ij}\right|\right)$ stand for node-update and edge-update radial basis sets, respectively. b. The MACE layer is used to aggregate the atom-level chemical environment feature representation, which encodes atom embedding, edge length, and orientation information to higher-body features. c. The node degree expansion (NDE) layer elevates the node feature degree in order to be compatible with the edge-wise irreducible representations (irreps) corresponding to the Hamiltonian subblocks. d. The edge update block converts the node-wise features and geometry information into the edge-wise features corresponding to the matrix block. This is identical to the edge update block in DeepH-E3 DeepH-E3. All abbreviations are defined in the main text.
• Figure 3: Demonstration of MACE-H for Materials Electronic Structure Prediction. The predicted band structure along the high-symmetry k-path and the density of states in comparison with DFT results for a two-dimensional bilayer of Bi2Te3 (a) and face-centered-cubic bulk Au (b).
• Figure 4: Model Performance Assessment of MACE-H Compared to DeepH-E3.a. The test set mean absolute error of matrix elements for different datasets of 2D monolayers, shifted 2D bilayers, and bulk systems. The 2D material datasets are generated by the OpenMX package, the Au all-electron dataset is generated with FHI-aims. b. The eigenvalue error and c. electronic entropy error for Au datasets. The eigenvalue error and electronic entropy error are defined in Section \ref{['sec:metrics']}. The hyperparameters are listed in Supplementary Table 3. The correlation order is $\nu=3$.
• Figure 5: Data- and Computational Efficiency of MACE-H.a. The matrix element MAE as a function of training data size of shifted Bi2Te3 for MACE-H and DeepH-E3. Here, DeepH-E3 uses spherical harmonics with $l$ up to 4, MACE-H uses the same settings in the edge-update block and node-wise blocks with spherical harmonics with $l$ up to 4, and hidden states with azimuthal number up to 2. The MACE-H with notation "high" uses spherical harmonics $l$ up to 5 and hidden state azimuthal numbers up to 5. The correlation order $\nu=3$ for both MACE-H models. b. The onsite and offsite-resolved matrix element MAE as a function of Bi2Te3 training data size, which consists of a total of 256 configurations with 90 atoms in each. Furthermore, shown are the single batch inference time of MACE-H on a single GPU (see c) and 64 CPU cores (see d) for graphene bilayers without spin-orbital coupling and Bismuthene and Bi2Te3 bilayers with spin-orbital coupling.