SE(3)-equivariant prediction of molecular wavefunctions and electronic densities
Authors: Oliver Thorsten Unke
Mihail Bogojeski
Michael Gastegger
Mario Geiger
Tess Smidt
Klaus Robert Muller
Venue: NeurIPS 2021 Poster (2021)
Keywords: Deep learning
quantum chemistry
steerable networks
message passing networks

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

-Euiveeqiv prediction of molecular wavefunctions and electronic densities
Oliver T. Unke 1 oliver. unke@gmail. com
Mihail Bogojeski! 1,8+
miihaibbgojjskti jesoiegmsiiillll.o.o com
Michael Gastegger 1,2
Mario Geiger
Tess Smidt
Klaus-Robert Müller 1,6,7,8,10
Zlaus-robert.mme www.leerrtu-errli..d.e... www.bberrlli.ee
Machine Learning Group, Technische Universität Berlin, 10587 Berlin, Germany
DFG Cluster of Excellence "Unifying Systems in Catalysis" (UniSysCat), Technische Universität Berlin,
10623 Berlin, Germany
Institute of Physies, École Polytechnique Fédérale de Lausanne, 1015 Lausanne, Switzerland
4 Computational Research Division, Lawrence Berkeley National Laboratory. Berkeley, CA 94720 Center for Advanced Mathematics for Energy Research Applications (CAMERA).
Lawrence Berkeley National Laboratory, Berkeley, CA 94720
Department of Artificial Intelligence, Korea University, Anam-dong, Scongbuk-gu Scoul 02841, Korea
Max Planck Institute for Informatics, Stuhlsatzenhausseeg. 66123 Saarbrücken, Germany
BIFOLD Berlin Institute for the Foundations of Learning and Data, Berlin, Germany
BASLEARN TU Berlin/BASF Joint Lab for Machine Learning. Technische Universität Berlin,
10587 Berlin, Germany
Google Research, Brain Team, Berlin, Germany
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 com- plicated by the fact that wavefunctions transform non-trivially under molecular rotations, which makes them challenging prediction target. To solve this issue, we introduce general Squi)-equivariant operations and building blocks for construct- ing deep learning architectures for geometric point cloud data and apply them reconstruct wavefunctions of atomistic systems with unprecedented accuracy. Our model achieves speedups of over three orders magnitude compared to ab initio methods and reduces prediction errors by up two orders of magnitude compared to the previous state-of-the-ant This accuracy makes it possible to derive properties such energies and forces directly from the wavefunction in end-to and-to-end manner. We demonstrate the potential of our approach where model trained 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
transfer learning application,
"These authors contributed equally. i Work done at TU Berlin prior to joining Google Research.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

level while drastically decreasing computational cost. Additionally, the predicted wavefunctions can serve as initial guess in conventional ab initio methods, decreas- ing the number of iterations required to arrive at 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 equívariant framework for deep learning on point clouds is promising also beyond, say, in computer vision or graphics.
Introduction
Machine learning (ML) methods are becoming increasingly popular in quantum chemistry as a means to eircumvent expensive ab initio calculations, and led to advances in 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 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 learn from data. To solve this issue, we propose several SE(3)-cquivaria) 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 con-equivariant models. For the first time, sufficient accuracy is reached to predict properties like energies and forces directly from the wavefunction and 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 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 SE3))-cquivarian operations and building blocks for constructing DL architectures for geometric point cloud data.
We propose PhiSNet, 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 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 novel transier-learning application, where model trained 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 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.

---

A
B
malondialdehydee
ethenol
uracil
asgirine
sel intractionn feabures
Fgy
pain inter actiion featurse
f
diggonal sclection
diagonal
HF-level
correction (amplified)
"CCllevel Hamiltonian'
Hamiltonian
Hockwine Homiltonian consmuction
HF-ievel
correction (amplified)
"CC-evell
wavefunction
wavefunction"
Figure 1: A: Illustration of an aspirine molecule and its highest occupied molecular orbital (HOMO) in three different orientations, showing how the wavefunction (left) and Hamiltonian matrix (right) change with respect rotations. B: Overview of the proposed PhiSNet architecture. The atomic representation network creates atom- wise equivariant features. which are used to produce self eeff-interaction and pair-interaction features (Fig. S2), from which the Hamiltonian matrix is constructed block-hy-black (Fig. S5). C: Visualisation electronic densities (squared wavefunction) of various molecules predicted with our approach. D: Illustration of a transfer learning application, where model pretrained on Hartree-Foct (HF) Hamiltonians is fine-tuned to match energies and forces derived from highly accurate coupled cluster (CC) calculations. The model achieves this by learning correction to Hamiltonian matrix, which mimics the effects of many-body electron correlation. The effective "CC-Ievel" Hamiltonian can be thought of as a HF-level Hamiltomian plus a correction term. The HOMO is shown to visualize subtle changes to the wavefunction (the correction is amplified in magnitude by factor of 10 for better visibility).
2 Related work
Only small number of studies apply ML to the challenging problem of modeling the wavefunction directly [31]. This usually done by predicting Hamiltonian matrices, from which the wavefunction be obtained by solving generalized eigenvalue problem. The earliest such study is is by Hegde and Bowen [32], who used kernel ridge regression to learn the Hamiltonian matrix for two simple case studies. Later, Schütt 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 deep neural network architecture for predicting the Hamiltonian matrix of simple periodic crystals. However, none of these models form their predictions rotationally equivariant manner, i.c. they need to learn how 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] special Hamiltonian representations [36] 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 range of SE(3) -E(3)-equivariant models [37-48] to ensure that predictions exactly preserve the physically correct dependence with respect to the orientation of inputs (for more detailed discussion see the supplement). We would like to remark that after the initial submission of this manuscript, Nigam et [49] introduced complementary method to construct equivariant representations for Hamiltonian matrices, e.g. for the use in kernel machines.
aware

---

