PaperlayerView exact revision on arXiv

Source-linked AI summary

SE(3)-equivariant prediction of molecular wavefunctions and electronic densities

Oliver T. Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, Klaus-Robert Müller

arXiv:2106.02347v2physics.chem-phstat.ML

TL;DR

Wavefunctions are a central prediction target because observables can be derived from them, but their nontrivial rotational transformations make them difficult to learn. The paper introduces SE(3)-equivariant operations and PhiSNet to predict wavefunctions and electronic densities, reporting substantial accuracy and speed gains. The accuracy also supports direct property derivation and transfer learning toward higher-level calculations.

  • Problem

    Wavefunction prediction could provide a route to deriving observables beyond a fixed set of properties, but wavefunctions transform nontrivially under molecular rotations, making them difficult to learn.

    Source: §1 · p. 2

  • Method

    The paper builds PhiSNet from SE(3)-equivariant operations to predict wavefunctions and electronic densities from equivariant atomic representations.

    Source: §1 · p. 2

  • Results

    Up to two orders of magnitude lower prediction errors than the previous state-of-the-art and over three orders of magnitude speedups compared with DFT calculations are reported for PhiSNet.

    Source: §5 · pp. 8–9

  • Takeaways & Limitations

    The reported accuracy permits direct end-to-end derivation of energies and forces from wavefunctions and supports transfer learning toward higher-level calculations.

    Source: §1 · p. 2, §5 · p. 9

    Hamiltonian prediction scales quadratically with atom count and solving the generalized eigenvalue problem scales cubically with basis-function count, limiting scalability to very large systems.

    Source: §5 · p. 10

Abstract

from arXiv · show

Machine learning has enabled the prediction of quantum chemical properties with high accuracy and efficiency, allowing to bypass computationally costly ab initio calculations. Instead of training on a fixed set of properties, more recent approaches attempt to learn the electronic wavefunction (or density) as a central quantity of atomistic systems, from which all other observables can be derived. This is complicated by the fact that wavefunctions transform non-trivially under molecular rotations, which makes them a challenging prediction target. To solve this issue, we introduce general SE(3)-equivariant operations and building blocks for constructing deep learning architectures for geometric point cloud data and apply them to reconstruct wavefunctions of atomistic systems with unprecedented accuracy. Our model achieves speedups of over three orders of magnitude compared to ab initio methods and reduces prediction errors by up to two orders of magnitude compared to the previous state-of-the-art. This accuracy makes it possible to derive properties such as energies and forces directly from the wavefunction in an end-to-end manner. We demonstrate the potential of our approach in a transfer learning application, where a model trained on low accuracy reference wavefunctions implicitly learns to correct for electronic many-body interactions from observables computed at a higher level of theory. Such machine-learned wavefunction surrogates pave the way towards novel semi-empirical methods, offering resolution at an electronic level while drastically decreasing computational cost. Additionally, the predicted wavefunctions can serve as initial guess in conventional ab initio methods, decreasing the number of iterations required to arrive at a converged solution, thus leading to significant speedups without any loss of accuracy or robustness.

1 Introduction

The paper introduces SE(3)-equivariant operations and PhiSNet to predict wavefunctions and electronic densities while respecting molecular translations and rotations. It reports large accuracy and speed gains and demonstrates transfer learning and wavefunction-based initialization for quantum chemistry methods.

