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