3 Background
The aim of most quantum chemistry methods is to solve the electronic Schrödinger equation
Ha V el Val
where Her is the Hamiltonian operator describing the interactions and motion of the electrons, Va is the electronic wavefunction and Eel is the ground state energy. After Vo is determined, all physical observables (beyond Ea)) can be derived by applying other operators (similar Hay) the wavefunction and reading out the corresponding eigenvalues [50]. In practice, Eq. 1 is usually solved by expressing Vel as an antisymmetric product of molecular orbitals
v/ = Cujo
which are written as linear combinations of atom-centered basis functions & This leads to the equation
HC cSC
where the Hamiltonian is written matrix H with entries = o (F)dr (r denotes the electronic coordinates). The overlap matrix S with entries S.j = [ f (r) (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 Vaa via the coefficients Ci of the molecular orbitals (Eq. 2) and the eigenvalues c are the corresponding orbital energies. Since eigenvectors are only defined up 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 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. AA). Spherical -l,.... ..... form complete orthonormal basis
harmonics Y of degree 1 = 0,- -..... 0o and order for functions on the surface of sphere and can be used to derive irreducible representations (irreps) of the 3D rotation group (3). In other words, they are rotationally equivariant, which means that when r is rotated, the values of Y! (r) change accordingly. Since Val 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 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 sum over atomic contributions predicted from features x; E RR for each atom i. Starting from initial element-specine embeddings, the features are constructed by iteratively exchanging "messages" between neighboring atoms i and j. which depend on the current x; and x, and their distance Tij- Since geometric information enters only in the form of pairwise distances, the final atomic features 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 be achieved, it is useful to think of the F entries of atomic feature vectors x; E 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 i.e. there are now F' x (L 1)² entries (each spherical harmonic degree / E {O.....4) L} contributes 27- - values for all possible orders m E {-l,... .......... (T). Our proposed model, which we call PhiSNet, shares basic design principles with PhysNet [23], but uses equivariam (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 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 x (L+ 1)² entries for all feature channels F and spherical harmonics degrees E {o..... L). whereas superscript in parentheses (e.g. x(") is used refer only to the F x (21 + 1) entries of degree 1. Similarly, (r) refers to the collection of all 1 x (L + 1)² spherical harmonics with distinct combinations of and m, whereas Y (r) refers to the x (2l + 1) values for degree 1. The notation denotes Hadamard product between matrices, i.e. c E RFx(L+1)* and x E RPx(L+1)) multiplied entry-wise. When there is no one-to-one correspondence between the entries in с and x, e.g. when E IR and x E (Fx((++)* c implies that is "broadcasted" across the missing dimensions, i.e. each of the (I. + 1)2 "slices" of x is multiplied entry-wise with с and the result e o has dimensions F x (T. 2)2. Double-struck digits denote an irreducible representation (irrep) with the corresponding number of dimensions 2/ + 1. For example, 1l refers to a one-dimensional irrep of degree / = 0 and 3 to three-dimensional irrep of degree = 1. By abuse of terminology, the term "irrep" also used for the individual (21 + --dimensiooal components along the F feature dimensions of x RFx (L++) When collection of equivariant features x(") E R x(2++) is multiplied by matrix M E FR X the result is Mx E R Fou x (2++1) i.e. the ordinary rules for matrix multiplication apply.
4.1 SE(3)-equivarian neural network building blocks
In the following, we describe general-purpon 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 activation functions) for keeping feature representations rotationally equivariant.
Activation functions may only be applied to scalar equivariant properties:
0) features, or else the output loses its
(x(x()) (1))
[= 0 1> 0
a (x)
(4)
Here, a can be any activation function and the notation oxx()) means that o is applied x()) entry-wise. In this work, generalized SiLU [52, 53] activation function (also known as Swish [54]) given by
o
Ba
is used, where and 3 are both learned and separate parameters are kept for all feature channels and instances [27] (see Section B. the supplement for additional details).
Linear layers are applied to each degree I according to
lincar Fin Four (x) (Wx( +b 1= 0 (
1> 0
where W e RFu Fi and b RFaut weights and biases, respectively. The subscript / is used to distinguish the weights for different degrees (, i.c. separate lincar transformations are applied to the features of each degree I. The bias term must be omitted for I so that output features stay rotationally equivariant.
Tensor product contractions are used to couple two equivariant feature representations x(1) and (!2) form new features (ta)) The (reducible) tensor product ((1)) (122) of two irreps has (2)(+)) dimensions and can be expanded into a direct sum of irreducible representations, e.g.
385 - 3 In general, the value for order mg of the irrep of degree la {14 - 12/, .a...... (2)
in the direct sum representation of the tensor product ((1)) (12) is given by
12
(x(h) s y (())) Is =
daadads 1h d2 Ving
myg.
my=
My==-
da where my.ma.m are Clebsch-Gordan coefficients (CGCs) [55]. The short-hand notation x (4)) (yal
la
used to refer to the irrep of degree 13 in the direct sum representation of x (t1) (la)))) In other words,
5

---

