PaperlayerView exact revision on arXiv

Source-linked AI summary

Informing Geometric Deep Learning with Electronic Interactions to Accelerate Quantum Chemistry

Zhuoran Qiao, Anders S. Christensen, Matthew Welborn, Frederick R. Manby, Anima Anandkumar, Thomas F. Miller III

arXiv:2105.14655v4cs.LGphysics.chem-ph

TL;DR

Predicting molecular properties in unfamiliar chemical and electronic settings remains a challenge for molecular modeling. OrbNet-Equi learns from tight-binding electronic structure in an equivariant atomic-orbital representation, and reports strong benchmark performance, including transfer to unseen electronic states.

  • Problem

    Predicting properties for molecules with chemical compositions or electronic states not represented in training remains a challenge for molecular modeling.

    Source: Transferability on downstream tasks · p. 5

  • Method

    OrbNet-Equi uses tight-binding mean-field electronic structure in an atomic-orbital representation with an equivariant neural network to predict molecular properties.

    Source: The OrbNet-Equi methodology · p. 2

  • Results

    OrbNet-Equi/SDC21 has the lowest median WTMAD across GMTKN55 tasks among methods executable at semi-empirical QM cost and predicts ionization potentials more accurately than semi-empirical QM methods despite rare open-shell training examples.

    Source: Transferability on downstream tasks · p. 5, Discussion · p. 6

  • Takeaways & Limitations

    The results support OrbNet-Equi as a possible alternative to conventional ab initio simulations for general-purpose downstream applications at a computational cost on par with semi-empirical methods.

    Source: Discussion · p. 6

    The explicit neural-network module expressions assume one-electron operators and a single feature matrix; extensions are discussed separately.

    Source: §1.1 · p. 8

Abstract

from arXiv · show

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.

Results

OrbNet-Equi learns molecular properties from tight-binding electronic interactions in an atomic-orbital representation, with equivariance to basis transformations. It shows strong performance across molecular-property and downstream benchmarks, including density prediction and transfer to unfamiliar chemical and electronic settings.

The OrbNet-Equi methodology

  • OrbNet-Equi uses mean-field tight-binding electronic structure, represented by atomic-orbital Fock, density, core-Hamiltonian, and overlap matrices, as input to an equivariant neural network.The network propagates information through off-diagonal matrix blocks and uses pooling to impose constraints such as size consistency.Source: The OrbNet-Equi methodology · pp. 2–3

Performance on benchmark datasets

  • On QM9, OrbNet-Equi matches FCHL18/GPR for total-energy test MAE, improves delta-learning MAE by about 45% over OrbNet, and averages about two-fold lower test MAE than atomistic deep-learning methods across 12 targets.For dipole moments, it exhibits steep learning-curve slopes under both direct- and delta-learning.Source: Performance on benchmark datasets · pp. 3–4
  • Adding energy-weighted hole- and particle-excitation density matrices reduces frontier-orbital-energy test MAE and avoids the learning-curve rank reversal seen with default ground-state features.The authors associate the reversal with incomplete spatially degenerate descriptors and report that breaking energy-space degeneracy restores the gap between delta- and direct-learning curves.Source: Performance on benchmark datasets · pp. 3–4

Accurate modelling for electron densities

  • OrbNet-Equi achieves average density error of 0.191±0.003% on BfDB-SSI with 2,000 training samples, versus 0.29% for SA-GPR, and 0.206±0.001% on QM9 with 123,835 samples, versus 0.28%–0.36% for DeepDFT.It also maintains low errors across real space and robust asymptotic decay in low-density regions.Source: Accurate modelling for electron… · p. 4
  • For a strongly interacting glutamic acid–lysine system, OrbNet-Equi’s density error is 0.211±0.001%, compared with 1.47±0.02% for monomer density superposition.The authors identify density learning as an opportunity for transferable DFT models of extended systems at a small fraction of the cost of solving the Kohn–Sham equations from scratch.Source: Accurate modelling for electron… · pp. 4–5

Transferability on downstream tasks

  • OrbNet-Equi/SDC21 transfers without fine-tuning to downstream benchmarks, including drug-like conformer ranking, torsions, non-covalent interactions, and geometry optimization.It also predicts adiabatic ionization potentials more accurately than semi-empirical QM methods despite open-shell signatures being rare in training, suggesting potential transfer to unseen electronic states.Source: Transferability on downstream tasks · p. 5
  • On neutral, closed-shell GMTKN55 reactions containing CHONFSCl, OrbNet-Equi/SDC21 is on par with DFT across five reaction classes and outperforms ANI-2x and GFN methods.For charged closed-shell reactions with training-set elements but unfamiliar interaction or bonding types, it remains competitive with DFT on several task groups, but trails DFT on reaction barriers and intermolecular non-covalent interactions.Source: Transferability on downstream tasks · p. 6
  • Across all GMTKN55 tasks, OrbNet-Equi/SDC21 has the lowest median WTMAD among methods executable at semi-empirical QM cost.Its challenging extrapolative errors on W4-11 atomization energies and G21EA electron affinities parallel GFN’s lack of explicit treatment of triplet or higher-spin species.Source: Transferability on downstream tasks · p. 6

