Schütt Et Al. - 2021 - Equivariant Message Passing For The Prediction of

Download as pdf or txt
Download as pdf or txt
You are on page 1of 13

Equivariant message passing for the prediction of tensorial properties and

molecular spectra
Kristof T. Schütt,1, 2 Oliver T. Unke,1, 2 and Michael Gastegger1, 3
1)
Machine Learning Group, Technische Universität Berlin, 10587 Berlin, Germany
2)
Berlin Institute for the Foundations of Learning and Data, 10587 Berlin,
Germanya)
3)
BASLEARN – TU Berlin/BASF Joint Lab for Machine Learning, 10587 Berlin,
Germany
(Dated: 8 June 2021)
Message passing neural networks have become a method of choice for learning on graphs, in particular the
prediction of chemical properties and the acceleration of molecular dynamics studies. While they readily scale
to large training data sets, previous approaches have proven to be less data efficient than kernel methods. We
identify limitations of invariant representations as a major reason and extend the message passing formulation
to rotationally equivariant representations. On this basis, we propose the polarizable atom interaction neural
arXiv:2102.03150v4 [cs.LG] 7 Jun 2021

network (PaiNN) and improve on common molecule benchmarks over previous networks, while reducing
model size and inference time. We leverage the equivariant atomwise representations obtained by PaiNN for
the prediction of tensorial properties. Finally, we apply this to the simulation of molecular spectra, achieving
speedups of 4-5 orders of magnitude compared to the electronic structure reference.

I. INTRODUCTION the representation of nodes (atoms) remains rotation-


ally invariant. While equivariant convolutions have been
Studying dynamics of chemical systems allows insight successfully applied in computer vision16–18 , previous ap-
into processes such as reactions or the folding of proteins, proaches to molecular prediction19,20 have not reached
and constitutes a fundamental challenge in computational the accuracy of their rotationally invariant counterparts.
chemistry. Since the motion of atoms is governed by the In this work, we propose rotationally equivariant mes-
laws of quantum mechanics, accurate ab initio molecu- sage passing and the polarizable atom interaction neural
lar dynamics (MD) simulations may require solving the network ( PaiNN) architecture as one instance of it. We
Schrödinger equation for millions of time steps. While examine the limited capability of rotation-invariant repre-
the exact solution is infeasible to compute for all but sentations to propagate directional information and show
the smallest systems, even fast approximations such as that equivariant representations do not suffer from this
density functional theory quickly become prohibitive for issue. PaiNN outperforms invariant message passing net-
large systems and the prediction of accurate spectra. works on common molecular benchmarks and performs
Recently, machine learning potentials1–3 have gained at small sample sizes on par with kernel methods that
popularity for studying systems ranging from small have been deemed to be more data-efficient than neural
molecules at high levels of theory4,5 to systems with networks. Beyond that, the rotationally equivariant rep-
thousands or millions of atoms6–8 . In particular, message- resentation of PaiNN enables the prediction of tensorial
passing neural networks9 (MPNNs) yield accurate predic- properties which we apply to the ring-polymer MD simu-
tions for chemical properties across chemical compound lation of infrared and Raman spectra. By an acceleration
space and can handle large amounts of training data. of 4-5 orders of magnitude, PaiNN makes the simulation
Albeit MPNNs have significantly increased in accuracy of these spectra feasible, reducing the runtime in one case
over the years (as well as in computational cost), kernel from projected 25 years to one hour.
methods with manually crafted features10–12 have still
proven to perform better when only small training sets
are available. II. RELATED WORK
While molecules are often represented as graphs, they
are in fact interacting particles in a continuous 3d Behler and Parrinello 21 introduced neural network po-
space. Consequently, SchNet13 modeled message passes as tentials taking atom-centered symmetry functions based
continuous-filter convolutions over that space, albeit with on distances and angles as features. Graph neural net-
rotationally invariant filters. As Miller et al. 14 pointed works22 for molecular graphs23,24 and 3d geometries9,25–28
out, this leads to a loss of relevant directional, equiv- do not require such manually crafted features, but learn
ariant information. Klicpera, Groß, and Günnemann 15 embeddings of atom types and use graph convolutions or,
have introduced directional message-passing, the angu- more general, message passing9 , to model atom interac-
lar information here is restricted to the messages while tions based on interatomic distances. Klicpera, Groß, and
Günnemann 15 introduced directional message passing by
including additional angular information in the message.
a) Electronic mail: [email protected] Steerable CNNs16,29 allow to build equivariant filter
2

banks according to known symmetries of associated fea- conventional CNNs can be cast in this general framework.
ture types. While these approaches work on grids, archi- Rotational invariance of the representation can be ensured
tectures such as Tensor Field Networks19 , Cormorant20 by choosing rotationally invariant message and update
and NequIP30 use equivariant convolutions based on functions, which by definition need to fulfill
spherical harmonics (SH) and Clebsch-Gordon (CG) trans-
forms for point clouds. Kondor and Trivedi 31 have de- f (~x) = f (R ~x), (3)
scribed a general framework on equivariance and convolu-
tions in neural networks focusing on irreducible representa- for any rotation matrix R ∈ R3×3 .
tions. In contrast, PaiNN models equivariant interactions
in Cartesian space which is conceptually simpler and does
not require tensor contractions with CG coefficients. A B. Building equivariant MPNNs
similar approach was proposed by Jing et al. 32 with the
GVP-GNN. Both approaches are designed for distinct To obtain more expressive representations of local envi-
applications resulting in crucial differences in the design ronments, neurons do not have to be scalar, but can be
of the neural network architectures, most notably the geometric objects such as vectors and tensors16,19,20,33 .
message functions. The GVP-GNN has been designed for For the purpose of this work, we restrict ourselves to scalar
single-point protein sequence predictions, e.g., allowing and vectorial representations sti and ~vit , respectively, such
the use of non-smooth components. In contrast, PaiNN that a corresponding message pass can be written as
is designed for simulations with millions of inference steps, X
requiring a fast message function and a smoothly differ- ~ v,t+1
m i = M~ t (st , st , ~vt , ~vt , ~rij )
i j i j (4)
entiable model. j∈N (i)

Message and update functions for scalar features can be


