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

∗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 . 1 Introduction . 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 . 2 Related work . 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 . 3 Background . 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 ) . 4 Deep learning architecture for molecular wavefunctions . 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 .

Summary:  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.