Discussion

OrbNet-Equi is a QM-informed geometric deep-learning framework that represents molecular and material properties in the atomic-orbital basis. The authors report data-efficient learning across energy- and real-space properties and competitive accuracy on main-group benchmarks at speeds on par with semi-empirical methods.

Discussion

  • The framework learns molecular or material properties using representations in the atomic-orbital basis.The authors describe possible applications including electrolyte-property screening, QM/MM simulations for catalysis, and modeling reactive biochemical processes.Source: Discussion · p. 6
  • OrbNet-Equi shows data efficiency for both energy-space and real-space properties, broadening the molecular properties modeled by QM-informed machine learning.Source: Discussion · p. 6
  • On comprehensive main-group quantum-chemistry benchmarks, OrbNet-Equi offers an accuracy alternative to DFT methods at computation speeds on par with semi-empirical methods.Source: Discussion · p. 6
  • The authors identify adding simple model systems for absent interaction types as a possible way to improve accuracy and chemical-space coverage.They also say the framework can extend to alternative quantum-chemistry models for molecular or material systems.Source: Discussion · p. 6

1 Methods

UNiTE maps atomic-orbital input features to property predictions through equivariant atom-wise representations, iterative interaction modules, and a final pooling layer. Its modules encode symmetry-aware feature construction and updates, while EvNorm separates representation scale from direction.

1.1 The UNiTE neural network

  • UNiTE constructs atom-wise representations and iteratively updates them with block convolution, message passing, and point-wise interaction modules before pooling final representations into predictions.The stated setup assumes one input feature matrix for the explicit module expressions; extensions and proofs are discussed in appendices.Source: §1.1 · pp. 7–8
  • Each atom representation concatenates neurons associated with O(3) irreducible representations, indexed by channel, degree, and parity.The network is described as satisfying translation invariance and rotation and parity-inversion equivariance for the specified inputs.Source: §1.1 · p. 7

1.1.1 The Diagonal Reduction module

  • The diagonal reduction module decomposes diagonal input blocks into spherical-tensor representations to initialize equivariant atom-wise representations.A physically motivated construction uses on-site overlap integrals to produce the embeddings.Source: §1.1.1 · p. 8

1.1.2 The Block Convolution module

  • Block convolution contracts off-diagonal input blocks with atom representations, treating the blocks as convolution kernels and producing block-wise representations.The passage describes this as a local convolution.Source: §1.1.2 · pp. 8–9

1.1.3 The Message Passing module

  • Message passing sums block-wise representations over neighboring block indices, with scalar weights parameterized as SE(3)-invariant multi-head attention.The attention mechanism is described as increasing network capacity without increasing memory costs compared with explicitly expanding the tensor.Source: §1.1.3 · p. 9

1.1.4 Equivariant Normalization (EvNorm)

  • EvNorm separates normalized frame-invariant scale information from direction-like information, then uses the direction in the subsequent representation update.The authors report that EvNorm stabilizes training and removes the need to tune weight initializations and learning rates across tasks.Source: §1.1 · p. 7, §1.1.4 · p. 9

1.1.5 The Point-wise Interaction module and representation updates

  • Point-wise interaction nonlinearly updates atom-wise representations by coupling them with an O(3)-equivariant representation, using normalization and learned transformations.For early updates, representations are combined with aggregated messages; later updates are based on local information.Source: §1.1.5 · p. 10

1.1.6 Pooling layers and training

  • A pooling layer reads out the target prediction from the final updated representations.For dipole prediction, the example combines predicted atomic charges and atomic dipoles with atom coordinates.Source: §1.1.6 · p. 10

1.2 QM-informed featurization details and gradient calculations

The method builds QM-informed features from mean-field electronic structure in an atomic-orbital basis, using GFN-xTB tight-binding calculations. The network learns target properties from these matrix features, and energy gradients can be differentiated to obtain forces.