III. EQUIVARIANT MESSAGE PASSING defined analogously. Rotational equivariance of vector
features ~vit+1 can be ensured by employing rotationally
a. Notation. To clearly distinguish between the two equivariant functions U~ t and M
~ t , fulfilling
concepts, we write feature vectors as x ∈ RF x1 and vectors
in 3d coordinate space as ~r ∈ R1x3 . Vectorial features will R ~f (~x) = ~f (R ~x) (5)
be written as ~x ∈ RF x3 . Norms k · k, scalar products h·, ·i
and tensor products ⊗ are calculated along the spatial for any rotation matrix R ∈ R3×3 , where the matrix-
dimension. All other operations are calculated along the vector product is applied over the spatial dimension. This
feature dimension, if not stated otherwise. We write the constitutes essentially a linearity constraint for directional
Hadamard product as ◦. information. Therefore, any nonlinearities of the model
need to be applied to scalar features. In particular, equiv-
ariant MPNNs may use the following operations:
A. Message passing for 3d-embedded graphs
• Any (nonlinear) function of scalars: f (s)
MPNNs build complex representations of nodes within • Scaling of vectors: s ◦ ~v
their local neighborhood in a graph through a repeated ex-
change of messages followed by updates of node features. • Linear combinations of equivariant vectors: W~v
Here, we consider graphs embbedded in 3d Euclidean
space, where edges are specified by the relative positions • Scalar products: s = k~vk2 , s = h~v1 , ~v2 i
~rij = ~rj − ~ri of nodes i, j within a local neighborhood • Vector products: ~v1 × ~v2
N (i) = {j | k~rij k ≤ rcut }. In a chemistry context, this
agrees with the fact that a large part of the energy vari- Thus, directional information is preserved during message
ations can be attributed to local interactions, often con- passing while invariant representations can be recovered
ceptualized as bond lengths and angles. Thus, only nodes from scalar products to obtain rotationally invariant pre-
within that range can interact directly, so the number of dictions.
messages does not scale quadratically with the number
of the nodes. A general MPNN for embedded graphs can
be written as C. Limits of rotationally invariant representations
X
mt+1
i = Mt (sti , stj , ~rij ) (1) In the following, we examine the expressiveness of in-
j∈N (i) variant and equivariant representations at the example
st+1 = Ut sti , mt+1

i i . (2) of simple 2d molecular structures. To ensure rotational
invariance of a predicted property, the message function
with the update and message functions Ut and Mt , re- is often restricted to depend only on rotationally invari-
spectively9 . Many neural network potentials and even ant inputs such as distances25,28 or angles15,34 . While a
3

TABLE I: Comparison of expressiveness and computa-


tional complexity of distances, angles and directions for a
simple message function of the oxygen atom of a water
molecule.

Features Distances Angles Directions


FIG. 1: Illustration of message passing using angles and
directions for two structures. All edges within the cutoff
H2 O range (dashed lines) have equal length. The representa-
tions of the blue and red node are the same using angles
(left), while directions allow to distinguish both structures
(right).
X X X X ~
rij
Message M at k~
rij k αjik
k~
rij k
atom i j∈Ni j∈Ni k∈Ni j∈Ni

Scaling with with the cutoff chosen such that only neighboring atoms
O(|N |) O(|N |2 ) O(|N |)
neighbors are within range. We observe that for the two arrange-
Resolve change yes no no
ments the angles are equal as well (Fig. 1, left). Therefore,
of k~
r1j k they are indistinguishable for invariant message passing
Resolve change no yes yes
with distances and angles. In contrast, the equivariant
of α213 representations differ in sign of their components (Fig. 1,
right). When contracting them to invariant representa-
tions, as in the previous example, this information is lost.
However, we may instead retain equivariance in the repre-
single pass is limited to interactions within a local en- sentation and design a message function M ~ t that does not
vironment, successive message passes are supposed to only propagate directions to neighboring atoms, but those
construct more complex representations and propagate lo- of equivariant representations as well. This enables the
cal information beyond the neighborhood. This raises the efficient propagation of directional information by scaling
question whether rotationally invariant representations of linearly with the number of neighbors and keeping the
atomic environments hi are sufficient here. required cutoff small.
Tab. I compares three simplified message functions
Note that neither many-body representations, using an-
for the example of a water molecule. Using a distance-
gles, dihedral angles etc., nor equivariant representations
based message function, the representation of the atom
corresponding to a multipole expansion are complete. It
1 (oxygen) is able to resolve changing bond lengths to
has been shown that even when including up to 4-body
atoms 2 and 3 (hydrogens), however it is not sensitive
invariants, there are structures with as few as eight atoms
to changes of the bond angle. On the other hand, using
which cannot be distinguished35 and that a many-body
the angle directly as part of the message function can not
expansion up to n-body contributions is necessary to guar-
resolve the distances. Therefore, a combination of dis-
antee convergence for a system consisting of n atoms36 .
tances and angles is required to obtain a more expressive
The same holds for multipole expansions, where scalars
message function15 . Unfortunately, including angles in
and vectors correspond to the 0th and 1st order. Thus,
the messages scales O(|N |2 ) with the number of neigh-
even spherical harmonics expansions with less than infi-
bors. Alternatively, directions to neighboring atoms may
~ t (~rij ) = ~rij /k~rij k. Employing the nite degree are unable to represent arbitrary equivariant
be used as messages M n-body functions. Instead, a practically sufficient and
update function Ut (m) = kmk2 , this is related to angles computationally efficient approach for the problem at
as follows: hand is desirable.
2 In the following, we propose a neural network archi-
N X  ~rij  X N XN
X ~rij ~rik tecture that takes advantage of these properties. It over-
= , = cos αjik .
k~rij k k~rij k k~rik k comes the limitations of invariant representations dis-
j=1 j,k j=1 k=1
cussed above, which we will demonstrate at the example
of an organometallic compound in Section V C 2.
Thus, using equivariant messages, the runtime complexity
remains O(|N |) while angular information can be resolved.
Note that this update function contracts the equivariant
messages to a rotationally invariant representation. IV. POLARIZABLE ATOM INTERACTION NEURAL
Beyond the computational benefits, equivariant rep- NETWORK (PAINN)
resentations allow to propagate directional information
beyond the neighborhood which is not possible in the in- The potential energy surface E(Z1 , . . . , ZN , ~r1 , . . . , ~rN ),
variant case. Fig. 1 illustrates this at a minimal example, with nuclear charges Zi ∈ N and atom positions ~ri ∈ R3
where four atoms are arranged in an equidistant chain exhibits certain symmetries. Generally, they include in-
4

FIG. 2: The architecture of PaiNN with the full architecture (a) as well as the message (b) and update blocks (c) of
the equivariant message passing. In all experiments, we use 128 features for si and ~vi throughout the architecture.
Other layer sizes are annotated in grey.

