Towards Harmonization of SO(3)-Equivariance and Expressiveness: a Hybrid Deep Learning Framework for Electronic-Structure Hamiltonian Prediction (2401.00744v11)
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.
- E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1):2453, 2022.
- Unconventional superconductivity in magic-angle graphene superlattices. Nature, 556(7699):43–50, 2018.
- Steerable cnns. In ICLR, 2017.
- Exploiting cyclic symmetry in convolutional neural networks. In ICML Workshop, volume 48, pp. 1889–1898, 2016.
- Ferreira, F. Multiscale modelling of twisted bilayers of 2d materials. Nature Reviews Physics, 4(10):632–632, 2022.
- Se(3)-transformers: 3d roto-translation equivariant attention networks. In NeurIPS, 2020.
- e3nn: Euclidean neural networks. CoRR, abs/2207.09453, 2022.
- General framework for e (3)-equivariant neural network representation of density functional theory hamiltonian. Nature Communications, 14(1):2848, 2023.
- Neural network representation of electronic structure from ab initio molecular dynamics. Science Bulletin, 67(1):29–37, 2022.
- Inhomogeneous electron gas. Physical Review, 136(3B):B864, 1964.
- Spatial transformer networks. In NeurIPS, pp. 2017–2025, 2015.
- Self-consistent equations including exchange and correlation effects. Physical Review, 140(4A):A1133, 1965.
- Covariant compositional networks for learning graphs. In ICLR Workshop, 2018.
- Deep-learning density functional theory hamiltonian for efficient ab initio electronic-structure calculation. Nature Computational Science, 2(6):367–377, 2022.
- Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1), 2023.
- Tunable strongly coupled superconductivity in magic-angle twisted trilayer graphene. Nature, 590(7845):249–255, 2021.
- Reducing SO(3) convolutions to SO(2) for efficient equivariant gnns. In ICML, pp. 27420–27438, 2023.
- Equivariance through parameter-sharing. In ICML, pp. 2892–2901, 2017.
- E(n) equivariant graph neural networks. In ICML, volume 139, pp. 9323–9332, 2021.
- Schrödinger, E. Quantisierung als eigenwertproblem. Annalen der Physik, 384(4):361–376, 1926.
- Unifying machine learning and quantum chemistry with a deep neural network for molecular wavefunctions. Nature Communications, 10(1):5024, 2019.
- Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219, 2018.
- Se (3)-equivariant prediction of molecular wavefunctions and electronic densities. volume 34, pp. 14434–14447, 2021.
- Attention is all you need. In NeurIPS, pp. 5998–6008, 2017.
- Deepmd-kit: A deep learning package for many-body potential energy representation and molecular dynamics. Computer Physics Communications, 228:178–184, 2018.
- Universal machine learning for the response of atomistic systems to external fields. Nature Communications, 14, 2023.
- Embedded atom neural network potentials: Efficient and accurate machine learning with a physically inspired representation. The journal of physical chemistry letters, 10(17):4962–4967, 2019.
- Spherical channels for modeling atomic interactions. In NeurIPS, 2022.