1.2 QM-informed featurization details and gradient calculations

  • The QM-informed representation uses atomic-orbital-basis features directly, without heuristic post-processing to enforce rotational invariance.Source: §1.2 · p. 10
  • GFN-xTB supplies the mean-field electronic structure through a self-consistent tight-binding calculation, with Fock, density, core-Hamiltonian, and overlap matrices used as input features.The inputs may be extended, for example, with energy-weighted density matrices.Source: §1.2 · pp. 10–11
  • The described electronic-structure setup assumes a closed-shell singlet ground state for simplicity.Source: §1.2 · p. 11
  • Differentiating the predicted energy with respect to nuclear coordinates yields energy-conserving forces, computed through an adjoint approach with the SAAO transformation matrix replaced by the identity.Source: §1.2 · p. 11

Training datasets

The training data combine selected published dimers with a larger collection of drug-like and biofragment conformers. DFT single-point energies provide labels for the SDC21 geometries used to train OrbNet-Equi/SDC21.

Training datasets

  • The BFDb-SSI subset contains 2,291 sidechain–sidechain dimers with at most 25 atoms and no sulfur, selected to support comparisons among methods.Source: Training datasets · p. 11
  • The SDC21 dataset combines ChEMBL-derived conformers, protonation and salt variants, enumerated bonding-pattern conformers, and molecules from BFDb-SSI and JSCH-2005.The dataset spans 16 listed elements and contains 237,298 geometries before filtering.Source: Training datasets · p. 11
  • DFT single-point energies were calculated for SDC21 geometries at the ωB97X-D3/def2-TZVP level, and filtering left 235,834 geometries for model training.Geometries were removed when DFT failed to converge or broken bonds were detected between equilibrium and non-equilibrium geometries.Source: Training datasets · p. 11

Electronic structure computational details

The QM9 dipole labels use B3LYP/def2-TZVP to match the published labels, while QM9 and BFDb-SSI density labels use ωB97X-D3/def2-TZVP. Density expansion coefficients are obtained using a def2-TZVP-JKFIT fitting basis.

Electronic structure computational details

  • QM9 dipole moments were calculated at B3LYP/def2-TZVP to match the level of theory used for published QM9 labels.The calculations used Entos Qcore version 1.1.0.Source: Electronic structure computational… · p. 12
  • QM9 and BFDb-SSI electron densities were computed at ωB97X-D3/def2-TZVP, using def2-TZVP-JKFIT for Coulomb and exchange fitting and as the density expansion basis.The density expansion coefficients are calculated from the DFT AO density matrix and overlap quantities.Source: Electronic structure computational… · p. 12

Benchmarking details and summary statistics

The benchmark statistics use specified test sets and evaluation conventions for density and downstream tasks. GMTKN55 reaction errors are referenced to CCSD(T)/CBS, with reaction-level and aggregate WTMAD-related measures.

Benchmarking details and summary statistics

  • Density-error tests use 291 BFDb-SSI dimers and 10,000 QM9 molecules, with voxel spacings of 0.2 and 1.0 Bohr, respectively.Both grids use a density cutoff of 10^-5 a^-3.Source: Benchmarking details and summary… · p. 12
  • Density-error averages follow each baseline’s normalization convention: molecule-normalized errors for QM9 and electron-normalized errors for BFDb-SSI.Electron normalization produced averages around 5% higher than molecule normalization in the authors’ results.Source: Benchmarking details and summary… · p. 12
  • Downstream conformer energies are compared with DLPNO-CCSD(T), while geometry optimization accuracy is symmetry-corrected RMSD against ωB97X-D3/def2-TZVP geometries.Median per-molecule R2 values use bootstrapped molecule-pool error bars; TorsionNet500 and s66x10 use 95% confidence intervals.Source: Benchmarking details and summary… · p. 12
  • GMTKN55 uses CCSD(T)/CBS reference values and reaction-level WTAD scores defined similarly to WTMAD-2.The supplement also gives subset-wise and overall WTMAD-2 measures.Source: Benchmarking details and summary… · p. 12

S1 Additional neural network details

The supplement details implementation choices and property-specific pooling for UNiTE. These include array-based spherical-harmonic operations, input-channel handling, restricted coupling indices, geometric messages, and readouts for several molecular properties.

S1.1 Efficient GPU evaluation of spherical harmonics and Clebsch-Gordan coefficients

  • Spherical harmonics are evaluated through element-wise array operations and gatherscatter operations, while Clebsch-Gordan coefficients are tabulated and transformed from complex to real harmonics.The spherical-harmonic scheme uses coordinate powers and pre-tabulated coefficients.Source: S1.1 Efficient GPU evaluation of… · p. 19