variance of the energy towards the permutation of atom Next, we define message and update functions as in-
indices i, as well as rotations and translations of the troduced in Sec. III. We use a residual structure of inter-
molecule. A neural network potential should encode these changing message and update blocks (Fig. 2a), resulting
constraints to ensure the symmetries of the predicted in coupled scalar and vectorial representations. For the
energy surface and increase data efficiency. A common residual of the scalar message function, we adopt the
inductive bias of the neural network potential is a de- feature-wise, continuous-filter convolutions introduced by
composition of the energy into atomwise contributions Schütt et al. 13
PN
E = i=1 (si ), where an output network  predicts en- ∆sm
i = (φs (s) ∗ Ws )i (7)
ergy contributions from atoms embedded within their X
chemical environment, represented by si ∈ RF 21,37 . = φs (sj ) ◦ Ws (k~rij k),
While properties of chemical compounds may be such j
rotationally invariant scalars, they can also be equivariant where φs consists of atomwise layers as shown in Fig. 2b.
tensorial properties, e.g. the multipole expansion of the The roationally-invariant filters Ws are linear combi-
electron density nations of radial basis functions sin( rnπ k~rij k)/k~rij k as
cut

~ >~r + ~r> Q~r + . . . ,


n(~r) = q + µ (6) proposed by Klicpera, Groß, and Günnemann 15 with
1 ≤ n ≤ 20. Additionally, we apply a cosine cutoff to the
with charge q, dipole µ ~ , quadrupole Q and so on. Simi- filters39 .
larly, one can interpret invariant and equivariant atomwise Analogously, we use continuous-filter convolutions for
representations as local charges and dipole moments of the residual of the equivariant message function
atom-centered multipole expansions38 . We leverage this ∆~vim =
X
~vj ◦ φvv (sj ) ◦ Wvv (k~rij k) (8)
in Section IV B to predict tensorial molecular proper- j
ties. We coin our proposed architecture polarizable atom
interaction neural network (PaiNN).
X
0 ~rij
+ φvs (sj ) ◦ Wvs (k~rij k) ,
j
k~rij k

A. Representation where the first term is a convolution


P of an invariant filter
with scaled, equivariant features j ~vj ◦φs (sj ), related to
the gating proposed by Weiler et al. 29 as an equivariant
The inputs to PaiNN are the nuclear charges Zi ∈
nonlinearity. This propagates directional information
N and positions ~ri ∈ R3 for each atom i. Similar to
obtained in previous message passes to neighboring atoms.
previous approaches, the invariant atom representations
The second term is a convolution of invariant features
are initialized to learned embeddings of the atom type
with an equivariant filter. This can be derived as the
s0i = aZi ∈ RF ×1 . We keep the number of features
gradient of an invariant filter
F constant throughout the network. The equivariant
representations are set to ~vi0 = ~0 ∈ RF ×3 , since there is 0 ~rij
∇Wvs (k~rij k) = Wvs (k~rij k) .
no directional information available initially. k~rij k
5

PaiNN keeps the pairwise message functions compara-


tively lightweight. Instead complex transformations are
restricted to the atomwise update function and tensor
output network.
Depending on the chemical property of interest, one
may want to include additional constraints or replace ~νk,m
with the atom positions to include the global molecular
geometry. In the following, we demonstrate how this
can be employed for two tensorial properties that are
important for the simulation of molecular spectra.
The molecular dipole moment µ ~ is the response of the
molecular energy to an electric field ∇F~ E and, at the
FIG. 3: Gated equivariant block. same time, the first moment of the electron density (see
Eq. 6). It is often predicted using latent atomic charges40 :
0
Since Wvs (k~rij k) is another invariant filter, we can model N
X
it directly without taking the derivative. Furthermore, µ
~= qatom (si )~
ri . (12)
we use a shared network φ to perform the transform of all i=1
message functions and split the features afterwards (see
Fig. 2b). This assumes the center of mass at ~r = ~0 for brevity.
After the features-wise message blocks, the update While this only leverages invariant representations, we
blocks are applied atomwise across features. The residual can take advantage of the equivariant features of PaiNN.
of the scalar update function is given by In this picture, the molecular dipole is constructed from
polarized atoms, i.e. both atomic dipoles and charges41 ,
∆sui = ass (si , kV~vi k) (9) yielding
+ asv (si , kV~vi k) hU~vi , V~vi i.
N
X
Again, we use a scaling functions computed by a shared µ
~= µ
~ atom (~vi ) + qatom (si )~
ri . (13)
network a(si , kV~vi k) as nonlinearity. In this case, the i=1
norm of a linear combination of features is also used to ob-
tain the scaling, thus, coupling the scalar representations In terms of Eq. 11, this is a simple case with M = 1, R = 2,
with contracted equivariant features. In a second term, where the first term predicts the latent local dipoles and
we use the scalar product of two linear combinations of the second term computes the dipole generated by the
equivariant features. Similarly, we define the residual for local charges.
the equivariant features As an example with order M = 2, we consider the
polarizability tensor, which describes the response of the
∆~viu = avv (si , kV~vi k) U~vi , (10) molecular dipole to an electric field JF~ (~
µ) = HF~ (E). We
construct polarizability tensors using
which is again a nonlinear scaling of linearly combined
equivariant features. N
X
α= α0 (si ) I3 + ~ν (~vi ) ⊗ r~i + r~i ⊗ ~ν (~vi ), (14)
i=1
B. Prediction of tensorial properties
where the first term of the sum models isotropic, atomwise
PaiNN is able to predict scalar properties using atom- polarizabilities. The other two terms add the anisotropic
wise output layers of rotationally invariant representations components of the polarizability tensor. Similar to the
si , in the same manner as previous approaches25,27 . Be- charges in Eq. 13, here the atom positions are used to
yond that, the additional equivariant representations en- incorporate the global structure of the molecule.
able the prediction of equivariant tensors of order M that
can be constructed using a rank-1 tensor decomposition
V. RESULTS
N X
X R
T = λ(si )~ν (~vi )k,1 ⊗ · · · ⊗ ~ν (~vi )k,M . (11)
PaiNN has been implemented using PyTorch42 and
i=1 k=1
SchNetPack43 . All models use F = 128 and two output
The scalar and vectorial components of this decomposition layers as shown in Fig. 2 (blue) with the type of output
can be obtained from an output network of gated equiv- layers depending on the property of interest. They were
ariant blocks, as shown in Fig. 3. These components again trained using the Adam optimizer44 , the squared loss and
make use gated equivariant nonlinearities17 and show sim- weight decay λ = 0.01? , if not stated otherwise. We decay
ilarities to geometric vector perceptrons (GVP)32 . How- the learning rate by a factor of 0.5 if the validation loss
ever, while GVP blocks are used as message functions, plateaus. We apply exponential smoothing with factor 0.9
6

