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

Abstract:
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 ∗These authors contributed equally. †Work done at TU Berlin prior to joining Google Research. 35th Conference on Neural Information Processing Systems (NeurIPS 2021). 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. While we focus on physics applications in this contribution, the proposed equivariant framework for deep learning on point clouds is promising also beyond, say, in computer vision or graphics.

text: ∗These authors contributed equally. †Work done at TU Berlin prior to joining Google Research.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
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. While we focus on physics applications in this contribution, the proposed equivariant framework for deep learning on point clouds is promising also beyond, say, in computer vision or graphics.

heading: 1 Introduction
text: Machine learning (ML) methods are becoming increasingly popular in quantum chemistry as a means to circumvent expensive ab initio calculations, and led to advances in a broad range of applications, including the construction of potential energy surfaces [1–8], prediction of electron densities and density functionals [9–14], and development of models capable of predicting a range of physical observables across chemical space [15–29]. Typically, such models are trained on reference data for a predetermined set of quantum chemical properties and need to be retrained if other properties are required. However, if a model is capable of predicting the wavefunction, expectation values for any observable can be derived from it. Unfortunately, such an approach is complicated by the fact that wavefunctions are typically expressed in terms of rotationally equivariant basis functions, introducing non-trivial transformations under molecular rotations, which are difficult to learn from data. To solve this issue, we propose several SE(3)-equivariant operations for deep learning architectures for geometric point cloud data, which capture the effects of translations and rotations without needing to learn them explicitly. We assemble these building blocks to construct PhiSNet, a novel deep learning (DL) architecture for predicting wavefunctions and electronic densities, which is significantly more accurate than non-equivariant models. For the first time, sufficient accuracy is reached to predict properties like energies and forces directly from the wavefunction and in end-to-end manner.
This makes it possible to learn wavefunctions that lead to modified properties, which is interesting from an inverse design perspective; or the development of novel machine-learned semi-empirical methods, for example by learning a correction to the wavefunction that mimics the effects of electron correlation. Such hybrid methods maintain the accuracy and generality of high level electronic structure calculations while drastically reducing their computational cost. In addition, the predicted wavefunctions can serve as initial guess to speed up conventional ab initio methods. Beyond physics, other applications of our proposed equivariant DL architecture to e.g. computer vision or graphics are conceivable – whenever accurate invariant analyses of high dimensional point clouds are of importance.
In summary, this work provides the following contributions:
• We describe general SE(3)-equivariant operations and building blocks for constructing DL architectures for geometric point cloud data.
• We propose PhiSNet, a neural network for predicting wavefunctions and electronic densities from equivariant atomic representations, ensuring physically correct transformation under translations and rotations.
• We apply PhiSNet to predict wavefunctions and electronic densities of several molecules and show that our model reduces prediction errors of electronic structure properties by a factor of up to two orders of magnitude compared to the previous state-of-the-art and achieves speedups of over three orders of magnitude compared to ab initio solutions.
• We showcase a novel transfer-learning application, where a model trained on low accuracy wavefunctions is adapted to predict properties computed at a higher level of theory by learning a correction that implicitly captures the effects of many-body electron correlation.
• We demonstrate that the predicted wavefunctions can serve as initial guess in conventional quantum chemistry methods, leading to significant speedups without sacrificing the accuracy or robustness of ab initio solutions.
In principle, our method could also be used to construct orbital features as inputs for methods like OrbNet [30], which otherwise rely on semi-empirical or ab initio methods.