the operation x @ y performs the tensor product and contracts the result to single irrep of degree 13- 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 (1aa into (211 + +1) x (212 + 1) matrix that represents its contribution the direct sum representation of the tensor product of two irreps of degree I and l2, where 1/2 41 s 1a
S la + 4
21,22
1s
xx)
Chilads my mm Tmy
myg my=-la
Mirroring the shorthand used for tensor product contractions, & (la)) will be used to refer to the (2/1 + 1) x (2/2 + 1) matrix that is obtained from the tensor product expansion of x(!s),
Selfimix layers are used to recombine ("mixx) the F features of a single input x E [F across different degrees and optionally allow changing the maximum degree from Lin Lout. The output features of degree 13 are given by
Ki₃ xl1s) + Eii Eben+ Sladadd o
x((1) @0 x((x)
13 S Tin
selfmix Loou (x)
Eneo 2tz S19.2241 (11) x(b)
13 > Lin
(9)
Here, k,s E IR are learnable coefficients and the subscripts are used to distinguish independent parameters k, for different degrees: In total, selfimix layer has Lout + different ki, (one for each possible value la E {O...., Lomt and (Lout 1) Lin (mn(hn+1) different 81g./241 (one for each valid combination of 13, 12)
Spherical linear layers Spherical linear layers are combination of linear (Eq. 6) and selfmix (Eq. 9) layers given by
sphlinear L. Los. (x) linear (selfmixi -To (x))
(10)
Chaining both operations allows arbitrary combinations across feature channels and degrees while still preserving rotational equivariance. In principle, whenever Lin Lours selfmmx layers are not strictly necessary and Eq. 10 may be replaced by Eq. 6 for a boost computational efficiency and reduction of memory footprint. However, this comes at the cost of reduced accuracy (see Section D in the supplement for details).
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 blocks
residual(x)
(pphiinearg(c2(sphlnnee) (x))))
(11)
Here, Fi Fout and Lin Lout for both spherical linear layers.
Pairmix layers used to combine pair of features E 1I) and y E (Fx,+)) with a scalar r (e.g. their Euclidean distance) to generate new features of degree Loout:
(x(h)) eyll)))
yai (x. v. y)
Σ (W (Wr))...g(()))
(12)
11 1==00=0
Here, g(r) E RK is the vector (r)(r) 91 (r) (r)i and (r) are radial basis functions. In this work, exponential Bernstein polynomials [27] are used (see Section B.2 and Eq. S9 in supplement). The weight matrices W E RF allow to learn radial functions as linear combinations of the basis functions (r) and subseripts are used distinguish independent weights for different combinations of 4 l, la (in total, there are (L, + 1) 1)(L 1) possible combinations).

---