S1.2 Multiple input channels

  • A learnable linear layer maps stacked AO features into fixed-size channels, with each channel assigned to a subset of convolution channels.In these experiments, each input channel is assigned its own convolution channel.Source: S1.2 Multiple input channels · p. 19

S1.3 Restricted summands in Clebsch-Gordan coupling

  • For computational efficiency, Clebsch-Gordan coupling retains only angular-momentum pairs satisfying l1 + l2 ≤ lmax and l1, l2 ≤ n.Here lmax is the maximum angular momentum in the implementation.Source: S1.3 Restricted summands in… · p. 19

S1.4 Incorporating geometric information

  • The message-passing scheme incorporates atomic-coordinate directions using spherical harmonics and learnable linear functions.The direction vector is defined between atomic centers A and B.Source: S1.4 Incorporating geometric information · p. 19

S1.5.1 Energetic properties

  • Molecular electronic energy is rotation-invariant and extensive, so its pooling sums atom-wise contributions with learnable layers and element-specific biases.Element biases are initialized by linear regression on training labels; this scheme is used for several QM9 targets, MD17 energies, and OrbNet-Equi/SDC21.Source: S1.5.1 Energetic properties · p. 20

S1.5.2 Dipole moment⃗µ

  • Dipole moments are pooled from atomic charges and atomic dipoles; the scheme ensures translation invariance through charge neutrality.For published QM9 labels, which provide only the dipole norm, training uses the loss on the predicted vector norm.Source: S1.5.2 Dipole moment⃗µ · p. 20

S1.5.3 Polarizability α

  • The supplement defines a pooling operation for isotropic polarizability α.The supplied passage introduces the operation but does not state its form.Source: S1.5.3 Polarizability α · p. 20

S1.5.4 Molecular orbital properties

  • Frontier orbital energies use global-attention pooling for intensive predictions, and the HOMO-LUMO gap is the predicted HOMO energy minus the predicted LUMO energy.Element-specific biases are initialized by linear fitting on the targets.Source: S1.5.4 Molecular orbital properties · p. 20

S1.5.5 Electronic spatial extent ⟨R2⟩

Electronic density is predicted through an atom-centered expansion with learned weights specific to element and angular momentum. UNiTE inference scales with input sparsity, and its convolution can be evaluated using only nonzero tensor coefficients.

S1.5.5 Electronic spatial extent ⟨R2⟩

  • The predicted electron density is compactly parameterized by an atom-centered expansion with learnable weights specific to atomic number and angular momentum.The density fitting basis is def2-TZVP-JKFIT; other UNiTE parameters are independent of atomic number.Source: S1.5.6 Electron densities ρ(⃗r) · p. 21

S1.6 Time complexity

  • UNiTE inference has O(BNI) complexity, scaling as O(N(nd)N) for dense inputs but potentially lower for sparse inputs such as those with long-range cutoffs.The convolution can be evaluated exactly using only nonzero tensor coefficients, and embedding and point-wise interaction blocks have O(d) complexity.Source: S1.6 Time complexity · p. 21

S2 Theoretical results

The supplement formalizes N-body tensors and the symmetries their coefficients must respect, then establishes equivariant decompositions and network building blocks for UNiTE. The construction relies on tensor symmetry for permutation-equivariant diagonal reduction; asymmetric tensors are left for future treatment.

S2.1 N-body tensors (informal)

  • N-body tensors encode relations among collections of geometric objects, with coefficients that change under local-frame transformations despite the underlying tensor being frame-independent.The definition also requires translation invariance and symmetry under tensor-index permutations.Source: S2.1 N-body tensors (informal) · pp. 21–22

S2.2 Equivariance

  • G-equivariance requires the network to respect local unitary basis transformations, tensor-index permutations, and global translations.For atomic orbitals, spherical harmonics provide bases for irreducible representations of SO(3).Source: S2.2 Equivariance · p. 23, S2.3 N-body tensors · p. 23

S2.3 N-body tensors

  • The formal N-body tensor definition extends beyond Cartesian point-cloud bases to irreducible-representation bases, including the spherical-harmonic basis used for atomic orbitals.The supplement also recovers the Cartesian cases when the representation spaces use C^n or R^n.Source: S2.3 N-body tensors · p. 23

S2.4 Decomposition of diagonals Tu

  • Tensor-product representation theory supplies an invertible map that decomposes diagonal sub-tensors into order-1 representations transforming under irreducible representations.For symmetric tensors, the diagonal reduction is permutation-invariant; the general asymmetric-tensor case is out of scope.Source: S2.4 Decomposition of diagonals Tu · pp. 24–25