to the validation loss to reduce the impact of fluctuations the atomic forces could be predicted directly from vec-
which are particularly common when training with both torial features, we employ the gradients of the energy
energies and forces. Please refer to the supplement for model F~i = −∂E/∂~ri to ensure conservation of energy.
further details on training parameters. This property is crucial to run stable molecular dynam-
ics simulations. To demonstrate the data efficiency of
PaiNN, we use the more challenging setting with 1k
A. Chemical compound space known structures of which we use 950 for training and 50
for validation, where a separate model is trained for each
We use the QM9 dataset of ≈130k small organic trajectory. Tab. 5 shows the comparison with sGDML4
molecules45 with up to nine heavy atoms to evaluate and NequIP30 , which were trained on forces only, as
the performance of PaiNN for the prediction of scalar well as SchNet, PhysNet, DimeNet and FCHL1911 ,
properties across chemical compound space. We predict that were trained on a combined loss of energies and
the magnitude of the dipole moment using Eq. 13 and forces. Christensen and von Lilienfeld 47 have found that
the electronic spatial extent by the energies of MD17 are noisy. Thus, depending on the
molecule and chosen tradeoff, using energies for training
N
X is not always beneficial. For this reason, we train two
2
R = qatom (si )k~ri k2 , PaiNN models per trajectory: only on forces and on a
i=1 combined loss including energy with a force error weight
of ρ = 0.95. PaiNN achieves the lowest mean absolute
as implemented by SchNetPack. The remaining prop- errors for 12 out 14 targets on models trained only on
erties are predicted as sums over atomic contributions. forces and exhibits errors in a similar range as Gaussian
PaiNN is trained on 110k examples while 10k molecules regression with the FCHL19 kernel. Overall, PaiNN
are used as a validation set for decaying the learning rate performs best or equal to FCHL19 on 9 out of 14 targets.
and early stopping. The remaining data is used as test This demonstrates that equivariant neural networks ap-
set and results are averaged over three random splits. For proaches such as PaiNN are able to compete with kernel
the isotropic polarizability α, we first observed validation methods in the small data regime, while being able to
MAEs of 0.054 a0 . Upon closer inspection, we notice scale to large data sets at the same time.
that for this property both the squared loss as well as the
MAE can be reduced when minimizing the MAE directly
(as done by Klicpera, Groß, and Günnemann 34 ). This
yields both validation and test MAEs of 0.045a0 that are
comparable to those of DimeNet++.
Tab. II shows the mean absolute error (MAE) of PaiNN
for 12 target properties of QM9 in comparison with previ- C. Advantages of equivariant features
ous approaches. SchNet13 and PhysNet27 are MPNNs
with distance-based interactions, DimeNet++34 includes
additional angular information and is an improved vari- 1. Ablation studies
ant of DimeNet15 . L1Net14 and Cormorant20 are
equivariant neural networks based on spherical harmonics
and Clebsch-Gordon coefficients.
We evaluate the impact of equivariant vector features
PaiNN achieves state-of-the-art results in six target
at the example of the aspirin MD trajectory from the
properties and yields comparable results to DimeNet++
previous section. Compared to the full model, we remove
on another two targets. On the remaining proper-
the scalar product of vector features in Eq. 9 from the
ties, PaiNN achieves the second best results after
update block and the convolution over vector features
DimeNet++. Note that PaiNN using about 600k param-
in Eq. 8 (i.e., Wvv = 0). Table IV show the results
eters is significantly smaller than DimeNet++ with about
for the various ablations. The number of parameters
1.8M parameters. For random batches of 50 molecules
is kept approximately constant by raising the number
from QM9, the inference time is reduced from 45 ms to
of node features F accordingly. We observe that all
13 ms, i.e. an improvement of more than 70%, when
ablated components contribute to the final accuracy of the
comparing PaiNN to the reference implementation of
model, where the convolution over equivariant features in
DimeNet++46 using an NVIDIA V100.
the message function has a slightly larger impact. This
component also enables the propagation of directional
information, which will be examined in Sec. V C 2. Finally,
B. Molecular dynamics trajectories we remove all vector features from the model, resulting
in an invariant model. Despite keeping the number of
We evaluate the ability to predict combined energies parameters constant by increasing the number of atoms
and forces on the MD17 benchmark10 including molecular features to F = 174, the mean absolute error of the forces
dynamics trajectories of small organic molecules. While increases beyond 1 kcal/mol/Å.
7

TABLE II: Mean absolute errors on QM9 dataset for various chemical properties. Results for PaiNN are averaged
over three random splits. Best in bold.

Target Unit SchNet PhysNet DimeNet++ Cormorant L1Net PaiNN


µ D 0.033 0.053 0.030 0.038 0.043 0.012
α a30 0.235 0.062 0.044 0.085 0.088 0.045
HOMO meV 41 32.9 24.6 34 46.0 27.6
LUMO meV 34 24.7 19.5 38 34.6 20.4
∆ meV 63 42.5 32.6 38 67.5 45.7
hR2 i a20 0.073 0.765 0.331 0.961 0.354 0.066
ZPVE meV 1.7 1.39 1.21 2.03 1.56 1.28
U0 meV 14 8.15 6.32 22 13.46 5.85
U meV 19 8.34 6.28 21 13.83 5.83
H meV 14 8.42 6.53 21 14.36 5.98
G meV 14 9.40 7.56 20 13.99 7.35
cal
cv mol K
0.033 0.028 0.023 0.026 0.031 0.024

TABLE III: Mean absolute errors on MD17 dataset for energy and force predictions in kcal/mol and kcal/mol/Å,
respectively. Batzner et al. 30 only reported force errors for NequIP. Results for PaiNN are averaged over three
random splits. Best in bold.

sGDML NequIP PaiNN SchNet PhysNet DimeNet FCHL19 PaiNN


trained on forces only trained on energies & forces
energy 0.19 – 0.167 0.37 0.230 0.204 0.182 0.159
Aspirin
forces 0.68 0.348 0.338 1.35 0.605 0.499 0.478 0.371
energy 0.07 – 0.064 0.08 0.059 0.064 0.054 0.063
Ethanol
forces 0.33 0.208 0.224 0.39 0.160 0.230 0.136 0.230
energy 0.10 – 0.100 0.13 0.094 0.104 0.081 0.091
Malondialdehyde
forces 0.41 0.337 0.344 0.66 0.319 0.383 0.245 0.319
energy 0.12 – 0.116 0.16 0.142 0.122 0.117 0.117
Naphthalene
forces 0.11 0.096 0.077 0.58 0.310 0.215 0.151 0.083
energy 0.12 – 0.116 0.20 0.126 0.134 0.114 0.114
Salicylic Acid
forces 0.28 0.238 0.195 0.85 0.337 0.374 0.221 0.209
energy 0.10 – 0.095 0.12 0.100 0.102 0.098 0.097
Toluene
forces 0.14 0.101 0.094 0.57 0.191 0.216 0.203 0.102
energy 0.11 – 0.106 0.14 0.108 0.115 0.104 0.104
Uracil
forces 0.24 0.172 0.139 0.56 0.218 0.301 0.105 0.140