heading: 2 Related work
text: Only a small number of studies apply ML to the challenging problem of modeling the wavefunction directly [31]. This is usually done by predicting Hamiltonian matrices, from which the wavefunction can be obtained by solving a generalized eigenvalue problem. The earliest such study we are aware of is by Hegde and Bowen [32], who used kernel ridge regression to learn the Hamiltonian matrix for two simple case studies. Later, Schütt et al. [33] proposed the SchNOrb neural network architecture, which constructs the Hamiltonian matrix of molecules in a block-wise manner from atom-pair features. Recently, Li et al. [34] presented a deep neural network architecture for predicting the Hamiltonian matrix of simple periodic crystals. However, none of these models form their predictions in a rotationally equivariant manner, i.e. they need to learn how to predict the Hamiltonian matrix for all possible orientations of the system of interest, which requires large amounts of training data. Even when data augmentation via random rotations [35] or special Hamiltonian representations [36] are used to mitigate this issue, the final model is only approximately equivariant, i.e. properties derived from the wavefunction can change unphysically when the system is rotated or the frame of reference is changed. Here, we draw upon insights from a range of SE(3)-equivariant models [37–48] to ensure that predictions exactly preserve the physically correct dependence with respect to the orientation of inputs (for a more detailed discussion see the supplement). We would like to remark that after the initial submission of this manuscript, Nigam et al. [49] introduced a complementary method to construct equivariant representations for Hamiltonian matrices, e.g. for the use in kernel machines.

heading: 3 Background
text: The aim of most quantum chemistry methods is to solve the electronic Schrödinger equation
ĤelΨel = EelΨel, (1)
where Ĥel is the Hamiltonian operator describing the interactions and motion of the electrons, Ψel is the electronic wavefunction and Eel is the ground state energy. After Ψel is determined, all physical observables (beyond Eel) can be derived by applying other operators (similar to Ĥel) to the wavefunction and reading out the corresponding eigenvalues [50]. In practice, Eq. 1 is usually solved by expressing Ψel as an antisymmetric product of molecular orbitals
ψi = ∑ j Cijφj , (2)
which are written as linear combinations of atom-centered basis functions φ. This leads to the equation
HC = SC , (3) where the Hamiltonian is written as a matrix H with entries Hij = ∫ φ∗i (r)Ĥelφj(r)dr (r denotes
the electronic coordinates). The overlap matrix S with entries Sij = ∫ φ∗i (r)φj(r)dr has to be introduced and a generalized eigenvalue problem must be solved, because the basis functions φ are typically not orthonormal. The eigenvectors C specify the wavefunction Ψel via the coefficients Cij of the molecular orbitals (Eq. 2) and the eigenvalues are the corresponding orbital energies. Since eigenvectors are only defined up to sign changes, predicting H (instead of C) is preferable for ML applications. This is a challenging task, because the basis functions are typically products of a radial component and spherical harmonics, which introduces non-trivial dependencies of the matrix elements with respect to the orientation of the chemical system (see Fig. 1A). Spherical harmonics Y ml of degree l = 0, . . . ,∞ and order m = −l, . . . , l form a complete orthonormal basis for functions on the surface of a sphere and can be used to derive irreducible representations (irreps) of the 3D rotation group SO(3). In other words, they are rotationally equivariant, which means that when r is rotated, the values of Y ml (r) change accordingly. Since Ψel is expressed with spherical harmonic basis functions, the entries of H transform predictably under rotations (a more detailed overview of quantum chemistry fundamentals, groups, equivariance, and spherical harmonics is given in Section A of the supplement).