1 Introduction

  • Accurate wavefunctions enable direct end-to-end derivation of properties such as energies and forces, and may support machine-learned semi-empirical methods or wavefunction-based inverse design.The introduction also notes possible use of wavefunctions as inputs to methods such as OrbNet.Source: §1 · p. 2
  • PhiSNet predicts wavefunctions and electronic densities from equivariant atomic representations, preserving physically correct transformations under translations and rotations.The paper presents general SE(3)-equivariant operations and building blocks for geometric point-cloud architectures.Source: §1 · p. 2
  • Up to two orders of magnitude lower prediction errors and over three orders of magnitude faster than ab initio solutions are reported for PhiSNet.The comparison is against the previous state-of-the-art for errors and ab initio solutions for speed.Source: §1 · p. 2
  • The authors demonstrate transfer learning from low-accuracy wavefunctions by learning a correction that implicitly captures many-body electron-correlation effects.The correction adapts predictions toward properties computed at a higher level of theory.Source: §1 · p. 2
  • Predicted wavefunctions can serve as initial guesses in conventional quantum chemistry methods, speeding them up without sacrificing ab initio accuracy or robustness.Source: §1 · p. 2

2 Related work

Earlier direct-wavefunction approaches predict Hamiltonian matrices, but the models described here do not form predictions rotationally equivariantly. The paper draws on SE(3)-equivariant models to preserve the physically correct orientation dependence.

2 Related work

  • Prior direct-wavefunction approaches predict Hamiltonian matrices, from which wavefunctions are obtained by solving a generalized eigenvalue problem.Source: §2 · p. 3
  • Earlier examples include kernel ridge regression for two case studies, SchNOrb for molecules, and a deep network for simple periodic crystals.The passage identifies Hegde and Bowen, Schütt et al., and Li et al., respectively.Source: §2 · p. 3
  • The discussed prior models lack rotational equivariance, requiring them to learn across orientations; augmentation and special representations only make equivariance approximate.The paper says approximate equivariance can leave derived properties changing unphysically with rotation or frame of reference; its approach draws on SE(3)-equivariant models [37] [38] [39] [40] [41] [42] [43] [45] [46] [47] [48].Source: §2 · p. 3

3 Background

Quantum chemistry seeks solutions to the electronic Schrödinger equation, whose wavefunction can be used to derive physical observables. In practice, molecular-orbital basis functions lead to a generalized eigenvalue problem, and their rotational behavior makes Hamiltonian prediction challenging.

3 Background

  • The electronic Schrödinger equation relates the electronic Hamiltonian, wavefunction, and ground-state energy.Source: §3 · p. 4
  • After determining the wavefunction, other physical observables can be derived by applying operators and reading out their eigenvalues.Source: §3 · p. 4
  • Expressing the wavefunction through molecular orbitals in atom-centered basis functions introduces a generalized eigenvalue problem because the basis functions are typically nonorthonormal.The Hamiltonian and overlap matrices enter this problem, and its eigenvectors specify the wavefunction through molecular-orbital coefficients.Source: §3 · p. 4
  • Hamiltonian prediction is challenging because spherical-harmonic basis functions make its matrix entries depend nontrivially on molecular orientation.The spherical harmonics transform equivariantly under rotations, so the Hamiltonian entries transform predictably.Source: §3 · p. 4

4 Deep learning architecture for molecular wavefunctions

PhiSNet uses equivariant atomic representations to construct molecular Hamiltonians while preserving known physical symmetries. Its architecture combines general SE(3)-equivariant building blocks with a block-wise construction of orbital interactions.

4 Deep learning architecture for molecular wavefunctions

  • PhiSNet uses equivariant representations to construct Hamiltonian matrices block-wise and preserve known physical symmetries.Source: §4 · p. 4

4.1 SE(3)-equivariant neural network building blocks

  • The building blocks combine degree-specific linear maps, tensor-product mixing, and neighbor interactions using relative vectors expanded in spherical harmonics to retain SE(3) equivariance.Source: §4.1 · pp. 5–7
  • Non-scalar equivariant features cannot pass through ordinary activation functions; generalized SiLU uses learned α and β parameters.Source: §4.1 · p. 5