TABLE IV: Ablation study for the prediction of energies 2. Propagation of directional information in substituted
[kcal/mol] and forces [kcal/mol/Å] for aspirin trajectories ferrocene
from MD17.

Ablation # params F energy MAE force MAE To demonstrate the advantages of equivariant over in-
no ablation 588.3k 128 0.159 0.371
variant representations in practice, we consider a ferrocene
derivative where one hydrogen atom in each cyclopenta-
no scalar product of 589.1k 134 0.173 0.420
vector features in Eq. 8 dienyl ring is substituted by fluorine (see Fig. 4). This
no vector propagation 589.2k 135 0.183 0.441 molecule has been chosen as it features small energy fluc-
(Wvv = 0 in Eq. 7)
remove both 590.1k 142 0.200 0.507 tuations (<1 kcal/mol) when the rings rotate relative to
no vector features 590.3k 174 0.449 1.194
each other. Since the torsional energy profile depends
mainly on the orientation of the distant fluorine atoms
(measured by the rotation angle θ), it is a challenging
prediction target for models without equivariant repre-
sentations. Fig. 4 shows the predicted energy profiles
for a full rotation of the cyclopentadienyl ring using cut-
offs rcut ∈ {2.5, 3.0, 4.0} Å. All models are trained on
energies of 10k structures sampling thermal fluctuations
(300 K) and ring rotations of substituted ferrocene. An-
8

cessive computations required. Since nuclear quantum


effects (NQE) need to be considered to obtain high qual-
ity spectra49 , ring-polymer molecular dynamics (RPMD)
simulations need to be performed, which treat multiple
copies of a molecule simultaneously and further increase
computational cost.
We train a joined model for energies, forces, dipole
moments and polarizability tensors on 8k ethanol confor-
mations using additional 1k molecules each as validation
and test set. A second model is trained for aspirin with
16k training examples and 2k molecules for validation and
testing. Energies and forces are predicted as described
above for MD17. The dipole moments and polarizability
tensors are obtained as described in Section IV B, employ-
ing two gated equivariant blocks each yielding atomwise
scalars and vectors, that are used to compute the outputs
as described in Eqs. 13 and 14. The joined model exhibits
FIG. 4: Rotational energy profile of substituted ferrocene
accurate predictions for energy, forces, dipole moments
obtained by varying the rotation angle θ while keeping
and polarizabilities, as shown in the inset table of Fig 5.
all bond distances fixed at their equilibrium values (H:
To evaluate the importance of equivariant features for the
white, C: black, F: purple, Fe: grey). For small cutoff
dipole prediction, we have trained an alternative model
rcut (dashed grey circles), MPNNs with scalar feature
using only latent partial charges, but no local dipoles.
representations are unable to represent information about
The results show that predicting molecular dipoles using
the rotation angle θ.
vector features (Eq. 13) reduces the mean abs. error by
more than 50%, reaching 0.009 D compared to 0.020 D
using only scalars and positions (Eq. 12). Simulated in-
other 1k structures are used for validation. SchNet13 , frared and Raman spectra using classical MD and RPMD
which uses scalar features, predicts a flat energy profile with 64 replicas and the respective PaiNN models are
for rcut ≤ 3.0 Å, because it is unable to resolve θ. Al- shown for both molecules in Fig. 5, alongside harmonic
though it is possible to encode the relevant information spectra obtained with the electronic structure reference
when using a larger cutoff, the learned representation and, for ethanol, experimental spectra recorded in the gas
seems to be susceptible to noise, as indicated by the devi- phase50,51 .
ations of the predicted energy profile. Since DimeNet15
The peak positions and intensities of ethanol infrared
includes terms that explicitly depend on angles between
spectra (Fig. 5a, top) computed with classical MD and
triplets of atoms, θ can be resolved for smaller cutoffs,
PaiNN agree closely with the static electronic structure
but still fails for rcut = 2.5 Å. In contrast, the equivariant
reference, e.g. in the C-H and O-H stretching regions
environment representation of PaiNN allows to faithfully
at 3000 cm−1 and 3900 cm−1 . This indicates, that
reproduce the torsional energy profile even for very small
PaiNN faithfully reproduces the original electronic struc-
cutoffs.
ture method in general. However, when compared to
experiment, both spectra are shifted towards higher fre-
quencies. The PaiNN RPMD spectrum, on the other
D. Molecular Spectra hand shows excellent agreement with experiment, illus-
trating the importance of NQEs. Similar observations
We apply PaiNN to the efficient computation of in- can be made for the ethanol Raman spectrum (Fig. 5a,
frared and Raman spectra of ethanol and aspirin. Al- bottom), where the PaiNN RMPD spectrum once again
though these spectra can in principle be obtained from offers the most faithful reproduction of the experiment.
a single molecular structure via the harmonic oscillator We observe similar trends for aspirin (Fig. 5b). RPMD
approximation, such an approach tends to neglect impor- spectra are once again shifted to lower frequencies while
tant effects, e.g. different molecular conformations48 . In the classical MD closely resembles the static electronic
order to obtain high quality spectra, molecular dynamics structure spectrum.
simulations have to be performed, requiring the prediction PaiNN reduces the time required to obtain these spec-
of forces at each time step. Additionally, one needs to tra by several orders of magnitude. For ethanol, an
compute dipole moments and polarizabilities along the individual reference computation takes ≈ 140 seconds
generated trajectory. The infrared and Raman spectra compared to 14 ms for a PaiNN evaluation on a V100
are obtained as the Fourier transform of the time auto- GPU. This corresponds to a reduction of the overall sim-
correlation functions of the respective property. Using ulation time from 400 days to approximately one hour.
electronic structure methods, such simulations become The speedup offered by PaiNN is even more pronounced
prohibitively expensive due to the large number of suc- for aspirin, where a simulation that would have taken 25
9

FIG. 5: IR (top) and Raman (bottom) spectra of ethanol and aspirin. Spectra calculated with the reference method
using the harmonic oscillator approximation are shown in black (QM harmonic). The inset table shows the mean
absolute errors on the respective test set.