heading: 4 Deep learning architecture for molecular wavefunctions
text: Deep message-passing neural networks (MPNNs) [51] for quantum chemistry applications, such as DTNN [18] or SchNet [19], model physical properties of chemical systems as a sum over atomic contributions predicted from features xi ∈ RF for each atom i. Starting from initial element-specific embeddings, the features are constructed by iteratively exchanging “messages” between neighboring atoms i and j, which depend on the current xi and xj and their distance rij . Since geometric information enters only in the form of pairwise distances, the final atomic features are rotationally invariant by construction. This is desirable when they are used to predict a quantity that itself is rotationally invariant, for example the potential energy. To predict observables that change under rotation, e.g. electric moments or the electronic Hamiltonian, a natural extension is to instead construct rotationally equivariant features. To see how this can be achieved, it is useful to think of the F entries of atomic feature vectors xi ∈ RF as different “channels”, where each channel carries scalar information about the chemical environment of atom i. To construct rotationally equivariant features, each scalar channel can be replaced by values derived from the spherical harmonics up to a maximum degree L, i.e. there are now F × (L+ 1)2 entries (each spherical harmonic degree l ∈ {0, . . . , L} contributes 2l + 1 values for all possible orders m ∈ {−l, . . . , l}). Our proposed model, which we call PhiSNet, shares basic design principles with PhysNet [23], but uses equivariant (instead of invariant) operations throughout its architecture. Contrary to most other MPNNs, instead of directly predicting chemical properties from atomic features, PhiSNet constructs the Hamiltonian matrix in a block-wise manner from equivariant representations. All known physical symmetries of the Hamiltonian are preserved by construction, an essential constraint which is not satisfied by existing DL architectures for chemical applications.
Notation Whenever equivariant features are discussed, bold symbols (e.g. x) refer to the collection of all F × (L+1)2 entries for all feature channels F and spherical harmonics degrees l ∈ {0, . . . , L}, whereas a superscript l in parentheses (e.g. x(l)) is used to refer only to the F × (2l + 1) entries of degree l. Similarly, Y(r) refers to the collection of all 1× (L+ 1)2 spherical harmonics with distinct combinations of l and m, whereas Y(l)(r) refers to the 1× (2l+ 1) values for degree l. The notation c ◦ x denotes a Hadamard product between matrices, i.e. c ∈ RF×(L+1)2 and x ∈ RF×(L+1)2 are multiplied entry-wise. When there is no one-to-one correspondence between the entries in c and x, e.g. when c ∈ RF and x ∈ RF×(L+1)2 , c ◦ x implies that c is “broadcasted” across the missing dimensions, i.e. each of the (L+ 1)2 “slices” of x is multiplied entry-wise with c and the result c ◦ x has dimensions F × (L+ 1)2. Double-struck digits denote an irreducible representation (irrep) with the corresponding number of dimensions 2l + 1. For example, 1 refers to a one-dimensional irrep of degree l = 0 and 3 to a three-dimensional irrep of degree l = 1. By abuse of terminology, the term “irrep” is also used for the individual (2l + 1)-dimensional components along the F feature dimensions of x ∈ RF×(L+1)2 . When a collection of equivariant features x(l) ∈ RFin×(2l+1) is multiplied by a matrix M ∈ RFout×Fin , the result is Mx(l) ∈ RFout×(2l+1), i.e. the ordinary rules for matrix multiplication apply.