4.2 PhiSNet architecture

  • Atomic features start from nuclear-charge embeddings and are refined through five modules whose outputs feed later modules and sum into the final representations.Source: §4.2 · p. 7
  • Tensor-product expansions map orbital-degree irreps into matrix blocks, while separate self- and pair-interaction features supply diagonal and off-diagonal content.Source: §4.2 · pp. 7–8
  • Constructing H as Htilde + Htilde^T enforces Hermitian symmetry, makes pair features contribute equally, and handles the required block swaps under equivalent-atom permutations.Source: §4.2 · p. 8

5 Results and discussion

PhiSNet predicts wavefunctions and electronic densities with higher accuracy and speed than prior approaches. The results also demonstrate transfer learning toward higher-level properties and faster conventional calculations using predicted wavefunctions as initial guesses.

5 Results and discussion

  • Training used DFT/PBE/def2-SVP Kohn-Sham and overlap matrices for non-equilibrium configurations of water, ethanol, malondialdehyde, uracil, and aspirin.The aspirin geometries were sampled from MD17; the other molecule datasets were taken from SchNOrb.Source: §5 · p. 8
  • Up to two orders of magnitude fewer prediction errors than SchNOrb, with the largest differences for uracil and aspirin.Table 1 reports prediction errors across molecules, matrices, energies, and occupied-orbital wavefunctions; PhiSNet also provides over three orders of magnitude speedup versus DFT.Source: §5 · pp. 8–9, Table 1 · p. 9
  • PhiSNet uses about one fifth of SchNOrb’s parameters and preserves exact equivariance, while SchNOrb uses rotation augmentation to approximate it.Exact equivariance allows faster convergence; speedups over DFT exceed three orders of magnitude.Source: §5 · p. 9
  • Fine-tuning HF/cc-pVDZ data on CCSD(T)/cc-pVTZ forces yielded 79 µEh energy and 0.85 mEha−1 force MAEs, versus 4266 µEh and 15 mEha−1 for the HF-level wavefunction.The correction adapts the Fock matrix; reported errors are reduced by over one order of magnitude without additional computational overhead.Source: §5 · p. 9
  • Using PhiSNet wavefunctions as initial guesses reduced solver iterations by 56–72% and wall-clock time by 24–40%, without loss of accuracy or robustness.Conventional HF and DFT calculations may still be preferable when exact solutions are required.Source: §5 · pp. 9–10
  • Hamiltonian prediction scales quadratically with atom count and solving the generalized eigenvalue problem cubically with basis-function count, limiting very large systems.The authors suggest exploiting small distant-atom orbital overlaps and matrix sparsity as a possible extension.Source: §5 · p. 10

6 Conclusion

The paper presents general SE(3)-equivariant building blocks and uses them in PhiSNet to predict molecular wavefunctions and densities. It reports strong accuracy and speed results for small-to-medium molecules, alongside potential uses in other point-cloud applications and quantum chemistry workflows.

6 Conclusion

  • Hard-coding known symmetries constrains learning to meaningful submanifolds, reducing learning complexity and the need to learn those properties explicitly.The paper notes that invariance or equivariance learned approximately can require more reference data or data augmentation.Source: §6 · p. 10
  • General SE(3)-equivariant operations for geometric point clouds form the building blocks used to construct PhiSNet for wavefunction and density prediction.They capture the wavefunction’s transformation under molecular rotations exactly rather than learning it approximately through data augmentation.Source: §6 · p. 10
  • On small-to-medium molecules, PhiSNet improves accuracy by up to two orders of magnitude over prior work and speeds ab initio solutions by over three orders of magnitude.The model also requires significantly fewer parameters.Source: §6 · p. 10
  • The presented operations are intended for point-cloud architectures beyond quantum chemistry, including applications outside physics.The paper names other geometric point-cloud applications as examples.Source: §6 · p. 10
  • SE(3), rather than rotation-only SO(3), matters because translating atoms relative to one another can substantially change chemical-system properties..Source: §6 · p. 10

Supplementary information for SE(3)-equivariant prediction of molecular wavefunctions and electronic

The supplement reviews the quantum-chemistry and equivariance foundations behind the approach, then describes the equivariant building blocks and radial basis functions used in the model.