years (3140 seconds / step) now takes one hour (15 ms / ACKNOWLEDGEMENTS
step).
KTS acknowledges support by the Federal Ministry of
Education and Research (BMBF) for the Berlin Center for
VI. CONCLUSIONS Machine Learning / BIFOLD (01IS18037A). MG works
at the BASLEARN – TU Berlin/BASF Joint Lab for
We have given general guidelines to design equivari- Machine Learning, co-financed by TU Berlin and BASF
ant MPNNs and discussed the advantages of equivariant SE. OTU acknowledges funding from the Swiss National
representations over angular features in terms of compu- Science Foundation (Grant No. P2BSP2 188147).
tational efficiency as well as their ability to propagate
1 J. Behler, “Perspective: Machine learning potentials for atomistic
directional information. On this basis, we have proposed
simulations,” J. Chem. Phys 145, 170901 (2016).
PaiNN that yields fast and accurate predictions of scalar 2 O. T. Unke, S. Chmiela, H. E. Sauceda, M. Gastegger, I. Poltavsky,
and tensorial molecular properties. Thereby, equivariant K. T. Schütt, A. Tkatchenko, and K.-R. Müller, “Machine Learn-
message passing allows us to significantly reduce both ing Force Fields,” arXiv preprint arXiv:2010.07067 (2020).
3 O. A. von Lilienfeld, K.-R. Müller, and A. Tkatchenko, “Exploring
model size and inference time compared to directional
message-passing while retaining accuracy. Finally, we chemical compound space with quantum-based machine learning,”
Nat. Rev. Chem. 4, 347–358 (2020).
have demonstrated that PaiNN can be applied to the 4 S. Chmiela, H. E. Sauceda, K.-R. Müller, and A. Tkatchenko,
prediction of tensorial properties, which we leverage to ac- “Towards exact molecular dynamics simulations with machine-
celerate the simulation of molecular spectra by 4-5 orders learned force fields,” Nat. Commun. 9, 3887 (2018).
5 J. Westermayr, M. Gastegger, and P. Marquetand, “Combining
of magnitude – from years to hours.
SchNet and SHARC: The SchNarc machine learning approach
In future work, the equivariant representation of
for excited-state dynamics,” J. Phys. Chem. Lett. 11, 3828–3834
PaiNN as well as the ability to predict tensorial proper- (2020).
ties may be leveraged in generative models of 3d geome- 6 T. Morawietz, A. Singraber, C. Dellago, and J. Behler, “How van

tries52–54 or the prediction of wavefunctions55–57 . We see der Waals interactions determine the unique properties of water,”
further applications of equivariant message passing in 3d Proc. Natl. Acad. Sci. 113, 8368–8373 (2016).
7 A. P. Bartók, J. Kermode, N. Bernstein, and G. Csányi, “Machine
shape recognition and graph embedding58 . learning a general-purpose interatomic potential for silicon,” Phys.
Many challenges remain for the fast and accurate predic- Rev. X 8, 041048 (2018).
tion of molecular properties, e.g. modeling of enzymatic 8 D. Lu, H. Wang, M. Chen, J. Liu, L. Lin, R. Car, W. Jia,

active sites or surface reactions. To describe such phenom- L. Zhang, et al., “86 PFLOPS Deep Potential Molecular Dynamics
ena, highly accurate reference methods are required. Due simulation of 100 million atoms with ab initio accuracy,” arXiv
preprint arXiv:2004.11658 (2020).
to their computational cost, reference data generation can 9 J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl,
become a bottleneck, making data-efficient MPNNs such “Neural message passing for quantum chemistry,” Interational
as PaiNN invaluable for future chemistry research. Conference on Machine Learning (2017).
10

10 S. Chmiela, A. Tkatchenko, H. E. Sauceda, I. Poltavsky, K. T. 31 R. Kondor and S. Trivedi, “On the Generalization of Equivariance
Schütt, and K.-R. Müller, “Machine Learning of Accurate Energy- and Convolution in Neural Networks to the Action of Compact
Conserving Molecular Force Fields,” Sci. Adv. 3, e1603015 (2017). Groups,” International Conference on Machine Learning , 2747–
11 A. S. Christensen, L. A. Bratholm, F. A. Faber, and O. Ana- 2755 (2018).
tole von Lilienfeld, “FCHL revisited: Faster and more accurate 32 B. Jing, S. Eismann, P. Suriana, R. J. L. Townshend, and R. Dror,

quantum machine learning,” J. Chem. Phys 152, 044107 (2020). “Learning from protein structure with geometric vector percep-
12 A. P. Bartók, M. C. Payne, R. Kondor, and G. Csányi, “Gaussian trons,” in International Conference on Learning Representations
approximation potentials: The accuracy of quantum mechanics, (2021).
without the electrons,” Phys. Rev. Lett. 104, 136403 (2010). 33 G. E. Hinton, A. Krizhevsky, and S. D. Wang, “Transforming
13 K. Schütt, P.-J. Kindermans, H. E. Sauceda, S. Chmiela, auto-encoders,” International Conference on Artificial Neural
A. Tkatchenko, and K.-R. Müller, “SchNet: A continuous-filter Networks , 44–51 (2011).
convolutional neural network for modeling quantum interactions,” 34 J. Klicpera, J. Groß, and S. Günnemann, “Fast and Uncertainty-

Advances in Neural Information Processing Systems , 991–1001 Aware Directional Message Passing for Non-Equilibrium
(2017). Molecules,” Machine Learning for Molecules Workshop at NeurIPS
14 B. K. Miller, M. Geiger, T. E. Smidt, and F. Noé, “Relevance (2020).
of rotationally equivariant convolutions for predicting molecular 35 S. N. Pozdnyakov, M. J. Willatt, A. P. Bartók, C. Ortner,

properties,” arXiv preprint arXiv:2008.08461 (2020). G. Csányi, and M. Ceriotti, “Incompleteness of atomic structure
15 J. Klicpera, J. Groß, and S. Günnemann, “Directional message representations,” Physical Review Letters 125, 166001 (2020).
passing for molecular graphs,” International Conference on Learn- 36 A. Hermann, R. P. Krawczyk, M. Lein, P. Schwerdtfeger, I. P.

ing Representations (2020). Hamilton, and J. J. Stewart, “Convergence of the many-body


16 T. S. Cohen and M. Welling, “Steerable CNNs,” International expansion of interaction potentials: From van der waals to covalent
Conference on Learning Representations (2017). and metallic systems,” Physical Review A 76, 013202 (2007).
17 M. Weiler, F. A. Hamprecht, and M. Storath, “Learning Steerable 37 A. P. Bartók, R. Kondor, and G. Csányi, “On representing chem-