heading: 4.1 SE(3)-equivariant neural network building blocks
text: In the following, we describe general-purpose operations for building SE(3)-equivariant MPNNs, which can also be used outside a chemical context to build feature representations for other point cloud data. Additionally, we discuss necessary modifications to established neural network components (e.g. linear layers or activation functions) for keeping feature representations rotationally equivariant.
Activation functions may only be applied to scalar (l = 0) features, or else the output loses its equivariant properties:
σ(x)(l) =
{ σ(x(l)) l = 0
x(l) l > 0 . (4)
Here, σ can be any activation function and the notation σ(x(l)) means that σ is applied to x(l) entry-wise. In this work, a generalized SiLU [52, 53] activation function (also known as Swish [54]) given by
σ(x) = αx
1 + e−βx (5)
is used, where α and β are both learned and separate parameters are kept for all feature channels and instances of σ [27] (see Section B.1 in the supplement for additional details).
Linear layers are applied to each degree l according to
linearFin→Fout(x) (l) =
{ Wlx (l) + b l = 0
Wlx (l) l > 0
, (6)
where W ∈ RFout×Fin and b ∈ RFout are weights and biases, respectively. The subscript l is used to distinguish the weights for different degrees l, i.e. separate linear transformations are applied to the features of each degree l. The bias term must be omitted for l > 0 so that output features stay rotationally equivariant.
Tensor product contractions are used to couple two equivariant feature representations x(l1) and y(l2) to form new features z(l3). The (reducible) tensor product x(l1) ⊗ y(l2) of two irreps has (2l1+1)(2l2+1) dimensions and can be expanded into a direct sum of irreducible representations, e.g. 3⊗5 = 3⊕5⊕7. In general, the value for orderm3 of the irrep of degree l3 ∈ {|l1−l2|, . . . , l1+l2} in the direct sum representation of the tensor product x(l1) ⊗ y(l2) is given by(
x(l1) ⊗ y(l2) )l3 m3 = l1∑ m1=−l1 l2∑ m2=−l2 Cl3,l2,l1m3,m2,m1x l1 m1y l2 m2 , (7)
whereCl3,l2,l1m3,m2,m1 are Clebsch-Gordan coefficients (CGCs) [55]. The short-hand notation x (l1)⊗ l3 y(l2) is used to refer to the irrep of degree l3 in the direct sum representation of x(l1)⊗y(l2). In other words,
the operation x(l1) ⊗ l3 y(l2) performs the tensor product and contracts the result to a single irrep of degree l3. A similar construction is also used in Clebsh-Gordan [56] and Tensor Field networks [42].
Tensor product expansions are inverse to tensor product contractions. Instead of contracting two irreps into one, CGCs are used to expand a single irrep x(l3) into a (2l1 + 1) × (2l2 + 1) matrix that represents its contribution to the direct sum representation of the tensor product of two irreps of degree l1 and l2, where |l2 − l1| ≤ l3 ≤ l2 + l1:
( ⊗x(l3) )l1,l2 m1,m2 = l3∑ m3=−l3 Cl1,l2,l3ml,m2,m3x l3 m3 . (8)
Mirroring the shorthand used for tensor product contractions, l1,l2
⊗ x(l3) will be used to refer to the (2l1 + 1)× (2l2 + 1) matrix that is obtained from the tensor product expansion of x(l3).
Selfmix layers are used to recombine (“mix”) the F features of a single input x ∈ RF×(Lin+1)2 across different degrees and optionally allow changing the maximum degree from Lin to Lout. The output features of degree l3 are given by
selfmixLin→Lout(x) (l3) =  kl3 ◦ x(l3) + ∑Lin l1=0 ∑Lin l2=l1+1 sl3,l2,l1 ◦ ( x(l1) ⊗ l3 x(l2) ) l3 ≤ Lin∑Lin
l1=0 ∑Lin l2=l1+1 sl3,l2,l1 ◦ ( x(l1) ⊗
l3 x(l2)
) l3 > Lin .
(9) Here, k, s ∈ RF are learnable coefficients and the subscripts are used to distinguish independent parameters k, s for different degrees: In total, a selfmix layer has Lout + 1 different kl3 (one for each possible value of l3 ∈ {0, . . . , Lout}) and (Lout + 1)Lin(Lin+1)2 different sl3,l2,l1 (one for each valid combination of l3, l2, l1).
Spherical linear layers Spherical linear layers are a combination of linear (Eq. 6) and selfmix (Eq. 9) layers given by
sphlinearLin→Lout,Fin→Fout(x) = linearFin→Fout (selfmixLin→Lout(x)) . (10)
Chaining both operations allows arbitrary combinations across feature channels and degrees while still preserving rotational equivariance. In principle, whenever Lin = Lout, selfmix layers are not strictly necessary and Eq. 10 may be replaced by Eq. 6 for a boost in computational efficiency and reduction of memory footprint. However, this comes at the cost of reduced accuracy (see Section D in the supplement for details).
Residual blocks are modules consisting of two sequential spherical linear layers (see Eq. 10) and activation functions (see Eqs. 4 and 5) inspired by the pre-activation residual block described in [57]:
residual(x) = x + sphlinear2(σ2(sphlinear1(σ1(x)))) . (11)
Here, Fin = Fout and Lin = Lout for both spherical linear layers.
Pairmix layers are used to combine a pair of features x ∈ RF×(Lx+1)2 and y ∈ RF×(Ly+1)2 with a scalar r (e.g. their Euclidean distance) to generate new features of degree Lout:
pairmixLx,Ly→Lout(x,y, r) (l3) = Lx∑ l1=0 Ly∑ l2=0 (Wl3,l2,l1g(r)) ◦ ( x(l1) ⊗ l3 y(l2) ) . (12)
Here, g(r) ∈ RK is the vector [g0(r) g1(r) . . . gK−1(r)]> and gk(r) are radial basis functions. In this work, exponential Bernstein polynomials [27] are used (see Section B.2 and Eq. S9 in the supplement). The weight matrices W ∈ RF×K allow to learn radial functions as linear combinations of the basis functions gk(r) and subscripts are used to distinguish independent weights for different combinations of l1, l2, l3 (in total, there are (Lx + 1)(Ly + 1)(Lout + 1) possible combinations).
Interaction blocks use message-passing to model interactions between the features c ∈ RF×(L+1)2 of a central point i with features n ∈ RF×(L+1)2 of neighboring points j within a local environment:
a(x, r)(l) = x0 ◦ (Wlg(‖r‖)) ◦ sphlinearL→L,1→F (Y(r))(l) ,
b(x, r)(l) = pairmixL,L→L(x, sphlinearL→L,1→F (Y(r)) , ‖r‖)(l) ,
interaction(c,n, r) (l) i = c (l) i + ∑ j 6=i ( a(nj , rij) (l) + b(nj , rij) (l) ) .
(13)
Here rij is the distance vector rij = rj − ri between the positions ri, rj of i and j. The radial basis function expansion g is the same as in pairmix layers and W ∈ RF×K are independent weight matrices for each degree l. Since geometric information enters Eq. 13 via relative distance vectors rij expanded in a spherical harmonics basis, interactions blocks are equivariant with respect to the SE(3) group of roto-translations.

heading: 4.2 PhiSNet architecture
text: PhiSNet takes as inputs nuclear charges Z and positions r of N atoms, which are used to construct equivariant feature representations encoding information about the chemical environment of each atom. These features are then further transformed and used to predict the entries of the Hamiltonian matrix, see below. An overview over the complete architecture is shown in Fig. 1B and more detailed diagrams of individual building blocks are given in Fig. S2 (see Section D of the supplement for ablation studies that explore the impact of possible simplifications of the PhiSNet architecture on prediction accuracy).
Atomic feature representations An embedding layer produces initial atomic feature representations x from the nuclear charges Z according to
embedding(Z)(l) =
{ WdZ + bZ l = 0
0 l > 0 , (14)
where W ∈ RF×4 is a weight matrix and bZ element-specific biases with learnable parameters. Here, dZ ∈ R4 are fixed vectors for each element that contain information about their nuclear charge and ground state electron configuration, similar to the embeddings described in [27] (see Section C.1 in the supplement for details). The features are then refined by five sequential modules, each consisting of identical building blocks with independent parameters:
t = residual(x) ,
i = sphlinearL→L,F→F (σ (residual(t))) ,
j = sphlinearL→L,F→F (σ (residual(t))) ,
v = sphlinearL→L,F→F (σ (residual(interaction(i, j, r)))) ,
x̃ = residual (t + v) ,
ỹ = residual (x̃) .
(15)
Each module produces two different outputs x̃ and ỹ. The first output x̃ serves as input to the next module in the chain (replacing x in Eq. 15), whereas the second output ỹ is summed with the outputs of other modules m to form the final atomic feature representations f = ∑ m ỹm (see Fig. S2A for a visual representation).
Hamiltonian matrix prediction The Hamiltonian matrix is constructed block-by-block, with each block corresponding to the interaction between two atoms i and j. Diagonal and off-diagonal blocks are treated separately, i.e. different atomic pair features are constructed to predict them. Since these transformations also involve interactions with neighboring atoms (similar to interaction blocks), separate representations for central c and neighboring atoms n are created from the atomic features f :
ci = residual(fi) ,
nj = residual(fj) . (16)
Self-interaction features fii (for diagonal blocks) are constructed according to
f (l) ii = residual c(l)i +∑ j 6=i ( n (l) j ◦Wlg(‖rij‖) ) , (17) where rij is the distance vector between the positions of atoms i and j, and the radial basis function expansion g is the same as in pairmix layers. The matrices W ∈ RF×K are separate trainable weight matrices for each degree l.
Similarly, pair-interaction features fij (for off-diagonal blocks) are obtained by combining the central representations of atoms i and j and interacting them with the neighbors of atom i according to
f (l) ij = residual pairmix(ci, cj , ‖rij‖)(l) + ∑ k/∈{i,j} ( n (l) k ◦Wlg(‖rik‖) ) . (18) All blocks Hii and Hij of the Hamiltonian matrix, each representing the interaction between two atoms, are themselves composed of smaller blocks corresponding to the interaction between atomic orbitals. Since atomic orbitals are expressed in a spherical harmonics basis, their interactions transform non-trivially (put predictably) under rotations. The correct equivariant behavior of a matrix block Ml1,l2 ∈ R(2l1+1)×(2l2+1) corresponding to the interaction between orbitals of degree l1 and l2 can be constructed as a sum over matrices obtained from tensor product expansions (Eq. 8) of irreps a (collected from specific channels of the pairwise features fii or fij , see below) of all valid degrees l3 ∈ {|l2 − l1|, . . . ,≤ l2 + l1}:
Ml1,l2 = l2+l1∑ l3=|l2−l1| l1,l2 ⊗ a(l3) . (19)
Two sets of indices Iself and Ipair count and keep track of the irreps necessary to construct all required matrices Ml1,l2 . For diagonal blocks Hii, the irreps of a given degree l for the interaction of orbitals n and m of atoms with nuclear charge Z are collected from specific channels of the self-interaction features fii via a unique index Iself(Z, n,m,L). Similarly, for off-diagonal blocks Hij , a unique index Ipair(Zi, Zj , ni, nj , L) selects irreps corresponding to orbitals ni and nj of atom pairs with nuclear charges Zi and Zj from the pair-interaction features fij . After all blocks have been constructed, the complete matrix H̃ is obtained by placing individual blocks at the appropriate positions (based on which atoms and orbitals interact). Finally, a Hamiltonian matrix satisfying the necessary Hermitian symmetry is constructed as H = H̃+ H̃T . This symmetrization guarantees that both pair-features fij and fji contribute equally to the corresponding off-diagonal blocks and also makes sure that sub-blocks of the Hamiltonian swap positions in the correct way when equivalent atoms are permuted. In cases where multiple Hamiltonian-like matrices need to be predicted, all parameters up to the final residual blocks in Eqs. 17 and 18 are shared. An exception are overlap matrices (see Section 3), for which simpler self- and pair-interaction features derived directly from the embeddings are sufficient (see Section C.3 in the supplement for details). The complete block-wise construction process of the Hamiltonian matrix from irreps is illustrated in Fig. S5 for a water molecule with minimal basis set.

heading: 5 Results and discussion
text: To assess the ability of PhiSNet to predict molecular wavefunctions and electronic densities (see Fig. 1C), we train it on Kohn-Sham (the Kohn-Sham matrix takes the role of the Hamiltonian in DFT methods, see Section A.1 in the supplement) and overlap matrices for various non-equilibrium configurations of water, ethanol, malondialdehyde, uracil, and aspirin computed at the density functional theory (DFT) level with PBE/def2-SVP. Datasets for all molecules are taken from [33], with the exception of aspirin, for which geometries were sampled from the MD17 dataset [5] (more details on the datasets, training procedure, and hyperparameter settings can be found in Sections F and G of the supplement).
The results are summarized in Tab. 5 and compared to the current state-of-the-art given by SchNOrb [33]. PhiSNet achieves accuracy improvements up to two orders of magnitude, with
the biggest differences arising in larger and more complex molecules like uracil and aspirin. Note that the training process for SchNOrb requires data augmentation via random rotations to approximate the equivariance relation between wavefunction and molecular orientation, while PhiSNet preserves exact equivariance. This not only allows for faster convergence, but also leads to much smaller model sizes, with our model requiring approximately one fifth of the parameters of SchNOrb while providing significantly more accurate results. In addition, PhiSNet provides speedups of over three orders of magnitude compared to DFT calculations (see Section E.1 for details).
The improved prediction accuracy provided by PhiSNet makes it possible to accurately derive properties such as energies and forces directly from the wavefunction in an end-to-end manner, enabling a number of interesting and novel ML applications for the molecular sciences. As an example, we showcase a transfer-learning application, where a model trained on low accuracy HartreeFock (HF) electronic structure calculations is fine-tuned to learn a correction to the wavefunction, such that energies and forces match those obtained via high-level coupled cluster with singles, doubles, and perturbative triple excitations (CCSD(T)) calculations. The CCSD(T) method models the effects of electronic many-body interactions, which are neglected in HF theory, and is often considered to be the “gold standard” of quantum chemistry [8]. However, its accuracy comes at a significantly increased computational complexity, going from O(N3) for HF to O(N7) for CCSD(T) (here, N is the number of basis functions). Thus, a machine-learned correction to HF theory, which mimics the effects of electron correlation in a computationally efficient manner, is a possible way towards novel hybrid methods that rival the accuracy of high level electronic structure calculations and combine the generality and robustness of ab initio methods with the efficiency of ML.
After pretraining PhiSNet on HF/cc-pVDZ data for Fock matrices F, core Hamiltonians Hcore, and overlap matrices S, we fine-tune it on forces computed at the CCSD(T)/cc-pVTZ level, but still keep loss terms for Hcore and S computed with HF/cc-pVDZ (see Section A.1 in the supplement for a brief overview of HF theory, where we explain the relevance of F, Hcore, and S, and how energies and forces are derived from them). This way, the model learns to adapt only the Fock matrix F, which embodies the electron-electron interactions in the HF formalism. Interestingly, very subtle changes to F seem to be sufficient to approximate the effects of electron correlation, resulting in an “effective CCSD(T) wavefunction” that, at first glance, appears to be almost identical to its original HF-level counterpart (see Fig. 1D). Nonetheless, the modified wavefunction reduces the mean absolute errors (MAEs) between energies and forces predicted with PhiSNet to just 79 µEh and 0.85 mEha−10 , respectively, compared to the CCSD(T) reference. In contrast, the original HF-level wavefunction leads to MAEs of 4266 µEh and 15 mEha−10 for energies and forces, respectively, i.e. prediction errors are reduced by over one order of magnitude with no additional computational overhead. More details on the transfer learning application can be found in Section G.2 of the supplement.
Although PhiSNet predicts electronic properties very accurately, some applications might require exact solutions, making existing quantum chemistry methods preferable over predictions of an
ML model. Even then, PhiSNet can be used to achieve speedups without any loss of accuracy or robustness: A significant fraction of compute time in HF and DFT calculations must be spent to arrive at a self-consistent solution of Eq. 3. A good initial guess for the wavefunction can reduce the number of required iterations significantly. We observe a decrease between 56–72% in the number of iterations, resulting in a total reduction of wall-clock time between 24–40%, when using PhiSNet-predicted wavefunctions in place of the default guess (see Section E.2 in the supplement for more details).
While the results are highly promising, the current approach also has limitations. Due to the fact that all pairwise combinations of atoms have to be considered for constructing the pair-interaction features fij , the Hamiltonian matrix prediction scales quadratically with the number of atoms. Further, to derive the coefficient matrix C defining the wavefunction, a generalized eigenvalue problem has to be solved (see Eq. 3), which scales cubically with the matrix size (the number of basis functions). For these reasons, PhiSNet does not scale well to systems with a very large number of atoms. However, possible extensions of our method could exploit the fact that orbital overlaps between distant atoms are very small, so the entries of corresponding matrix blocks are approximately zero [58]. Then, only pair-interaction features for non-zero blocks need to be computed and the solution of Eq. 3 can exploit sparsity.

heading: 6 Conclusion
text: For learning problems with known invariances, equivariances, symmetries, or other constraints, as is common in physics applications, it is useful to include such properties directly into the model architecture. This effectively reduces the complexity of the learning problem (cf. [59, 60]), increasing model performance and decreasing the required amount of training data. While invariance and equivariance properties could also be (approximately) learned, this typically requires much more reference data and/or data augmentation. In contrast, by “hard coding” domain knowledge, the learning is constrained to a meaningful submanifold, e.g. reflecting rotational equivariance [38], energy conservation, physical laws or symmetries (e.g. [4, 5, 23, 27, 61]), group equivariance (e.g. [37]), graph properties (e.g. [40, 46]) or alike. Thus, known properties do not need to be learned explicitly, because the submanifold where learning takes place already embodies them.
Our present contribution follows this design principle, specifically, we describe a series of general SE(3)-equivariant operations and building blocks for deep learning architectures operating on geometric point cloud data, which we used here to construct PhiSNet, a novel neural network architecture capable of accurately predicting wavefunctions and electronic densities. Unlike previous models, which need to approximately learn how the wavefunction transforms under molecular rotations and rely on data augmentation, the SE(3)-equivariant3 building blocks of our network allow to exactly capture the correct transformation without needing to learn it explicitly. By applying PhiSNet on a range of small to medium-sized molecules, we demonstrated that our model achieves accuracy improvements of up to two orders of magnitude compared to the previous state-of-the-art (while at the same time requiring significantly less parameters), and speedups with respect to ab initio solutions of over three orders of magnitude.
For the first time, sufficient accuracy is reached to derive quantum mechanical observations directly from the predicted wavefunction in an end-to-end manner, which also allows to adapt the predicted wavefunction such that it leads to desired physical properties. To showcase such an application, we fine-tuned a model trained on low accuracy wavefunctions to predict properties computed at a much higher level of theory, thereby learning a correction that implicitly captures the effects of many-body electron correlation. This paves the way towards the development of novel semi-empirical methods that are capable of providing highly accurate quantum chemical calculations at a drastically reduced computational cost. In addition, we demonstrated that existing ab initio approaches can benefit from guess wavefunctions provided by our method, achieving significant speedups without any loss of accuracy or robustness.
Although we focus on quantum chemistry applications in this contribution, we would like to reiterate that the presented SE(3)-equivariant operations are general and can be used to construct other deep learning architectures for geometric point cloud data beyond physics, e.g. [37, 38, 40, 45, 46].
3Note that SE(3)-equivariance is important (in contrast to just SO(3)-equivariance), because the properties of chemical systems may change substantially when individual atoms are translated relative to each other.
Acknowledgments and Disclosure of Funding
OTU acknowledges funding from the Swiss National Science Foundation (Grant No. P2BSP2_188147). MB acknowledges support by the Federal Ministry of Education and Research (BMBF) for BIFOLD (01IS18037A). KRM was supported in part by the Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea Government (No. 2019-0-00079, Artificial Intelligence Graduate School Program, Korea University), and was partly supported by the German Ministry for Education and Research (BMBF) under Grants 01IS14013A-E, 01GQ1115, 01GQ0850, 01IS18025A and 01IS18037A; the German Research Foundation (DFG) under Grant Math+, EXC 2046/1, Project ID 390685689. M. Gastegger works at the BASLEARN Joint Lab for Machine Learning, co-financed by TU Berlin and BASF SE. There are no competing interests to declare.

Review summaries:
 This contribution describes a method that approximates quantum mechanical total energies and wavefunctions with compute costs that are similar to those of Hartree Fock calculations. The paper describes a novel architecture, PhiSNet, that predicts wavefunctions from nuclear charges and positions: an equivariant featurization of the atomic information feeds a network that is built using a number SE(3) equivariant building blocks that in turn feeds to a block that builds the Hamiltonian matrix. The authors demonstrate that PhiSNet can train to predict DFT total energies to accuracies of about 0.11 kcal/mol for aspirin and better for smaller molecules. They further demonstrate that the network can "correct" a HF level wavefunction to an effective wavefunction that includes the effects of electron correlation. 

The paper describes a method to predict the wavefunctions and electron densities using a rotationally equivariant method. The method achieves sufficient accuracy for the model to be used to predict properties such as energies and forces directly from atom positions and charges. The approach is evaluated by comparing to an existing rotationally-invariant approach using prediction errors in Kohn-Sham matrices, energies and wavefunction cosine similarity as metrics to compare the methods. Background on relevance and significance as well as implementation details are provided in the supplementary information along with sample code.

This paper describes PhiSNet, a model for predicting the Hamiltonians of systems of atoms in the context of quantum chemistry. The output Hamiltonian can then be used to derive the wavefunction using traditional methods, and the wavefunction can in turn be used to derive physically observable properties of interest. A core part of the design of this model is built-in equivariance to rotation in three dimensions. The paper defines the architecture of the PhiSNet and each of its equivariant component layers in detail, and results are presented from a comparison to the competing method SchNOrb.