Efficient Prediction of SO(3)-Equivariant Hamiltonian Matrices via SO(2) Local Frames

TL;DR

Addresses efficient, symmetry-aware Hamiltonian matrix prediction by introducing SO(2) local frames to achieve global SO(3) equivariance without the costly $O(L_{\max}^6)$ SO(3) tensor products, reducing to $O(L_{\max}^3)$ with SO(2) operations. Proposes QHNetV2 with SO(2) Linear, Gate, LayerNorm, and Tensor Product in local frames, plus a minimal frame averaging mapping between $SO(3)$ and $SO(2)$ irreps. Demonstrates state-of-the-art MAE on Hamiltonian blocks and eigenvalues on QH9 and MD17, with substantial speedups (e.g., ~4.34x) over prior SO(3)-TP approaches. Shows strong generalization to diverse molecular structures and trajectories, suggesting a scalable path for symmetry-aware electronic-structure learning.

Abstract

We consider the task of predicting Hamiltonian matrices to accelerate electronic structure calculations, which plays an important role in physics, chemistry, and materials science. Motivated by the inherent relationship between the off-diagonal blocks of the Hamiltonian matrix and the SO(2) local frame, we propose a novel and efficient network, called QHNetV2, that achieves global SO(3) equivariance without the costly SO(3) Clebsch-Gordan tensor products. This is achieved by introducing a set of new efficient and powerful SO(2)-equivariant operations and performing all off-diagonal feature updates and message passing within SO(2) local frames, thereby eliminating the need of SO(3) tensor products. Moreover, a continuous SO(2) tensor product is performed within the SO(2) local frame at each node to fuse node features, mimicking the symmetric contraction operation. Extensive experiments on the large QH9 and MD17 datasets demonstrate that our model achieves superior performance across a wide range of molecular structures and trajectories, highlighting its strong generalization capability. The proposed SO(2) operations on SO(2) local frames offer a promising direction for scalable and symmetry-aware learning of electronic structures. Our code will be released as part of the AIRS library https://github.com/divelab/AIRS.

Figures (2)

• Figure 1: SO(2) equivariant operations. (a) SO(2) Linear. For SO(2) irreps with order $m>0$, this operation uses weight matrices $\mathbf{w}_1^{(m)}, \mathbf{w}_2^{(m)} \in \mathbb{R}^{C \times C}$ where $C$ is the number of channels for input irreps. (b) SO(2) Gate. For the $m=0$ features, a multi-layer perceptrons (MLP) is used to update them. Simultaneously, for each irrep with order $m>0$, the MLP outputs a gate value passed through a sigmoid function, which modulates the corresponding SO(2)-equivariant features. (c) SO(2) LayerNorm (LN). For the $m=0$ features, a standard LN is applied. For $m>0$ features, LN is applied on the norm of SO(2) irreps according to Eq. \ref{['eq:SO(2)_LN']}. (d) SO(2) Tensor Product (TP). The SO(2) tensor product fuses features by combining irreps under the constraints $m_3 = m_1 +m_2$ (shown as solid lines) or $m_3 = | m_1 - m_2|$ (shown as dashed lines). The color of the path corresponds to its originating SO(2) irreps. A more general case containing $v-1$ TPs for $v$ set of SO(2)-equivariant features is shown in Eq. \ref{['eq:SO(2)_TP_v_interactions']}. Each valid combination defines a path, ensuring the resulting features remain SO(2)-equivariant.
• Figure 2: The overall architecture of the proposed QHNetV2. In this figure, $\times$ denotes element-wise multiplication, $\langle\cdot, \cdot \rangle$ denotes inner product. Gray color denotes scaler values, red color denotes SO(2) irreps, and blue color denotes SO(3) irreps.