Filters for Rotation Equivariant CNNs,” Proceedings of the IEEE ical environments,” Phys. Rev. B 87, 184115 (2013).
Conference on Computer Vision and Pattern Recognition , 849– 38 M. Gastegger, K. T. Schütt, and K.-R. Müller, “Machine learning

858 (2018). of solvent effects on molecular spectra and reactions,” arXiv


18 D. Worrall and G. Brostow, “CubeNet: Equivariance to 3D Rota- preprint arXiv:2010.14942 (2020).
tion and Translation,” Proceedings of the European Conference 39 J. Behler, “Atom-centered symmetry functions for constructing

on Computer Vision (ECCV) , 567–584 (2018). high-dimensional neural network potentials,” J. Chem. Phys 134,
19 N. Thomas, T. Smidt, S. Kearnes, L. Yang, L. Li, K. Kohlhoff, 074106 (2011).
and P. Riley, “Tensor field networks: Rotation-and translation- 40 M. Gastegger, J. Behler, and P. Marquetand, “Machine learning

equivariant neural networks for 3D point clouds,” arXiv preprint molecular dynamics for the simulation of infrared spectra,” Chem.
arXiv:1802.08219 (2018). Sci. 8, 6924–6935 (2017).
20 B. Anderson, T. S. Hy, and R. Kondor, “Cormorant: Covariant 41 M. Veit, D. M. Wilkins, Y. Yang, R. A. DiStasio Jr, and M. Ceri-

Molecular Neural Networks,” Advances in Neural Information otti, “Predicting molecular dipole moments by combining atomic
Processing Systems 32, 14537–14546 (2019). partial charges and atomic dipoles,” J. Chem. Phys. 153, 024113
21 J. Behler and M. Parrinello, “Generalized neural-network repre- (2020).
sentation of high-dimensional potential-energy surfaces,” Phys. 42 A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan,

Rev. Lett. 98, 146401 (2007). T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al., “Pytorch: An
22 F. Scarselli, M. Gori, A. C. Tsoi, M. Hagenbuchner, and G. Mon- imperative style, high-performance deep learning library,” arXiv
fardini, “The graph neural network model,” IEEE Trans. Neural preprint arXiv:1912.01703 (2019).
Netw. 20, 61–80 (2008). 43 K. Schütt, P. Kessel, M. Gastegger, K. Nicoli, A. Tkatchenko, and
23 D. K. Duvenaud, D. Maclaurin, J. Iparraguirre, R. Bombarell, K.-R. Müller, “SchNetPack: A deep learning toolbox for atomistic
T. Hirzel, A. Aspuru-Guzik, and R. P. Adams, “Convolutional systems,” J. Chem. Theory Comput. 15, 448–455 (2018).
networks on graphs for learning molecular fingerprints,” Advances 44 D. P. Kingma and J. Ba, “Adam: A method for Stochastic

in Neural Information Processing Systems 28, 2224–2232 (2015). Optimization,” arXiv preprint arXiv:1412.6980 (2014).
24 S. Kearnes, K. McCloskey, M. Berndl, V. Pande, and P. Riley, 45 R. Ramakrishnan, P. O. Dral, M. Rupp, and O. A. Von Lilien-

“Molecular graph convolutions: moving beyond fingerprints,” J. feld, “Quantum chemistry structures and properties of 134 kilo
Comput. Aided Mol. Des. 30, 595–608 (2016). molecules,” Sci. Data 1, 1–7 (2014).
25 K. T. Schütt, F. Arbabzadah, S. Chmiela, K. R. Müller, and 46 https://github.com/klicperajo/dimenet.

A. Tkatchenko, “Quantum-chemical insights from deep tensor neu- 47 A. S. Christensen and O. A. von Lilienfeld, “On the role of

ral networks,” Nat. Commun. 8 (2017), 10.1038/ncomms13890. gradients for machine learning of molecular energies and forces,”
26 K. T. Schütt, H. E. Sauceda, P.-J. Kindermans, A. Tkatchenko, Mach. Learn.: Sci. Technol. 1, 045018 (2020).
and K.-R. Müller, “SchNet–A deep learning architecture for 48 M. Thomas, M. Brehm, R. Fligg, P. Vöhringer, and B. Kirchner,

molecules and materials,” J. Chem. Phys 148, 241722 (2018). “Computing vibrational spectra from ab initio molecular dynamics,”
27 O. T. Unke and M. Meuwly, “PhysNet: A Neural Network for Pre- Phys. Chem. Chem. Phys. 15, 6608–6622 (2013).
dicting Energies, Forces, Dipole Moments, and Partial Charges,” 49 H. E. Sauceda, V. Vassilev-Galindo, S. Chmiela, K.-R. Müller,

J. Chem. Theory Comput. 15, 3678–3693 (2019). and A. Tkatchenko, “Dynamical strengthening of covalent and
28 N. Lubbers, J. S. Smith, and K. Barros, “Hierarchical modeling of non-covalent molecular interactions by nuclear quantum effects
molecular energies using a deep neural network,” J. Chem. Phys at finite temperature,” Nat. Commun. 12, 1–10 (2021).
148, 241715 (2018). 50 P. Linstrom and W. G. Mallard, NIST Chemistry WebBook NIST
29 M. Weiler, M. Geiger, M. Welling, W. Boomsma, and T. S. Cohen, Standard Reference Database Number 69 (National Institute
“3D Steerable CNNs: Learning Rotationally Equivariant Features of Standards and Technology, doi:10.18434/T4D303, (retrieved
in Volumetric Data,” Advances in Neural Information Processing September 24, 2020), Gaithersburg MD, 20899, 2020).
Systems 31, 10381–10392 (2018). 51 J. Kiefer, “Simultaneous acquisition of the polarized and depo-
30 S. Batzner, T. E. Smidt, L. Sun, J. P. Mailoa, M. Kornbluth, larized Raman signal with a single detector,” Anal. Chem. 89,
N. Molinari, and B. Kozinsky, “SE (3)-Equivariant Graph Neural 5725–5728 (2017).
Networks for Data-Efficient and Accurate Interatomic Potentials,” 52 N. Gebauer, M. Gastegger, and K. Schütt, “Symmetry-adapted

arXiv preprint arXiv:2101.03164 (2021). generation of 3d point sets for the targeted discovery of molecules,”
Advances in Neural Information Processing Systems , 7564–7576
11