A.1 Quantum chemistry

  • Hartree–Fock neglects electron correlation, while post-HF accuracy comes at high cost; DFT also relies on an approximate exchange-correlation functional.The supplement notes O(N^3) scaling for HF and DFT, versus O(N^7) for coupled cluster theory.Source: §A.1 · p. 16

A.2 Group representations and equivariance

  • Composing equivariant functions preserves equivariance, and SO(3) representations decompose into irreducible components of dimension 2l + 1.The supplement identifies Wigner-D matrices as the irreducible linear operators of SO(3).Source: §A.2 · p. 16

A.3 Spherical harmonics

  • Real spherical harmonics provide the rotational basis used to represent functions on a sphere and relative atomic positions.The supplement defines real-valued harmonics by degree l and order m.Source: §A.3 · pp. 16–17

A.4 Group equivariance in machine learning

  • PhiSNet’s equivariant operations aim to adapt invariant message-passing design principles, adding selfmix-based coupling across degrees and channels and tensor-product expansions for higher-order tensors.The spherical linear layer is presented as a drop-in replacement for linear layers in invariant architectures.Source: §A.4 · p. 18

B.1 Properties of the generalized SiLU activation function

  • The generalized SiLU activation interpolates between linear and ReLU behavior, with α = 1.0 and β = 1.702 chosen to approximate GELU.As β approaches infinity it approaches a scaled ReLU; at β = 0 it is linear.Source: §B.1 · p. 18

B.2 Exponential Bernstein polynomial basis functions

  • Distance-dependent equivariant operations use exponential Bernstein radial basis functions with a learnable length scale and a smooth cutoff.The mapping x = e^−γr provides an inductive bias for exponential decay, and basis expansions across layers share γ for efficiency.Source: §B.2 · pp. 18–19

C Additional details on the PhiSNet architecture

The architecture uses element descriptors for embeddings, assembles Hamiltonian blocks from irreducible representations, and predicts overlap features with simplified representations.

C.1 Element descriptors used in embeddings

  • Element descriptors encode nuclear charge and ground-state electron configuration; the implementation covers periods 1 and 2 with entries for Z, 1s, 2s, and 2p electrons.The descriptors can be extended with additional shell entries for higher-period elements.Source: §C.1 · pp. 19–20

C.2 Details on the block-wise Hamiltonian matrix prediction

  • Hamiltonian blocks for orbital pairs are assembled from irreps whose degrees span the angular-momentum coupling range, then accumulated into the full matrix.For water in a minimal basis, the supplement describes a 7 × 7 Hamiltonian and the irreps required for diagonal and off-diagonal blocks.Source: §C.2 · pp. 20–21, Table S2 · p. 21

C.3 Overlap matrix prediction

  • Because overlap integrals depend on element types and relative orientation and distance rather than many-body effects, the overlap matrix uses simplified self- and pair-interaction features.These features are transformed through spherical linear and residual layers before the overlap matrix is assembled.Source: §C.3 · pp. 21–22

D Ablation studies

Ablations show that most tested simplifications reduce prediction accuracy without meaningful efficiency gains; removing selfmix is the exception, trading accuracy for faster inference and training.

D Ablation studies

  • The tested simplifications share matrix features or interaction features, remove selfmix, or construct self- and pair-interaction features without neighboring-atom interactions.The table reports Kohn–Sham matrix prediction errors for ablated variants against the baseline.Source: D Ablation studies · pp. 22–23, Table S3 · p. 23
  • The no-selfmix variant is 14–77% less accurate than the baseline but reduces inference speed and training time by roughly 50%.The other considered modifications significantly reduce accuracy, while their efficiency gains are negligible (<1%).Source: D Ablation studies · p. 23, Table S3 · p. 23