S2.5 Generalized neural network building blocks

  • UNiTE generalizes its point-wise interaction, matching, convolution, message-passing, and EvNorm operations to the formal N-body tensor setting.The supplement defines feature-channel and representation indices for these generalized building blocks.Source: S2.5 Generalized neural network… · pp. 25–26

S2.6 G-equivariance

  • The supplement proves G-equivariance for UNiTE by combining the diagonal-reduction result with equivariance of its convolution, message-passing, normalization, and point-wise interaction blocks.Permutation equivariance uses the symmetric-tensor condition and permutation-invariant diagonal reduction.Source: S2.6 G-equivariance · pp. 26–28

S3 Supplementary numerical results

The supplementary results document datasets, evaluation settings, and additional benchmarks for OrbNet-Equi. On MD17 and rMD17, direct learning is competitive with reported methods, while delta learning yields the lowest test MAEs for energy and forces on most molecules.

S3.1 The QM9 dataset

  • Table S1 reports QM9 test MAEs for 12 properties, with OrbNet-Equi trained on 110k samples using delta learning for the first eight tasks and direct learning for the last four.Test-MAE uncertainties are estimated by bootstrapping with sample size 5000 and 100 iterations.Source: Table S1 · p. 28, S3.1 The QM9 dataset · p. 28

S3.3 Inference timings

  • Inference timings on the Hutchison drug-like molecule dataset are reported as energy-and-force calculation breakdowns in Table S4.Timing measurements use 16 cores of an Intel Xeon Gold 6130 CPU.Source: Table S4 · p. 29, S3.3 Inference timings · p. 30

S3.2 MD17 and rMD17 datasets

  • Delta learning yields the lowest test MAEs for both energy and forces on most MD17 and rMD17 test molecules, while direct learning is competitive with reported kernel methods and graph neural networks.The study notes that MD17 is an interpolative setting with abundant high-level reference calculations for a single chemical system.Source: S3.2 MD17 and rMD17 datasets · pp. 29–30

S3.4 Downstream benchmarks

  • Tables S5–S7 summarize downstream benchmark performance, including GMTKN55 comparisons against CCSD(T)/CBS and a filtered OrbNet-Equi/SDC21 evaluation on reactions with training-seen elements and electronic states.The filtered column restricts evaluation to reactions containing elements and electronic states present in SDC21 training.Source: S3.4 Downstream benchmarks · pp. 30–31

S4 Hyperparameters and training details

The supplement specifies OrbNet-Equi hyperparameters and training procedures across its benchmarks. It details optimizer settings and dataset-specific losses, schedules, training sizes, and compute requirements.

S4.1 Model hyperparameters

  • OrbNet-Equi benchmark experiments use shared model hyperparameters, with the SDC21 model locally optimized on a 4219-sample validation set and modified normalization settings for robustness on extrapolative geometries.The hyperparameters are summarized in Tables S8 and S9.Source: S4.1 Model hyperparameters · p. 32, Table S8 · p. 32, Table S9 · p. 32

S4.2 Training

  • Training uses Adam with a maximal learning rate of 5 × 10^-4 and SmoothL1Loss, while batch sizes and epoch counts vary by benchmark.No additional weight decay or early stopping is used.Source: S4.2 Training · p. 32

S4.2.1 QM9

  • QM9 models use target-specific losses, batch sizes of 64 or 16 depending on training-set size, and 100 warmup epochs followed by 200 cosine-annealing epochs.Runs with 110k training samples take around 36 hours on one Nvidia Tesla V100-SXM2-32GB GPU.Source: S4.2.1 QM9 · p. 32

S4.2.2 MD17

  • For MD17, models use energy/force loss weights cE=1, cF=1000 on rMD17, versus cE=0, cF=1000 on original MD17, training on 1000 ‘train 01’ geometries per molecular system.Each training run takes around 30 hours for 1500 epochs on a single Nvidia Tesla V100-SXM2-32GB GPU.Source: S4.2.2 MD17 · p. 33

S4.2.3 Electron densities

  • Electron-density models use an analytic L2 density loss and sparse COO overlap matrices for batched training.Training takes 10 hours on BFDb-SSI and 120 hours on QM9 using one Nvidia Tesla V100-SXM2-32GB GPU.Source: S4.2.3 Electron densities · p. 33

S4.2.4 The OrbNet-Equi/SDC21 model

  • OrbNet-Equi/SDC21 training samples a geometry from each molecule's available geometries within each mini-batch and uses cG = 10.Training runs for 125 epochs with batch size 64 and cosine annealing, taking 64 hours on one Nvidia Tesla V100-SXM2-32GB GPU.Source: S4.2.4 The OrbNet-Equi/SDC21 model · p. 33

Related papers

Loading 2105.14655v4…