(2019).
53 J. Köhler, L. Klein, and F. Noé, “Equivariant Flows: sampling
configurations for multi-body systems with symmetric energies,”
Proceedings of the 37th International Conference on Machine
Learning (2019).
54 G. N. Simm, R. Pinsler, G. Csányi, and J. M. Hernández-Lobato,

“Symmetry-Aware Actor-Critic for 3D Molecular Design,” arXiv


preprint arXiv:2011.12747 (2020).
55 G. Hegde and R. C. Bowen, “Machine-learned approximations

to Density Functional Theory Hamiltonians,” Sci. Rep. 7, 1–11


(2017).
56 K. Schütt, M. Gastegger, A. Tkatchenko, K.-R. Müller, and

R. J. Maurer, “Unifying machine learning and quantum chemistry


with a deep neural network for molecular wavefunctions,” Nat.
Commun. 10, 1–10 (2019).
57 J. Hermann, Z. Schätzle, and F. Noé, “Deep-neural-network so-

lution of the electronic Schrödinger equation,” Nat. Chem. 12,


891–897 (2020).
58 P. Goyal and E. Ferrara, “Graph embedding techniques, applica-

tions, and performance: A survey,” Knowledge-Based Systems


151, 78–94 (2018).
59 C. Bannwarth, S. Ehlert, and S. Grimme, “GFN2-xTB – An accu-

rate and broadly parametrized self-consistent tight-binding quan-


tum chemical method with multipole electrostatics and density-
dependent dispersion contributions,” J. Chem. Theory Comput.
15, 1652–1671 (2019).
60 https://github.com/grimme-lab/xtb.
61 J. S. Smith, O. Isayev, and A. E. Roitberg, “ANI-1: an exten-

sible neural network potential with DFT accuracy at force field


computational cost,” Chem. Sci. 8, 3192–3203 (2017).
62 C. Adamo and V. Barone, “Toward reliable density functional

methods without adjustable parameters: The PBE0 model,” J.


Chem. Phys. 110, 6158–6170 (1999).
63 F. Weigend and R. Ahlrichs, “Balanced basis sets of split valence,

triple zeta valence and quadruple zeta valence quality for H to Rn:
Design and assessment of accuracy,” Phys. Chem. Chem. Phys.
7, 3297–3305 (2005).
64 F. Neese, “The ORCA program system,” WIREs Comput. Mol.

Sci. 2, 73–78 (2012).


65 F. Weigend, “A fully direct RI-HF algorithm: Implementation,

optimised auxiliary basis sets, demonstration of accuracy and


efficiency,” Phys. Chem. Chem. Phys. 4, 4285–4291 (2002).
66 G. J. Martyna, M. L. Klein, and M. Tuckerman, “Nosé–Hoover

chains: The canonical ensemble via continuous dynamics,” J.


Chem. Phys. 97, 2635–2643 (1992).
67 M. Ceriotti, M. Parrinello, T. E. Markland, and D. E. Manolopou-

los, “Efficient stochastic thermostatting of path integral molecular


dynamics,” J. Chem. Phys. 133, 124104 (2010).
68 N. Wiener, “Generalized harmonic analysis,” Acta Math. 55,

117–258 (1930).
69 R. B. Blackman and J. W. Tukey, “The measurement of power

spectra from the point of view of communications engineering –


Part I,” Bell Syst. Tech. J. 37, 185–282 (1958).
12

Appendix A: Data Appendix B: Training details

All models were trained using the Adam optimizer44 .


1. Reference data for substituted ferrocene The learning rate is decayed by a factor of 0.5 if the
validation loss plateaus, starting with the largest learning
rate that does not diverge in steps 1e-4, 5e-4, 1e-3, etc.
Reference computations for substituted ferrocene were We apply exponential smoothing with factor 0.9 to the
carried out with the semi-empirical GFN2-xTB method59 validation loss to reduce the impact of fluctuations which
using the xtb package60 ). Training structures were gener- are particularly common when training with both energies
ated by normal mode sampling61 at 300 K starting from and forces. We use smaller batches for datasets where we
the global minimum structure and randomly rotating train on both energies and forces, since it is commonly
the cylopentadienyl (Cp) moieties relative to each other observed that larger batch sizes converge to larger errors
around the Cp-Fe-Cp axis by angles uniformly sampled in this setting.
from [0, 2π]. The rotational potential energy profile was
generated by sampling a full rotation around the Cp-Fe-
Cp axis using 1k steps, starting from the global minimum Appendix C: Computation of infrared and Raman spectra
structure and keeping all other degrees of freedom fixed
at their equilibrium positions. All simulations were carried out with the molecular dy-
namics module implemented in SchNetPack26 . Classical
molecular dynamics simulations for ethanol and aspirin
were carried out for 50 ps at a temperature of 300 K con-
trolled via Nose-Hoover chain66 thermostat with a chain
length of 3 and time constant of 100 fs. The first 10 ps
of these trajectories were then discarded. Ring polymer
molecular dynamics using 64 replicas were performed us-
2. Reference data for infrared and Raman spectra ing the same simulation and equilibration periods and
temperature, but employed a specially adapted global
Nose-Hoover chain as introduced in Ref.67 for thermostat-
Electronic structure reference computations for aspirin ting instead. The overall chain settings were kept the
were carried out at the PBE0/def2-TZVP62,63 level of the- same as above. In all cases, the velocity Verlet algo-
ory using the ORCA quantum chemistry package64 . SCF rithm and a time step of 0.2 fs were used to integrate
convergence was set to tight and integration grid levels of the equations of motion. Infrared and Raman spectra
4 and 5 were employed during SCF iterations and the final were computed from the time-autocorrelation functions
computation of properties, respectively. Computations of the dipole moment and polarizability time derivatives
were accelerated using the RIJK approximation65 . The (see Ref.48 ). All autocorrelation functions were computed
reference data for aspirin was generated by selecting 20 000 with the Wiener-Khinchin theorem68 , using an autocorre-
random configurations from the MD17 database10 and lation depth of 2048 fs. A Hann window function69 and
recomputing them at the above level of theory. Reference zero-padding were applied to the autocorrelation func-
data for the ethanol molecule was taken from reference38 , tions in order to enhance the spectra. Raman spectra
which employed the same level of theory bar the RIJK were calculated using a laser frequency of 514 nm and
approximation. temperature of 300 K.
13

Supplementary Table 5: Training parameters for all experiments.

Data set batch size learning rate decay patience stopping patience rcut [Å]
QM9 100 5 · 10−4 5 30 5.0
MD17 10 1 · 10−3 50 150 5.0
Ferrocene 10 1 · 10−3 10 30 2.5-4.0
Spectra 10 5 · 10−4 15 50 2.7

You might also like