Interaction blocks use message-passiinnne to model interactions between the features e E R RFx(1++1) of central point i with features ((*((++)) of neighboring points within local environment:
a(x, = x (Wig([ sphhineer,,-4,.--1 (
b(x, r) pairmix,, -.- sph.incar (Yrr). I=)) (13)
interaction(c, n, n,r), C
Σ j
a(n ( b(n),
Here raj is the distance vector rij rj r; between the positions rú, ry of i and j. The radial basis function expansion is the same as in pairmix layers and W E RF are independent weight matrices for each degree 1. Since geometric information enters Eq. 13 via relative distance vectors rij expanded in spherical harmonics basis, interactions blocks are equivariant with respect to the SE(3) group roto-translations.
4.2 PhiSNet architecture
PhiSNet takes as inputs nuclear charges Z and positions of N atoms, which are used construct cquivariant feature representations encoding information about the chemical environment of each atom. These features are then further transformed and used 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 representa- tions x from the nuclear charges Z according to
embedding(2) Wdz + bz I -= 0 1>
(14)
0
where W FR4 is weight matrix and by element-specific biases with learnable parameters. Here, dz E R$ 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.I in the supplement for details). The features are then refined by five sequential modules, each consisting of identical building blocks with independent parameters:
residual(x),
= sphlineary-1,FF=F (a (residua((t))
L-
= sphlinear
(residual(t))).
(15)
v sphlinear,,s1.F-41 (a (residual(iterrattio((....
X = residual (t+v),
y residual (x)
Each module produces two different outputs X and ÿ. The first output x serves as input to the next module in the chain (replacing x in Eq. 15), whereas the second output ý is summed with the outputs of other modules m to form the final atomic feature representations f = Σm m ym (see Fig. S2A for visual representation).
Hamiltonian matrix prediction The Hamiltonian matrix is constructed block- block-by-block, with each block corresponding to the interaction between two atoms 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 с and neighboring atoms are created from the atomic features
C; residual(f.).
(16)
n, = residual(f,).
7

---

Self-interaction features fę (for diagonal blocks) are constructed according to
residual
Σ 11) n
Wig((
(17)
where Fij is the distance vector between the positions of atoms and j, and the radial basis function expansion g is the same as in pairmix layers. The matrices W E RPx are separate trainable weight matrices for each degree
Similarly, pair-interaction features lyy (for off-diagonal blocks) are obtained by combining the central representations of atoms and j and interacting them with the neighbors of atom according to
-residual pairmix((; Cj [m)l))
(m)) Wig([raa)) (E
Σ
(18)
All blocks H;; and H:j 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 spherical harmonics basis, their interactions transform non-trivially (put predictably) under rotations. The correct equivariant behavior of a matrix corresponding to the interaction between orbitals of degree and
block M E +1)x++))x(22>>> l2 can be constructed as sum over matrices obtained from tensor product expansions (Eq, 8) of irreps (collected from specific channels of the pairwise features f or fay. see below) of all valid degrees 13 (1/2 11/p.... I <22 1 l
latts
4 X (fa)
M
Σ
(19)
T== 12-111
Two sets of indices psolf and Iaair count and keep track of the irreps necessary to construct all required matrices M For diagonal blocks Ha, the irreps of given degree / for the interaction of orbitals n and m of atoms with nuclear charge 7 are collected from specific channels of the self-interaction features fú via unique index T (Z, n. m, I.). Similarly, for off-diagonal blocks Hy, a unique index Tpair (Z,, Zj, Tuj n2 L) selects irreps corresponding to orbitals n2 and n, of atom pairs with nuclear charges Zi and Zj from the pair-interaction features fej. After all blocks have been constructed, the complete matrix H is obtained by placing individual blocks ar the appropriate positions (based on which atoms and orbitals interact). Finally, Hamiltonian matrix satisfying the necessary Hermitian symmetry is constructed as H H + H This symmetrization guarantees that both pair-features fúj and fjz 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, parameters up to the final residual blocks in Eqs. 17 and 18 are shared. An exception matrices (see Section 3), for which simpler self- and pair-interaction features derived directly from the embeddings 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.
overlap
5 Results and discussion
To assess the ability of PhiSNet to predict molecular wavefunctions and electronic densitics (sce Fig. IC), we train it Kohn-Sham (the Kohn-Sham matrix takes the role of the Hamiltonian in DFT methods, see Section A.I I in the supplement) and overlap matrices for various non couuiiiiii configurations of water, ethanol, malondialdehyde, uracil, and aspirin computed at the density functional theory (DFT) level with PBE/def2-SVi 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 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

---

Table 1: Prediction errors of PhiSNet for various molecules compared to SchNOrb [33]. In addition Kohn-Shan K and overlap matrices, we also report errors for energies and the cosine similarity between predicted and reference wavefunction z for all occupied orbitals. Best results in bold.
K
S
v
Data set
[10 En]
[10-6]
[10 6 En.]
SchNOrb PhiSNet
165.4
79.1
279.3
1.00
Water
17.59
1.56
85.53
1.00
SchNOrb
187.4
67.8
334.4
1.00
Ethanol
PhiSNet
12.15
0.626
62.75
1.00
SchNOrb
191.1
67.3
400.6
0.99
Malondialdehyde
PhiSNet
12.32
0.567
73.50
1.00
SchNOrb
227.8
82.4
1760
0.90
Uracil
PhiSNet
10.73
0.533
84.03
1.00
SchNOrb 506.0
110
48689
0.57
Aspirin
PhiSNet 12.84
0.406
176.6
0.98
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 DFT calculations (see Scction 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 end-to-end manner, enabling a number of interesting and novel ML applications for the molecular sciences. As an example, showcase transfer-learning application, where a model trained on low accuracy Hartree- Fock (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, 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 often considered to be the "gold standard" of quantum chemistry [8]. However, its accuracy comes at significantly increased computational complexity, going from ( for HF to O(N7) for CCSD(T) (here, N is the number of basis functions). Thus, a machine-learnen 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 MI..
doubles,
After pretraining PhiSNet on HF/cc-pVDZ data for Fock matrices F, core Hamiltonians Heore, and forces computed at the CCSD(T)cc-p)) VTZ level, but still keep
overlap matrices S, we fine-tune loss terms for H and S computed with HF/cc-p VDZ (see Section A. in the supplement for brief overview of HF theory, where we explain the relevance of F, H 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, first glance, appears be almost identical to its original HF-level counterpart (see Fig. ID). Nonetheless, the modified wavefunction reduces the mean absolute errors (MAEs) between energies and forces predicted with PhiSNet to just 79 μEE and 0.85 mE,,a,' respectively, compared the CCSD(T) reference. In contrast, the original HF-level wavefunction leads to MAEs of 4266 μdB, and 15 mE,a, for energies and forces, respectively, i.e. prediction errors are reduced by over one order of magnitude with no additional computational overhead. More details 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 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 total reduction of wall-clock time between 24 40%, when using PhiSNct-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 f.y, 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 very large number of atoms. However, possible extensions of our method could exploit the fact that orbital overlaps between distant atoms very small, 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. can exploit sparsity.
6 Conclusion
For learning problems with known invariances, equivariances, symmetries, or other constraints, as 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 series of general -quuivaaui operations and building blocks for deep learning architectures operating on geomet- ric 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)--quivariant 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 reached to derive quantum mechanical observations directly from the predicted wavefunction in an end-to-end manner, which also allows to adapt the predicted wavefunetion 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 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).
NNoe that SE(3) equivariane is important (in contrast to just equiivaraaa))) because the properties of chemical systems may change substantially when individual atoms translated relative each other.
10

---

Acknowledgments and Disclosure of Funding
OTU acknowledges funding from the Swiss National Science Foundation (Grant No. P2BSP2_188 MB acknowledges support by the Federal Ministry of Education and Research (BMBF) for BIFOLD (011S18037A). KRM was supported in part by the Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea Govern- ment (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 011$14013-EE, 01 GQ1115, 01GQ0850, 011S 18025A and 01 18037A; the German Research Foun- dation (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.
References
[1] Jörg Behler and Michele Parrinello. Generalized neural-network representation of high- dimensional potential-energy surfaces. Physical Review Letters, 98(14): 2007.
[2] Albert P. Bartók, Mike C. Payne, Risi Kondor, and Gábor Csányi Gaussian approximation potentials: The accuracy of quantum mechanics, without the electrons. Physical Review Letters, 104(13):136403, 2010.
[3] Justin S. Smith, Olexandr Isayev, and Adrian E. Roitberg. ANI-1: extensible neural network potential with DFT accuracy at force field computational cost. Chemical science, 8(4):3192- 3203, 2017.
[4] Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Müller, Machine learning of accurate energy-conserving molecular force fields. Science Advances,
2017.
[5] Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards exact molecular dynamics simulations with machine-learned force fields. Nature Communica- tions, (1))3887, 2018.
[6] Anders S. Christensen, Felix A. Faber, and O. Anatole von Lilienfeld. Operators in quantum machine learning: Response properties in chemical space. Journal of Physical Chemistry, 150 (6):064105, 2019.
[7] Simon Batzner, Tess E. Smidt, Lixin Sun, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, and Boris Kozinsky. -Eu3)--quivariant graph neural networks for data-efficient and accurate interatomic potentials. arXiv preprint arXi::2101.03164 2021.
[8] Oliver T. Unke, Stefan Chmiela, Huziel E. Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T. Schütt, Alexandre Tkatchenko, and Klaus-Robert Müller. Machine learning force fields. Chemi- cal Reviews, 121(16):10142-000
2021.
[9] John C. Snyder, Matthias Rupp, Katja Hansen, Klaus-Robert Müller, and Kieron Burke. Finding density functionals with machine learning. Physical Review Letters, 108(25):253002, 2012.
[10] Felix Brockherde, Leslie Vogt, Li Li, Mark Tuckerman, Kieron Burke, and Klaus- Müller. Bypassing the Kohn- Sham equations with machine learning. Nature Communications, 8:872, 2017.
[11] Kevin Ryczko, David Strubbe, and Isaac Tamblyn. Deep learning and density functional theory. arXiv preprint arXiv: /8/1. 1.08928, 2018
[12] Andrea Grisafi, David M. Wilkins, Benjamin A. R. Meyer, Alberto Fabrizio, Clemence Cormin- bocuf, and Michele Ceriotti. transferable machine-learnin model of the electron density. arXiv preprint FrXXiv:809.05344 2018.
[13] Ryo Nagai, Ryosuke Akashi, and Osamu Sugino. Completing density functional theory by machine learning hidden messages from molecules. npj Computational Materials, 6(1):1-8, 2020.
11

---

[14] Mihail Bogojeski, Leslie Yogt-Maranic Mark E. Tuckerman, Klaus-Robert Müller, and Kieron Burke. Quantum chemical accuracy from density functional approximations via machine learning. Nature communications, 11:5223, 2020.
[15] Matthias Rupp, Alexandre Tkatchenko, Klaus-Rober Müller, and O. Anatole von Lilienfeld. Fast and accurate modeling of molecular atomization energies with machine learning. Physical Review Letters, 108(5):05830, 2012.
[16] Grégoire Montavon, Matthias Rupp. Vivekanand Gobre, Alvaro Vazquez-Mayaaoittaa Katja Hansen, Alexandre Tkatchenko, Klaus-Robert Müller, and O. Anatole von Lilienfeld. Machine learning of molecular electronic properties in chemical compound space. New Journal of Physics, 159))095003,, 2013.
[17] Sandip De, Albert P Bartók, Gábor Csányi, and Michele Ceriotti. Comparing molecules and solids across structural and alchemical space. Physical Chemistry Chemical Physics, 18(20): 13754-13769, 2016.
[18) Kristof Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus-Robers Müller, and Alexandre Tkatchenko. Quantum-chent chemmcal insights from deep tensor neural networks. Nature Communi- cations, 8: 13890, 2017.
[19] Kristof T. Schütt, Pieter-Jan Kindermans, Huziel E. Sauceda, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing Systems, pages 991-1001,
[20] Kristof T. Schütt, Huziel E. Sauceda, Pieter-Jan Kindermans, Alexandre Tkatchenko, and Klaus-Robert Müller. SchNet deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148((4):241722, 2018.
[21] Wiktor Pronobis, Alexandre Tkatchenko, and Robert Müller. Many-body descriptors for predicting molecular properties with machine learning: Analysis of pairwise and three-body interactions in molecules. Journal of Chemical Theory and Computation, 14(())2991-300
2018.
[22] Michael Eickenberg, Georgios Exarchakis, Matthew Hirn, Stéphane Mallat, and Louis Thiry. Solid harmonic wavelet scattering for predictions of molecule properties. The Journal of Chemical Physics, 148(24):24132,, 2018.
[23] Oliver T. Unke and Markus Meuwly. PhysNet: A neural network for predicting energies, forces, dipole moments, and partial charges. Journal of Chemical Theory and Computation, 15(6): 3678-3693, 2019.
[24] Anatole von Lilienfeld, Klaus-Robers Müller, and Alexandre Tkatchenko. Exploring chemical compound space with quantum- based machine learning. Nature Reviews Chemistry, 4:347-358, 2020.
[25] Zhantao Chen, Nina Andrejevic, Tess E. Smidt, Zhiwei Ding, Qian Xu, Yen-Ting Chi, Quynh T Nguyen, Ahmet Alatas, Jing Kong, and Mingda Li. Direct prediction of phonon density of states with Euclidean neural networks. Advanced Science, page 2004214, 2021.
[26) Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger. Equivariant message passing for the prediction tensorial properties and molecular spectra. Proceedings of the 38th International Conference Machine Learning, PMLR 139:9377-9388, 2021.
[27] Oliver T. Unke, Stefan Chmiela, Michael Gastegger, Kristof T. Schütt, Huziel E. Sauceda, and Klaus-Robert Müllerr. SpookyNet: Learning force fields with electronic degrees freedom and nonlocal effects. arXiv preprint arXiv:210S.00300 2021.
[28] John A. Keith, Valentin Aassilev-Galindo, Bingqing Cheng, Stefan Chmiela, Michael Gastegger, Klaus-Robert Müller, and Alexandre Tkatchenko. Combining machine learning and com- putational chemistry for predictive insights into chemical systems. Chemical Reviews, 121 (16):9816-987)) 2021. URI. https: : //pbbs.u....//// acs. wwgg/do//bbaaeeoooooooooooooooot awtp//bbbb/...o// 1021/ace chemrev. 1c00107.
12

---

[29] Michael Gastegger, Kristof T. Schütt. and Klaus-Robert Müller. Machine learning of solvent effects on molecular spectra and reactions. Chemical Science, 12(34):11473-112 2021.
[30] Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and Thomas F. Miller III. OrbNet: Deep learning for quantum chemistry using symmetry-adapter atomic- orbital features. The Journal of Chemical Physics, 153(12);1241) 2020.
[31] Giuseppe Carleo and Matthias Troyer. Solving the quantum many-body problem with artificial neural networks. Science, 355(6325)x602-6666, 2017.
[32] Ganesh Hegde and Chris Bowen. Machine-learned approximations to density functional theory Hamiltonians. Scientific Reports, 7:42669, 2017.
[33] Kristof T. Schütt, Michael Gastegger, Alexandre Tkatchenko, Klaus-Robert Müller, and Rein- hard J. Maurer. Unifying machine learning and quantum chemistry with a deep neural network for molecular wavefunctions. Nature Communications, 10:5024, 2019.
[34] He Li, Zun Wang, Nianlong Zou, Meng Ye, Wenhui Duan, and Yong Xu. Deep neural network representation of density functional theory Hamiltonian. arXiv preprint NXXi:22144007866 2021.
[35] Grégoire Montavon, Katja Hansen, Siamac Fazli, Matthias Rupp. Franziska Biegler, Andreas Ziche, Alexandre Tkatchenko, Anatole Lilienfeld, and Klaus-Rober Müller. Learning invariant representations of molecules for atomization energy prediction. Advances in neural information processing systems, 25:440 448, 2012.
[36] Michael Gastegger, Adam McSloy, M Luya, Kristof T. Schütt, and Reinhard J. Maurer. deep neural network for molecular wave functions in quasi-atomic minimal basis representation. The Journal of Chemical Physics, 153(4):044123, 2020.
[37] Taco S. Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning, pages 2990-2999 PMLR, 2016.
[38] Diego Marcos, Michele Volpi, Nikos Komodakis, and Devis Tuia. Rotation equivariant vector field networks. In Proceedings of the IEEE International Conference on Computer Vision, pages 5048-5057, 2017.
[39] Taco S Cohen, Mario Geiger, Jonas Köhler, and Max Welling, Spherical CNNs. arXiv preprint 1rXiv:1801.10130 2018
[40] Risi Kondor, Hy Truong Son, Horace Pan, Brandon Anderson, and Shubhendu Trivedi. Covari- compositional networks for learning graphs. arXiv preprint arXiv:1801.021 2018.
[41] Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3D steer- able CNNs: Learning rotationally equivariant features in volumetric data. arXiv preprint arXiv: 2018.
[42] Nathaniel Thomas, Tess E. Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation and translation-equivarian neural networks for 3D point clouds. arXiv preprint arXiv:1802.082/9. 2018
[43] Andrea Grisafi, David M. Wilkins, Gábor Csányi. and Michele Ceriotti. conmetry-adapter machine learning for tensorial properties of atomistic systems. Physical Review Letters, 120: 036002, Jan 2018. doi: 10.1 103//hyssevl..et....mm wtt.:/lee........eeee 120.036002.
[44] Truong Son Hy, Shubhendu Trivedi, Horace Pan, Brandon M. Anderson, and Risi Kondor. Predicting molecular properties with covariant compositional networks. The Journal of Chemical Physics, 148(24):241745, 2018.
[45] Taco S. Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant CNNs on homogeneous spaces. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
13

---

[46] Nicolas Keriven and Gabriel Peyré. Universal invariant and equivariant graph neural networks. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[47] Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling. Gauge equivariant convolutional networks and the icosahedral CNN. In International Conference on Machine Learning, pages 1321-1330. PMLR. 2019.
[48] Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. transforrassssmr:s:: 3D roto-translation equivariant attention networks. arXiv preprint arXiv:2006. 10503, 2020.
[49] Jigyasa Nigam, Michael Willatt, and Michele Ceriotti. Equivariant representations for molecular Hamiltonians and N-center atomic-scale properties. arXiv preprint arXiv:2109.12088 2021.
[50] Martin Schechter. Operator methods in quantum mechanics. Flsevier, 2012.
[51] Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In International Conference Machine Learning, pages 1263-1272. 1272. PMLR, 2017.
[52] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). arXiv preprint arXiv:1606.08415 2016.
[53] Stefan Elfwing. Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Netw., 107:3-11, 2018.
[54) Prajit Ramachandran, Barret Zoph, and Quoc V Le.. Searching for activation functions. arXiv preprint arXiv: 1710.05941, 2017.
[55] D. A. Varshalovich, A. N. Moskalev, and V. K. Khersonskii. Quantum theory of angular momentum. World Scientific. 1988.
[56] Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch-gordan nets: fully fourier space spherical convolutional neural network. arXiv preprint arXiv:1806.0925:: 2018.
[57] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. European conference on computer vision, pages 630-645. Springer, 2016.
[58] Stefan Goedecker and GE Seuserza. Linear scaling electronic strueture methods chemistry and physics. Computing in Science & Engineering, 5(4):14-21, 2003.
[59] Fabio Anselmi, Lorenzo Rosasco, and Tomaso Poggio. On invariance and selectivity in representation learning. Information and Inference: A Journal of the IMA, 15) 2016. doi: 10.199//mmiia/iaw0000
[60] Mikio L. Braun, Joachim M. Buhmann, and Klaus-Rober Müller. On relevant dimensions in kernel feature spaces. The Journal of Machine Learning Research,
2008
[61] Huziel E. Sauceda, Luis E. Gälvez-González, Stefan Chmiela, Lauro Oliver Paz-Borbón, Klaus- Robert Müller, and Alexandre Tkatchenko. BIGDML: Towards exact machine learning force fields for materials. arXiv preprint arXiv:2106.04222 2021.
14


Reviews:
[{"review_id": "cUOjBJ253AK", "reviewer": "Reviewer_ZUAT", "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.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Main paper**\n\n1. Related works: The following highly related references should be cited since they are part of very related threads of work:\n\n   [1] \"SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks\" by Fabian B. Fuchs and Daniel E. Worrall and Volker Fischer and Max Welling\n\n   [2] \"OrbNet: Deep learning for quantum chemistry using symmetry-adapted atomic-orbital features\" by Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and  Thomas F. Miller\n\n   At the very least some comments on how the proposed method differs from these approaches should be included. Reference [1] presents an equivariant network for property prediction and reference [2] includes a 1000x speedup over DFT by leveraging symmetry-aware features for orbital data. \n    \n2. The abstract, Section 1 and Section 3 (theory development) all mention that the underlying geometries of molecules need to to rotation-equivariant to capture the various rotational symmetries without losing the information about the rotation. This suggests that an SO(3)-equivariant network would be sufficient. However, there is an abrupt jump to SE(3)-equivariance in Section 4 (Deep learning architecture). What is the reasoning behind that? Note that SE(3) expresses $[R | t]$ and SO(3) expresses $R$, where $R$ is the rotation matrix and $t$ is the translation vector in $\\mathbb{R}^3$. We can ignore the homogeneous formulation's scale factor for now. How/When do translations need to be preserved and predicted? How do translations even manifest in a representation of a molecule? Do individual bonds within a molecule elongate or contract (in response to temperature or pressure, for example) and if so, then how is that data fed into the network? Or is SE(3)-equivariance being used as a superset of SO(3)? Or is it the case that properties of groups of atoms within a molecule are of interest, and the bonds attaching these groups to rest of the molecule can change length, thereby encoding some useful information in the translation components? \n\n3. Figure 1 (D). The HF-level function, the correction (amplified) and the CC-level function are presented. By eye, the HF-level function and CC-level function look the same to me. Is it correct to expect 'CC-level function = HF-level function + correction'? If not, how should we interpret Figure 1-D? A clarifying caption would be helpful. Alternatively, the differences in HF-level function and CC-level function could be highlighted. \n\n4. Evaluation [Table 1]: The paper clearly states that the approach is built on PhySNet, yet the proposed method (PhiSNet) is not compared to PhySNet. It is instead compared only to SchNOrb. The results presented in Table 1 for SchNOrb do correlate with the results in the supplementary material of the SchNOrb publication, but only partially. So I presume authors ran additional experiments. Great. \n\n5. Evaluation [Table 1]: It would be good to have a comparison to PhySNet here as it would allow comparison to a whole host of methods. OrbNet evaluates a number of methods against DFT including PhySNet and at a 1000x speedup on DFT. It would be good to get a rough estimate of speedup of the proposed method wrt DFT as a general guideline. \n\n6. Evaluation: Several other works in this area present results for property prediction (QM9, QM7, DrugBank etc). It is curious to not find these included in the paper as the paper explicitly claims to provide good transferability to property prediction. I do note that some results wrt HF to CCSD transfer learning are provided in the supplementary material but as these measure very different quantities from those in related works, it's unclear how to interpret the results. (Apologies, I am still getting my head round the level of theory business).\n\n7. Section 4.1 is missing a sentence or two in way of introduction. It's quite unclear whether the layers defined in Section 4.1 refer to PhiSNet (this paper) or PhySNet (a previous work upon which this paper builds). I am assuming they refer to PhiSNet (this paper) based on the cited references inline and line 144-145, but just staing this upfront would be easier for the reader. \n\n8. Section 4.1 [Interaction blocks]: This is the first point where positions $\\mathbf{r}$ of the atoms and the inter-atom distances $\\mathbf{r}_{ij}$ are used. So coming back to my earlier point about SO(3) vs SE(3), there could be an argument here for needing SE(3) to encode the bond lengths between the atoms. Much as before, aren't these known and fixed quantities or do they vary substantially, and if they do, do datasets actually have that variation? Else, why not just assume some units of 'one C-H bond away'? (Of course, things get complicated when the bonds are not unique, i.e. valence allows either double or triple bond between C and N etc). But this can largely by disambiguated by using valence constraints.\n\n9. In Eq (16), it might be good to add some suffix or words to clarify that the $\\mathbf{f}$ in the definition of $\\mathbf{c} = residual(\\mathbf{f})$ and $\\mathbf{n} = residual(\\mathbf{f})$ operates on different atoms within the expressions itself. \n\n10. Network size: A comparison of the number of parameters is provided wrt SchNOrb in the supplementary information (17M for PhiSNet and 34M for SchNOrb). How many parameters does PhySNet have, seeing as that's the network this paper is most closely based on?\n\n11. In the contributions section, the authors say the proposed method improves over state of the art by two orders of magnitude. It would be helpful to clarify whether this is the case when comparing to equivariant methods or non-equivariant methods? I presume it is the latter. Going back to the point on evaluation, as the main evaluation only compares against one other method, it would be good to add a comparison to PhySNet and perhaps also, where possible a partial/alternative comparison to other methods, such as the ones OrbNet uses for evaluation with a accuracy-effort tradeoff to highlight the benefits of the current approach (end-end trainable). Ballpark number could also suffice for such evaluation. \n\n12. There is a slight contradiction in the contributions section (Lines 49-53): Contribution 1 suggests the need for SE(3) whereas contribution 2 suggests only transformation under rotations are required to be preserved. \n\n13. Evaluation: PhySNet seems to improve with ensembles. Does it help with PhiSNet? If it's not needed, that could be a good thing to highlight. \n\n14. Would the method scale to work on proteins (larger molecules but more repeated structures) as it does for small molecules. Most of the examples included in the paper are for small molecules. There is a passing comment about computational complexity scaling with the number of atoms in Section 5, but my question is more around the repeated structures aspect.  \n\n\n**Supplementary information:**\n\n1. Thank you! The background information on quantum chemistry (Section A) leading up to the embedding details (Section C) along with the network architecture diagrams and Hamiltonian matrix construstion with minimal basis sets (Section D and Figure S5) are the clearest I have seen in this field. To make this even better, I would suggest using terminology from both ML/CV and Physics where relevant. Group equivariance is heavily studied in both physics and computer vision, for example, but different terminology is used. Perhaps for electronic dissemination, it would make sense to have a long version of this paper that interweaves the main paper and the supplementary info so as to give the reader an uninterrupted (if longer) and lucid flow of reading.\n\n2. Why do the batch sizes used for different molecules (Table S4) vary so much? Is it in direct proportion to the number of atoms per molecule? Do you not have to change the learning rate or the number of epochs to adapt? In other words, are the number of atoms within a batch for the different molecules roughly constant? \n\n**Notes**\n\nWriting style: At several points while reading the paper, I had questions along the lines of 'What do the different levels that the authors refer to mean?' or 'What is **S** matrix'? All these questions are answered nicely in the supplementary information, but then it's unclear what order to read the paper and supplementary info in? Given the topic of the paper, quite a bit of background from quantum chemistry/physics is required, I think it's fine for there to be a substantial supplementary information for those who may need it, but it would be nicer for all readers if the supplementary info was referred to in the main paper whenever the background from supplementary info is required for understanding the content in the main paper. For example, 'We predict **S** (see Supp. info Section .. for explanation of **S**)'. \n\n**Overall assessment**\n\nThis paper presents a very good motivation of DL methods for quantum chemistry,specifically learning wavefunctions from electronic orbital information. The method presented is novel even though it builds on previous work in the area. The method addresses a key challenge and demonstrates significant improvement in accuracy of wavefunction and energy prediction with an end-end approach. However, some comparisons and comments on differences wrt key existing works would be helpful in understanding the importance of this work. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, to a reasonable extent. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "SF_WqKxlYD1", "reviewer": "Reviewer_5CND", "summary": "\nThis 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.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nThe method is a clear improvement over SchNOrb: it incorporates SO3 symmetries in its core and it shows significantly better approximation accuracy, better data efficiency, and reduced model size. The paper is written in a clear and concise way, though the demonstrations are somewhat lacking compared to what I'd have expected for the NeurIPS conference. The possible combinations of the building blocks all make sense, but I'd have liked to see an ablation study or additional arguments in the supplement for the particular combinations of architectural choices. Importantly, the authors do not discuss how/if this architecture could generalize to predict results on a system not seen in the training set. They also don't show performance on total energy when using the previously established 1K MD17 training datasets, so it's unclear how much this contribution adds to practical applications. Although the prediction of DFT level wavefunctions for new molecular conformations is potentially interesting, I was more excited about the possibility to correct these wavefunction to an effective CCSD(T) wavefunction. Unfortunately, the example with a single illustration of a single water molecule is not very convincing that this architecture can provide useful practical tool in the near future.\n\nEdit: following the additional work that demonstrates the performance gains on DFT-level calculations, I've updated my overall score to 6.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors discuss the scaling limitations of their work as they related to the cubic scaling with the number of atoms; there are no potential negative societal impacts of this work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}, {"review_id": "SB6ke_xOcB", "reviewer": "Reviewer_1y2Z", "summary": "The paper proposes PhiSNet, an SE3 equivariant NN whose inputs are distance vectors between atoms of a given molecule with nuclear charges, and outputs a prediction of the entries of the Hamiltonian matrix that is used to solve the Schrodinger equation to find the electronic wavefunction of the molecule. It is claimed that training PhiSNet on a molecular dataset reaches state of the art prediction accuracies on various molecules.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Without a background in quantum chemistry, it is difficult for me to assess the significance of the results in the paper, hence I will focus on assessing the proposed methodology.\n\nThe paper attributes the success of PhiSNet over the baseline SchNOrb to its SE3 equivariance, but it is unclear how much of the improvements over the baseline are due to its equivariance and how much is from design choices of the network. Of course these two come hand in hand, so it's difficult to separate their contributions to improving prediction accuracy, but a more detailed description of the baseline method along with a description of the difference between this and PhiSNet would be desirable. In addition, given the claims of SOTA, it would be useful to have ablations to explain the various design choices of PhiSNet, e.g. parameterisation of Interaction blocks (eq 13), using a different set of parameters for the central point features and the neighbouring point features (eq 15, 16), or using the other neighbours of atom i for computing the pair interaction features between atoms i & j (eq 18). In summary, a better breakdown of what design choices have led to how much improvement in prediction accuracy would be very helpful for explaining the success of PhiSNet and will encourage follow up work to make further improvements.\n\nThe paper introduces PhiSNet as a novel SE3 equivariant architecture for point cloud data, but it's not clear where the novelty is: it's not clear whether the authors are claiming some of the building blocks in Section 4.1 are novel or that the combination of these existing building blocks is novel. The building blocks described in Section 4.1 are lacking in citations e.g. Tensor product contractions have been used frequently in the SE3 equivariance literature, so relevant work should be cited e.g. Clebsch-Gordan Nets [Kondor et al, 2018].\n\nI appreciate the paper's attempt to make the methodology clear, with details on notation and the methodology, but there remain some points of confusion \n- In line 207 it is said 'tilde{x} serves as input to the next module in the chain' but I don't see tilde{x} appearing subsequently.\n- In line 209, what is the sum across?\n- I find the use of Tensor product expansions to construct matrix predictions interesting, but the explanation in Section 4.2 is not so clear (e.g. what is used as 'a' in eq 15?). It would be helpful to include Figure S5 as part of the main paper for this.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "n/a", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "Iyq4QjCEDb", "reviewer": "Reviewer_jstg", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In general the paper is well-written and well-structured, with about the right amount of space devoted to each part (e.g. the background in section 3 is pretty minimal, but I think this is the correct decision here, and I appreciated the longer-form background provided in the supplement). As with almost all papers, it would benefit from another proofreading pass to catch minor errors (e.g. \"calculations and led to\" on line 25 should be \"calculations, and have led to\"). The detailed description of PhiSNet and its component parts is clear and well laid out. Figure 1 is excellent, I found it extremely helpful in clarifying the problem setting and the architecture.\n\nThere are, as the authors note in section 2, a number of existing SE(3)-equivariant machine learning methods. In section 4.1, it is unclear which of the building blocks the authors are claiming as novel and which are taken from (or heavily based on) these earlier works. They say \"we describe novel general-purpose operations\" on line 118-119, but clearly not everything that follows is novel (e.g. the chosen activation function and the linear layers). I would appreciate more clarity on this.\n\nA point about terminology: I'm unclear exactly what the authors mean by \"irrep\". Strictly an irrep is a representation, i.e. a homomorphism from SO(3) to GL(N) (which satisfies some additional properties to be \"irreducible\"), but the authors also seem to refer to feature vectors as \"irreps\" (e.g. in section 4 on line 152). Is the idea that SO(3) acts on these feature vectors by an irrep?\n\nThe results appear impressive, though since I am unfamiliar with this application domain I am taking it on good faith that the results in Table 1 represent a fair comparison against a state-of-the-art competitor. The results are the most compelling part of the paper, and they are the main thing that convinces me of the value of the model architecture described in section 4.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately address some limitations of their work in section 5. They do not discuss potential societal impact - it is not clear to me that such a discussion is necessary for this paper.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors present an SE(3) equivariant architecture for predicting symmetric Hamiltonian matrices, which can then be used to derive physical properties of molecules. This architecture provides a dramatic improvement in accuracy over baseline methods. The majority of reviewers agreed it was an excellent paper deserving acceptance. The one holdout reviewer had concerns about the scale of the experiments and recommended that results be presented for the entirety of the MD17 dataset. However I am convinced by the author response that this would not be feasible with the computational resources at their disposal, and I am willing to cut them some slack on this given the quality of results on the systems presented. Therefore I recommend that the paper be accepted.
Decision: Accept (Poster)