E.1 Model inference times compared to ab initio calculations

  • The speed comparison measures single-structure wall-clock times for PhiSNet inference and PBE/def2-SVP DFT calculations using ORCA 4.0.1.PhiSNet runs on an NVIDIA A100 GPU with batch size 64; DFT runs on Intel Xeon E5-2690 CPUs.Source: §E.1 · p. 23, Table S4 · p. 23

E.2 Speeding up ab initio calculations using predicted wavefunctions as initial guess

The supplement evaluates predicted wavefunctions as initial guesses for conventional calculations, which can help when direct predictions are not accurate enough for exact-solution applications.

E.2 Speeding up ab initio calculations using predicted wavefunctions as initial guess

  • When small prediction errors remain unacceptable, PhiSNet-predicted wavefunctions can serve as initial guesses for conventional ab initio calculations.The supplement compares convergence iterations and wall-clock time against the default guess.Source: §E.2 · p. 23, Table S5 · p. 24
  • The quality of the predicted guesses also makes it possible to use quadratically converging SOSCF and NRSCF algorithms instead of standard DIIS.The reported comparison covers different molecules and SCF algorithms.Source: §E.2 · p. 24, Table S5 · p. 24
  • Table S5 compares default and PhiSNet-predicted guess wavefunctions by convergence-iteration reduction and total wall-clock time across molecules and SCF algorithms.The supplied passage identifies the comparison scope but gives no numerical table values.Source: Table S5 · p. 24

F Datasets

The datasets cover five molecules and combine DFT reference data with higher-level water data for transfer learning.

F Datasets

  • DFT datasets cover water, ethanol, malondialdehyde, uracil, and aspirin; the aspirin data came from a separate source.Except for water, structures are subsets of MD17; water structures were sampled with a classical force field.Source: F Datasets · p. 24
  • DFT quantities were computed at PBE/def2-SVP, while water transfer-learning data also use HF/cc-pVDZ and CCSD(T)/cc-pVTZ.The higher-level calculations use the same water structures as the DFT data.Source: F Datasets · p. 24
  • The datasets include molecular structures and quantum-chemical properties; matrix quantities have dimensions Nbasis × Nbasis.Table S6 describes molecule counts, atom counts, basis-function counts, and the properties available for each dataset.Source: Table S6 · p. 24

G Training procedure and hyperparameters

The models share a common architecture configuration and were trained with AMSGrad on a single GPU.

G Training procedure and hyperparameters

G.1 DFT datasets

DFT models are trained against matrix targets, while the HF-to-CCSD(T) transfer procedure adds energy and force supervision and requires matrix-loss warm-up.

G.1 DFT datasets

  • DFT models minimize a loss over predicted and reference Kohn-Sham and overlap matrices, using the squared Frobenius norm.Batch, training, validation, and test set sizes are reported in Table S7; batch sizes vary with matrix size and hardware memory constraints.Source: §G.1 · p. 25, Table S7 · p. 25

G.2 Transfer learning from HF to CCSD(T)

  • Transfer learning uses the same structures with HF/cc-pVDZ and CCSD(T)/cc-pVTZ data, split into 500 training, 500 validation, and 3,999 test structures.Source: §G.2 · p. 25
  • After adding energy and force information to the loss, energy and force predictions improve substantially, while matrix-element prediction deteriorates slightly.The initial matrix-only loss produces matrix errors comparable to DFT results, but much larger errors in derived energies and forces.Source: §G.2 · pp. 25–26
  • Direct training with the combined loss does not converge; a warm-up phase using only matrix loss is necessary to avoid numerical instabilities.The instabilities occur when solving the generalized eigenvalue problem with matrices predicted early in training.Source: §G.2 · p. 26
  • The HF-to-CCSD(T) energy comparison requires a fitted constant shift because the two levels have a constant energy offset.The shift is fit on training-set energies and added before test-set evaluation; absolute energies can be shifted without changing the underlying physics.Source: §G.2 · p. 26
Loading 2106.02347v2…