Informing Geometric Deep Learning with Electronic Interactions to Accelerate Quantum Chemistry

TL;DR

OrbNet-Equi presents a QM-informed geometric deep learning approach that integrates tight-binding mean-field electronic structure with a 3D equivariant neural network to predict a wide range of quantum-chemical properties with high data efficiency and speed. The UNiTE architecture processes atomic-orbital features through diagonal reduction, block convolution, message passing, and point-wise interaction modules, all constructed to be equivariant under roto-translation, enabling consistent predictions across reference frames. Delta-learning with tight-binding featurization yields competitive accuracy to state-of-the-art DFT methods while delivering substantial speed gains, and the model demonstrates robust transfer to diverse chemical spaces (e.g., GMTKN55) and complex phenomena such as charge transfer and open-shell systems; it also accurately predicts electron densities and can generalize to unseen electronic states (zero-shot IP predictions on G21IP). These capabilities imply that a physics-guided ML-hybrid paradigm can broaden the scope and efficiency of quantum chemistry and materials discovery, enabling high-throughput screening and more affordable multiscale simulations. The work highlights data-efficient, symmetry-aware representations as a path toward transferable, scalable quantum-chemical modelling with practical implications for catalysis, battery design, and biomolecular research.

Abstract

Predicting electronic energies, densities, and related chemical properties can facilitate the discovery of novel catalysts, medicines, and battery materials. By developing a physics-inspired equivariant neural network, we introduce a method to learn molecular representations based on the electronic interactions among atomic orbitals. Our method, OrbNet-Equi, leverages efficient tight-binding simulations and learned mappings to recover high fidelity quantum chemical properties. OrbNet-Equi models a wide spectrum of target properties with an accuracy consistently better than standard machine learning methods and a speed orders of magnitude greater than density functional theory. Despite only using training samples collected from readily available small-molecule libraries, OrbNet-Equi outperforms traditional methods on comprehensive downstream benchmarks that encompass diverse main-group chemical processes. Our method also describes interactions in challenging charge-transfer complexes and open-shell systems. We anticipate that the strategy presented here will help to expand opportunities for studies in chemistry and materials science, where the acquisition of experimental or reference training data is costly.

Figures (8)

