Towards Harmonization of SO(3)-Equivariance and Expressiveness: a Hybrid Deep Learning Framework for Electronic-Structure Hamiltonian Prediction

TL;DR

The paper tackles predicting electronic-structure Hamiltonians while enforcing $SO(3)$-equivariance and preserving nonlinear expressiveness. It introduces HarmoSE, a two-stage framework that couples a group-theory–based equivariant module with a highly expressive non-linear 3D graph Transformer to learn a residual correction to baseline Hamiltonians. Key contributions include a cascaded regression scheme, integration of covariant features, and demonstration of state-of-the-art results across six benchmarks, including twisted bilayer structures. This approach enables accurate, rotation-robust Hamiltonian predictions with strong generalization for 3D atomic systems, facilitating scalable quantum-material modeling and design.

Abstract

Deep learning for predicting the electronic-structure Hamiltonian of quantum systems necessitates satisfying the covariance laws, among which achieving SO(3)-equivariance without sacrificing the non-linear expressive capability of networks remains unsolved. To navigate the harmonization between equivariance and expressiveness, we propose a deep learning method synergizing two distinct categories of neural mechanisms as a two-stage encoding and regression framework. The first stage corresponds to group theory-based neural mechanisms with inherent SO(3)-equivariant properties prior to the parameter learning process, while the second stage is characterized by a non-linear 3D graph Transformer network we propose, featuring high capability on non-linear expressiveness. The novel combination lies in the point that, the first stage predicts baseline Hamiltonians with abundant SO(3)-equivariant features extracted, assisting the second stage in empirical learning of equivariance; and in turn, the second stage refines the first stage's output as a fine-grained prediction of Hamiltonians using powerful non-linear neural mappings, compensating for the intrinsic weakness on non-linear expressiveness capability of mechanisms in the first stage. Our method enables precise, generalizable predictions while capturing SO(3)-equivariance under rotational transformations, and achieves state-of-the-art performance in Hamiltonian prediction on six benchmark databases.

Figures (4)

• Figure 1: Left part: overview of the two-stage encoding and regression framework for Hamiltonian prediction. Right part: The internal architecture of the non-linear 3D graph Transformer network.
• Figure 2: Visualization of challenging testing samples, which exhibit structural deformations caused by thermal motions and inter-layer twists, calling for strong capabilities on expressiveness and SO(3)-equivariance of a regression model.
• Figure 3: The MAE values (denoted as $MAE^H_{block}$) on various basic blocks of the Hamiltonian matrix in direct product state on the experimental monolayer structures. The MAE values ($MAE^H_{cha\_b}$) on the Hamiltonian block where the baseline model DeepHE3 performs the worst are highlighted for comparison.
• Figure 4: The MAE values (denoted as $MAE^H_{block}$) on various basic blocks of the Hamiltonian matrix in direct product state on the non-twisted (marked with superscripts $nt$) and twisted (marked with superscripts $t$) bilayer structures. The MAE values ($MAE^H_{cha\_b}$) on the Hamiltonian block where the baseline model DeepHE3 performs the worst are highlighted for comparison.