• Figure 1: QM-informed machine learning for modelling molecular properties. (a) Conventional ab initio quantum chemistry methods predict molecular properties based on electronic structure theory through computing molecular wavefunctions and interaction terms, with general applicability but at high computational cost. (b) Atomistic machine learning approaches use geometric descriptors such as interatomic distances, angles, and directions to bypass the procedure of solving the electronic structure problem, but often requires vast amounts of data to generalize toward new chemical species. (c) In our approach, features are extracted from a highly coarse-grained QM simulation to capture essential physical interactions. An equivariant neural network efficiently learns the mapping, yielding improved transferability at an evaluation speed that is competitive to Atomistic ML methods. (d) Characteristics of the atomic orbital features considered in OrbNet-Equi. Every pair of atoms $(A, B)$ is mapped to a block in the feature matrix, with the row dimension of the block matching the atomic orbitals of the source atom $A$ and the column dimension matching the atomic orbitals of the destination atom $B$. (e) OrbNet-Equi is equivariant with respect to isometric basis transformations on the atomic orbitals (Equations \ref{['eq:basis_rot']}-\ref{['eq:rotsym']}), yielding consistent predictions (illustrated as the dipole moment vector of a HSF molecule) at different viewpoints.
• Figure 2: Schematic illustration of the OrbNet-Equi method. The input atomic orbital features $\mathbf{T}[{\Psi_{0}}]$ are obtained from a low-fidelity QM simulation. A neural network termed UNiTE first initializes atom-wise representations through the diagonal reduction module, and then updates the representations through stacks of block convolution, message passing, and point-wise interaction modules. A programmed pooling layer reads out high-fidelity property predictions $\hat{\mathbf{y}}$ based on the final representations. Neural network architecture details are provided in Methods \ref{['sec:unite']}.
• Figure 3: Model performance on the QM9 dataset. (a-b) Test mean absolute error (MAE) of OrbNet-Equi is shown as functions of the number of training samples, along with previously reported results from task-specific ML methods (FCHL18Faber2018, FCHL19Christensen2019, SLATMhuang2017efficient, SOAPbartok2017machine, FCHL18*christensen2019operators, MuMLveit2020predicting) and deep-learning-based methods (SchNetschutt2017schnet, PhysNetunke2019physnet, OrbNet orbnet1) for targets (a) electronic energy $U_0$ and (b) molecular dipole moment vector $\vec{\mu}$ on the QM9 dataset. Results for OrbNet-Equi models trained with direct-learning and delta-learning are shown in dashed and solid lines, respectively. (c) Incorporating energy-weighted density matrices to improve data efficiency on learning frontier orbital properties. The HOMO, LUMO, and HOMO-LUMO gap energy test MAEs of OrbNet-Equi are shown as functions of the number of training samples. For models with the default feature set (red curves), the reduction in test MAE for delta-learning over direct-learning models gradually diminishes as the training data size grows. The LUMO and gap energy MAE curves exhibit a crossover around 32k-64k training samples, thereafter direct-learning models outperform delta-learning models. In contrast, when the energy-weighted density matrix features are supplied (blue curves), the test MAE curves between direct-learning and delta-learning models remain gapped when the training data size is varied. The black stars indicate the lowest test MAEs achieved by Atomistic ML methods (SphereNet spherenet) trained with 110k samples.
• Figure 4: Learning electron charge densities for organic and biological motif systems. (a) 2D heatmaps of the log-scale reference density $\rho(\vec{r})$ and the log-scale OrbNet-Equi density prediction error $\lvert\hat{\rho}(\vec{r})-\rho(\vec{r})\rvert$ (both in $a_0^{-3}$). The heatmaps are calculated by sampling real-space query points $\vec{r}\in\mathbb{R}^3$ for all molecules in the (red) BfDB-SSI test set and (blue) QM9 test set. The nearly-linear relationship for $\log_{10}(\rho(\vec{r}))<-4$ low-density regions reveals that OrbNet-Equi-predicted densities possess a physical long-range decay behavior. Distributions of $\log_{10}(\rho(\vec{r}))$ and $\log_{10}(\lvert\hat{\rho}(\vec{r})-\rho(\vec{r})\rvert)$ are plotted within the marginal charts. (b) The $L^1$ density errors $\varepsilon_{\rho}$ of OrbNet-Equi are plotted against the $\varepsilon_{\rho}$ of densities obtained through monomer density superposition (MDS), across the BfDB-SSI test set. Error bars mark the 99% confidence intervals of $\varepsilon_{\rho}$ for individual samples. The inset figure shows the average $\varepsilon_{\rho}$ for MDS, an Atomistic ML method fabrizio2019electron, and OrbNet-Equi predictions on the BfDB-SSI test set. OrbNet-Equi yields the lowest average prediction error and consistently produces accurate electron densities for cases where inter-molecular charge transfer is substantial. (c-d) Visualization of density deviation maps for (c) MDS and (d) OrbNet-Equi-predicted densities on the $\textrm{Glu}^{-}/\textrm{Lys}^{+}$ system (SSI-139GLU-144LYS-1), a challenging example from the BfDB-SSI test set. Red isosurfaces correspond to $\Delta\rho = -0.001\ a_0^{-3}$ and blue isosurfaces correspond to $\Delta\rho = +0.001\ a_0^{-3}$, where $\Delta\rho$ is the model density subtracted by the DFT reference density.
• Figure 5: OrbNet-Equi/SDC21 infers diverse downstream properties at an efficiency of semi-empirical tight-binding calculations. (a) Conformer energy ranking on the Hutchison dataset of drug-like molecules. The horizontal axis is labelled with acronyms indicating each method (O: OrbNet-Equi/SDC21 (this work); G: GFN-xTB; G2: GFN2-xTB; A: ANI-2x; B: B97-3c; $\omega$: $\omega$B97X-D3/def2-TZVP). The y-axis corresponds to the molecule-wise $R^2$ between predictions and the reference (DLPNO-CCSD(T)) conformer energies. Violin plots display the distribution of $R^2$ scores for each method over the (left) neutral, (middle) charged, and (right) all molecules from the Hutchison dataset. Medians and first/third quantiles are shown as black dots and vertical bars. (b) A torsion profiles example from the TorsionNet500 benchmark. All predicted torsion scans surfaces are aligned to the true global minima of the highest level of theory ($\omega$B97X-D3/def2-TZVP) results, with spline interpolations. (c) A uracil-uracil base pair example for non-covalent interactions. The dimer binding energy curves are shown as functions of the intermolecular axis ($r_e$) where $r_e=1.0$ corresponds to the distance of optimal binding energy. (d) Geometry optimization results on the (left) ROT34 and (right) MCONF datasets. Histograms and kernel density estimations of the symmetry-corrected RMSD scores (Methods \ref{['si_summary']}) with respect to the reference DFT geometries are shown for each test dataset. (e) Evidence of zero-shot model generalization on radical systems. OrbNet-Equi/SDC21 yields prediction errors drastically lower than semi-empirical QM methods for adiabatic ionization potential on the G21IP dataset, achieving accuracy comparable to DFT on 7 out of 21 test cases.

Theorems & Definitions (10)

• Definition S1: $N$-body tensor
• Definition S2
• Corollary S1
• proof
• Theorem S1: Theorem 2.1 and Lemma 2.2 of klink1996multiplicity
• Corollary S2
• proof
• Corollary S3
• proof
• Definition S3