Original: So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems
Revised: So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems

Original: Abstract
Revised: Abstract

Original: The application of machine learning methods in quantum chemistry has enabled the study of numerous chemical phenomena, which are computationally intractable with traditional ab-initio methods.
Revised: The application of machine learning methods in quantum chemistry has enabled the study of numerous chemical phenomena, which are computationally intractable with traditional ab-initio methods.

Original: However, some quantum mechanical properties of molecules and materials depend on non-local electronic effects, which are often neglected due to the difficulty of modeling them efficiently.
Revised: However, some quantum mechanical properties of molecules and materials depend on non-local electronic effects, which are often neglected due to the difficulty of modeling them efficiently.

Original: This work proposes a modified attention mechanism adapted to the underlying physics, which allows to recover the relevant non-local effects.
Revised: This work proposes a modified attention mechanism adapted to the underlying physics, which allows to recover the relevant non-local effects.

Original: Namely, we introduce spherical harmonic coordinates (SPHCs) to reflect higher-order geometric information for each atom in a molecule, enabling a non-local formulation of attention in the SPHC space.
Revised: Namely, we introduce spherical harmonic coordinates (SPHCs) to reflect higher-order geometric information for each atom in a molecule, enabling a non-local formulation of attention in the SPHC space.

Original: Our proposed model S O 3 KRATES 3 – a self-attention based message passing neural network – uncouples geometric information from atomic features, making them independently amenable to attention mechanisms.
Revised: Our proposed model S O 3 KRATES 3 – a self-attention based message passing neural network – uncouples geometric information from atomic features, making them independently amenable to attention mechanisms.

Original: Thereby we construct spherical filters , which extend the concept of continuous filters in Euclidean space to SPHC space and serve as foundation for a spherical self-attention mechanism.
Revised: Thereby we construct spherical filters , which extend the concept of continuous filters in Euclidean space to SPHC space and serve as foundation for a spherical self-attention mechanism.

Original: We show that in contrast to other published methods, S O 3 KRATES is able to describe nonlocal quantum mechanical effects over arbitrary length scales.
Revised: We show that in contrast to other published methods, S O 3 KRATES is able to describe nonlocal quantum mechanical effects over arbitrary length scales.

Original: Further, we find evidence that the inclusion of higher-order geometric correlations increases data efficiency and improves generalization.
Revised: Further, we find evidence that the inclusion of higher-order geometric correlations increases data efficiency and improves generalization.

Original: S O 3 KRATES matches or exceeds stateof-the-art performance on popular benchmarks, notably, requiring a significantly lower number of parameters (0.25–0.4x) while at the same time giving a substantial speedup (6–14x for training and 2–11x for inference) compared to other models.
Revised: S O 3 KRATES matches or exceeds stateof-the-art performance on popular benchmarks, notably, requiring a significantly lower number of parameters (0.25–0.4x) while at the same time giving a substantial speedup (6–14x for training and 2–11x for inference) compared to other models.

Original:  Introduction
Revised:  Introduction

Original: Atomistic simulations use long time-scale molecular dynamics (MD) trajectories to predict macroscopic properties that arise from interactions on the microscopic scale [1–3].
Revised: Atomistic simulations use long time-scale molecular dynamics (MD) trajectories to predict macroscopic properties that arise from interactions on the microscopic scale [1–3].

Original: Their predictive reliability is determined by the accuracy of the underlying force field (FF), which needs to be queried at every time step.
Revised: Their predictive reliability is determined by the accuracy of the underlying force field (FF), which needs to be queried at every time step.

Original: This quickly becomes a computational bottleneck if the forces are determined from first principles, which may be required for accurate results.
Revised: This quickly becomes a computational bottleneck if the forces are determined from first principles, which may be required for accurate results.

Original: To that end, machine learning
Revised: To that end, machine learning

Original: [Figure]
Revised: [Figure]

Original: v i[Figure] r [Figure]
Revised: i

Original: p a r a

Original: [Figure]
Revised: [Figure]

Original: e t
Revised: p a

Original: [Figure]
Revised: [Figure]

Original: s
Revised: t

Original: [Figure]
Revised: [Figure]

Original: to local models, S O 3 KRATES is able to learn the long-range correlations between hydrogen rotors in cumulene and reproduces the ab-initio ground truth faithfully.
Revised: to local models, S O 3 KRATES is able to learn the long-range correlations between hydrogen rotors in cumulene and reproduces the ab-initio ground truth faithfully.

Original: FFs (MLFFs) offer a computationally more efficient, yet accurate empirical alternative to expensive ab-initio methods [2, 4–24].
Revised: FFs (MLFFs) offer a computationally more efficient, yet accurate empirical alternative to expensive ab-initio methods [2, 4–24].

Original: In recent years, Geometric Deep Learning has become a popular design paradigm, which exploits relevant symmetry groups of the underlying learning problem by incorporating a geometric prior [12, 25, 26].
Revised: In recent years, Geometric Deep Learning has become a popular design paradigm, which exploits relevant symmetry groups of the underlying learning problem by incorporating a geometric prior [12, 25, 26].

Original: This effectively restricts the learnable functions of the model to a subspace with a meaningful inductive bias.
Revised: This effectively restricts the learnable functions of the model to a subspace with a meaningful inductive bias.

Original: Prominent examples for such models are e.g. convolutional neural networks (CNNs) [27], which are equivariant w. r. t. the group of translations, or graph neural networks (GNNs) [28], which are invariant w. r. t. node permutation.
Revised: Prominent examples for such models are e.g. convolutional neural networks (CNNs) [27], which are equivariant w. r. t. the group of translations, or graph neural networks (GNNs) [28], which are invariant w. r. t. node permutation.

Original: For molecular property prediction, it has been shown that equivariance w. r. t. the 3D rotation group SO(3) greatly improves data efficiency and accuracy of the learned FFs [29–32].
Revised: For molecular property prediction, it has been shown that equivariance w. r. t. the 3D rotation group SO(3) greatly improves data efficiency and accuracy of the learned FFs [29–32].

Original: To achieve equivariance, architectures either rely on feature expansions in terms of spherical harmonics (SH) [33] or explicitly include (dihedral) angles [29, 34].
Revised: To achieve equivariance, architectures either rely on feature expansions in terms of spherical harmonics (SH) [33] or explicitly include (dihedral) angles [29, 34].

Original: While the latter scales quadratically (cubically) in the number of neighboring atoms and has been shown to be geometrically incomplete [35], the calculation of spherical harmonics scales only linear in the number of neighboring atoms, which makes them a fast and accurate alternative [30, 32, 36].
Revised: While the latter scales quadratically (cubically) in the number of neighboring atoms and has been shown to be geometrically incomplete [35], the calculation of spherical harmonics scales only linear in the number of neighboring atoms, which makes them a fast and accurate alternative [30, 32, 36].

Original: However, current higher-order geometric representations based on SHs usually result in expensive transformations, since an individual feature channel per SH degree (and order) is required.
Revised: However, current higher-order geometric representations based on SHs usually result in expensive transformations, since an individual feature channel per SH degree (and order) is required.

Original: As a result, going to higher degrees is computationally expensive and comes at the price of increasing complexity, resulting in state-of-the-art (SOTA) models with millions of parameters [30, 32, 34].
Revised: As a result, going to higher degrees is computationally expensive and comes at the price of increasing complexity, resulting in state-of-the-art (SOTA) models with millions of parameters [30, 32, 34].

Original: However, in order to be applicable to large molecular structures, models are required to be both efficient and accurate on all length scales.
Revised: However, in order to be applicable to large molecular structures, models are required to be both efficient and accurate on all length scales.

Original: Non-local electronic effects have been outlined as one of the major challenges for a new generation of
Revised: Non-local electronic effects have been outlined as one of the major challenges for a new generation of

Original: MLFFs [21].
Revised: MLFFs [21].

Original: They result in non-local, higher-order geometric relations between atoms.
Revised: They result in non-local, higher-order geometric relations between atoms.

Original: Most current architectures implicitly assume locality of interactions (expressed through a local neighborhood), which prohibits an efficient description of all relevant atomic interactions at larger scales.
Revised: Most current architectures implicitly assume locality of interactions (expressed through a local neighborhood), which prohibits an efficient description of all relevant atomic interactions at larger scales.

Original: Simply increasing the cutoff radius used to determine local neighborhoods is not an adequate solution, since it only shifts the problem to larger length scales [30].
Revised: Simply increasing the cutoff radius used to determine local neighborhoods is not an adequate solution, since it only shifts the problem to larger length scales [30].

Original: In this work, we propose spherical harmonic coordinates (SPHCs), which encode higher-order geometric information for each node in a molecular graph (Fig. 1e).
Revised: In this work, we propose spherical harmonic coordinates (SPHCs), which encode higher-order geometric information for each node in a molecular graph (Fig. 1e).

Original: This is in stark contrast to current approaches, which consider molecules as three-dimensional point clouds with learned features and fixed atomic coordinates: we propose to make the SPHCs themselves a learned quantity.
Revised: This is in stark contrast to current approaches, which consider molecules as three-dimensional point clouds with learned features and fixed atomic coordinates: we propose to make the SPHCs themselves a learned quantity.

Original: Through localization in the space of SPHCs (Fig. 1f), models are able to efficiently describe electronic effects that are non-local in three-dimensional Euclidean space (Fig. 1g).
Revised: Through localization in the space of SPHCs (Fig. 1f), models are able to efficiently describe electronic effects that are non-local in three-dimensional Euclidean space (Fig. 1g).

Original: We then present S O 3 KRATES (Fig. 1b), a self-attention based message-passing neural network (MPNN), which decouples atomic features and SPHCs and updates them individually (Fig. 1a).
Revised: We then present S O 3 KRATES (Fig. 1b), a self-attention based message-passing neural network (MPNN), which decouples atomic features and SPHCs and updates them individually (Fig. 1a).

Original: This resembles ideas from equivariant graph neural networks [25], but allows to go to arbitrarily high geometric orders.
Revised: This resembles ideas from equivariant graph neural networks [25], but allows to go to arbitrarily high geometric orders.

Original: The separation of higher-order geometric and feature information allows to overcome the parametric and computational complexity usually encountered in models with higherorder geometric representations, since we only require a single feature channel (instead of one per
Revised: The separation of higher-order geometric and feature information allows to overcome the parametric and computational complexity usually encountered in models with higherorder geometric representations, since we only require a single feature channel (instead of one per

Original: SH degree and order).
Revised: SH degree and order).

Original: Thus, S O 3 KRATES resembles some early architectures like S CH N ET [10] or
Revised: Thus, S O 3 KRATES resembles some early architectures like S CH N ET [10] or

Original: P HYS N ET [19] in parametric simplicity.
Revised: P HYS N ET [19] in parametric simplicity.

Original: We further show that S O 3 KRATES outperforms the popular
Revised: We further show that S O 3 KRATES outperforms the popular

Original: S GDML [38] kernel model by a large margin in the low-data regime, a domain which has so far been considered to be dominated by kernel machines [21].
Revised: S GDML [38] kernel model by a large margin in the low-data regime, a domain which has so far been considered to be dominated by kernel machines [21].

Original: Numerical evidence suggests that the data efficiency of S O 3 KRATES is directly related to the maximal degree of geometric information encoded in the SPHCs.
Revised: Numerical evidence suggests that the data efficiency of S O 3 KRATES is directly related to the maximal degree of geometric information encoded in the SPHCs.

Original: We then apply S O 3 KRATES to the well-established MD17 benchmark and show that our model achieves SOTA results, despite is light-weight structure and having only 0.25–0.4x the number of parameters of competitive architectures (Fig. 1c), while achieving speedups of 6–14x and 2–11x for training and inference, respectively (Fig. 1d).
Revised: We then apply S O 3 KRATES to the well-established MD17 benchmark and show that our model achieves SOTA results, despite is light-weight structure and having only 0.25–0.4x the number of parameters of competitive architectures (Fig. 1c), while achieving speedups of 6–14x and 2–11x for training and inference, respectively (Fig. 1d).

Original: Although we focus on quantum chemistry applications in this work, the developed methods are also applicable to other fields where long-ranged correlations in three-dimensional data are relevant.
Revised: Although we focus on quantum chemistry applications in this work, the developed methods are also applicable to other fields where long-ranged correlations in three-dimensional data are relevant.

Original: For example, models based on SPHCs may also be applicable to tasks like 3D shape classification or computer vision.
Revised: For example, models based on SPHCs may also be applicable to tasks like 3D shape classification or computer vision.

Original:  Preliminaries and Related Work
Revised:  Preliminaries and Related Work

Original: In the following, we review the most important concepts our method is based upon and relate it to prior work.
Revised: In the following, we review the most important concepts our method is based upon and relate it to prior work.

Original: Message Passing Neural Networks MPNNs [14] carry over many of the benefits of convolutions to unstructured domains and have thus become one of the most promising approaches for the description of molecular properties.
Revised: Message Passing Neural Networks MPNNs [14] carry over many of the benefits of convolutions to unstructured domains and have thus become one of the most promising approaches for the description of molecular properties.

Original: Their general working principle relies on the repeated iteration of message passing (MP) steps, which can be phrased as follows [14, 25]
Revised: Their general working principle relies on the repeated iteration of message passing (MP) steps, which can be phrased as follows [14, 25]

Original: [Equation]
Revised: [Equation]

Original: Here, m ij is the message between atoms i and j computed with the message function m ( · ) , m i is the aggregation of all messages in the neighborhood
Revised: Here, m ij is the message between atoms i and j computed with the message function m ( · ) , m i is the aggregation of all messages in the neighborhood

Original: N ( i ) of atom i , and u ( · ) is an update function returning updated features
Revised: N ( i ) of atom i , and u ( · ) is an update function returning updated features

Original: f ′ i based on the current features f i and message m i .
Revised: f ′ i based on the current features f i and message m i .

Original: The neighborhood
Revised: The neighborhood

Original: N ( i ) consists of all atoms which lie within a given cutoff radius around the atomic position r i , which ensures linear scaling in the number of atoms n .
Revised: N ( i ) consists of all atoms which lie within a given cutoff radius around the atomic position r i , which ensures linear scaling in the number of atoms n .

Original: While earlier variants parametrized messages only in terms of inter-atomic distances [13, 19], more recent approaches also take higher-order geometric information into account [25, 29, 31, 32, 39, 40].
Revised: While earlier variants parametrized messages only in terms of inter-atomic distances [13, 19], more recent approaches also take higher-order geometric information into account [25, 29, 31, 32, 39, 40].

Original: Molecules as Point Clouds
Revised: Molecules as Point Clouds

Original: A molecule can be considered as a point cloud of n atoms P 3 D
Revised: A molecule can be considered as a point cloud of n atoms P 3 D

Original: ( R , F ) ,
Revised: ( R , F ) ,

Original: where R =
Revised: where R =

Original: ( r 1 , ... r n ) denotes the set of atomic positions r i ∈ R 3 and F = ( f 1 , . .
Revised: ( r 1 , ... r n ) denotes the set of atomic positions r i ∈ R 3 and F = ( f 1 , . .

Original: , f n ) is the set of rotationally invariant atomic descriptors, or features, f i ∈ R F .
Revised: , f n ) is the set of rotationally invariant atomic descriptors, or features, f i ∈ R F .

Original: We write the distance vector pointing from the position of atom i to the position of atom j as r ij = r
Revised: We write the distance vector pointing from the position of atom i to the position of atom j as r ij = r

Original: j − r i , the distance as r ij = ∥ r ij ∥and the normalized distance vector as ˆ r ij = r ij /r ij .
Revised: j − r i , the distance as r ij = ∥ r ij ∥and the normalized distance vector as ˆ r ij = r ij /r ij .

Original: Given the point cloud, a density over Euclidean space assigning a vector value to each point r can be constructed as
Revised: Given the point cloud, a density over Euclidean space assigning a vector value to each point r can be constructed as

Original: [Equation]
Revised: [Equation]

Original: where δ is the Dirac delta function.
Revised: where δ is the Dirac delta function.

Original: It can be shown that applying a convolutional filter on ρ ( r ) resembles the update steps used in MPNNs [41].
Revised: It can be shown that applying a convolutional filter on ρ ( r ) resembles the update steps used in MPNNs [41].

Original: Equivariance Given a set of transformations that act on a vector space A as S g : A (cid:55)→ A to which we associate an abstract group G , a function f : A (cid:55)→ B is said to be equivariant w. r. t. G if
Revised: Equivariance Given a set of transformations that act on a vector space A as S g : A (cid:55)→ A to which we associate an abstract group G , a function f : A (cid:55)→ B is said to be equivariant w. r. t. G if

Original: [Equation]
Revised: [Equation]

Original: where T g : B (cid:55)→ B is an equivalent transformation on the output space [25].
Revised: where T g : B (cid:55)→ B is an equivalent transformation on the output space [25].

Original: Thus, in order to say that f is equivariant, it must hold that under transformation of the input, the output transforms “in the same way”.
Revised: Thus, in order to say that f is equivariant, it must hold that under transformation of the input, the output transforms “in the same way”.

Original: While equivariance has been a popular concept in signal processing for decades (cf.
Revised: While equivariance has been a popular concept in signal processing for decades (cf.

Original: e.g. [42] or wavelet neural networks [43]), recent years have seen efforts to design group equivariant NNs and kernel methods, since respecting relevant symmetries builds an important inductive bias [12, 44, 45].
Revised: e.g. [42] or wavelet neural networks [43]), recent years have seen efforts to design group equivariant NNs and kernel methods, since respecting relevant symmetries builds an important inductive bias [12, 44, 45].

Original: Examples are CNNs [27] which are equivariant w. r. t. translation, GNNs [14, 28] which are invariant ( T g = I ) w. r. t. permutation, or architectures which are equivariant w. r. t. the SO(3) group [25, 31, 33, 36, 46].
Revised: Examples are CNNs [27] which are equivariant w. r. t. translation, GNNs [14, 28] which are invariant ( T g = I ) w. r. t. permutation, or architectures which are equivariant w. r. t. the SO(3) group [25, 31, 33, 36, 46].

Original: In this work, we consider the SO(3) group of rotations, such that A is the Euclidean space R 3 , where the corresponding group actions are given by rotation matrices R θ ∈ R 3 × 3 .
Revised: In this work, we consider the SO(3) group of rotations, such that A is the Euclidean space R 3 , where the corresponding group actions are given by rotation matrices R θ ∈ R 3 × 3 .

Original: Spherical Harmonics
Revised: Spherical Harmonics

Original: The spherical harmonics are special functions defined on the surface of the sphere S 

Revised: The spherical harmonics are special functions defined on the surface of the sphere S

Original: = { ˆ r ∈ R 3 :
Revised: = { ˆ r ∈ R 3 :

Original: ∥ ˆ r ∥ 2 = 1 } and form an orthonormal basis for the irreducible representations (irreps) of SO(3) .
Revised: ∥ ˆ r ∥ 2 = 1 } and form an orthonormal basis for the irreducible representations (irreps) of SO(3) .

Original: In the context of tensor field networks [33], they have been introduced as elementary building blocks for SO(3) -equivariant neural networks.
Revised: In the context of tensor field networks [33], they have been introduced as elementary building blocks for SO(3) -equivariant neural networks.

Original: The spherical harmonics are commonly denoted as Y ml (ˆ r ) :
Revised: The spherical harmonics are commonly denoted as Y ml (ˆ r ) :

Original: S 2 (cid:55)→ R , where the degree l determines all possible values of the order m ∈ {− l, . .
Revised: S 2 (cid:55)→ R , where the degree l determines all possible values of the order m ∈ {− l, . .

Original: + l } .
Revised: + l } .

Original: They transform under rotation as
Revised: They transform under rotation as

Original: [Equation]
Revised: [Equation]

Original: (6) where D lmm ′
Revised: (6) where D lmm ′

Original: ( R θ ) are the entries of the Wigner-D matrix D l ( R θ ) ∈ R (2 l +1) ×
Revised: ( R θ ) are the entries of the Wigner-D matrix D l ( R θ ) ∈ R (2 l +1) ×

Original: (2 l +1)
Revised: (2 l +1)

Original: Based on the spherical harmonics, we define a vector-valued function Y ( l ) : S 2 (cid:55)→ R 2 l +1 for each degree l , with entries Y ml for all valid orders m of a given degree l .
Revised: Based on the spherical harmonics, we define a vector-valued function Y ( l ) : S 2 (cid:55)→ R 2 l +1 for each degree l , with entries Y ml for all valid orders m of a given degree l .

Original: Since Y ( l ) ( R θ ˆ r )
Revised: Since Y ( l ) ( R θ ˆ r )

Original: =
Revised: =

Original: D l ( R θ ) Y ( l ) (ˆ r ) (cf. eq.
Revised: D l ( R θ ) Y ( l ) (ˆ r ) (cf. eq.

Original: , Y ( l ) is equivariant w. r. t. SO(3) .
Revised: , Y ( l ) is equivariant w. r. t. SO(3) .

Original: Tensor Product Contractions The irreps Y ( l 1 ) and Y ( l 2 ) can be coupled by computing their tensor product Y ( l 1 ) ⊗ Y ( l 2 ) , which can equivalently be expressed as a direct sum [33, 48]
Revised: Tensor Product Contractions The irreps Y ( l 1 ) and Y ( l 2 ) can be coupled by computing their tensor product Y ( l 1 ) ⊗ Y ( l 2 ) , which can equivalently be expressed as a direct sum [33, 48]

Original: [Equation]
Revised: [Equation]

Original: where the entry of order m 3 for the coupled irreps ˜ Y
Revised: where the entry of order m 3 for the coupled irreps ˜ Y

Original: ( l 3 ) is given by
Revised: ( l 3 ) is given by

Original: [Equation]
Revised: [Equation]

Original: and C l 1 ,l 2 ,l 3 m 1 ,m 2 ,m 3 are the so-called Clebsch-Gordon coefficients .
Revised: and C l 1 ,l 2 ,l 3 m 1 ,m 2 ,m 3 are the so-called Clebsch-Gordon coefficients .

Original: In the following, we will denote the tensor product of degrees l 1 and l 2 followed by “contraction” to l 3 (meaning the irreps of degree l 3 in the direct sum representation of their tensor product) as (cid:0) Y ( l 1 ) ⊗ l 3 Y ( l 2 ) (cid:1) , which is a mapping of the form R (2 l 1 +1) ×
Revised: In the following, we will denote the tensor product of degrees l 1 and l 2 followed by “contraction” to l 3 (meaning the irreps of degree l 3 in the direct sum representation of their tensor product) as (cid:0) Y ( l 1 ) ⊗ l 3 Y ( l 2 ) (cid:1) , which is a mapping of the form R (2 l 1 +1) ×

Original: (2 l 2 +1)
Revised: (2 l 2 +1)

Original: (cid:55)→ R 2 l 3 +1 , since m 3 ∈ {− l 3 , . . .
Revised: (cid:55)→ R 2 l 3 +1 , since m 3 ∈ {− l 3 , . . .

Original: l 3 } .
Revised: l 3 } .

Original:  Methods
Revised:  Methods

Original: In the following, we describe the main methodological contributions of this work.
Revised: In the following, we describe the main methodological contributions of this work.

Original: We introduce the concept of an adapted point cloud P 3 D ( R , X , F ) , which incorporates the set of spherical harmonics coordinates (SPHCs)
Revised: We introduce the concept of an adapted point cloud P 3 D ( R , X , F ) , which incorporates the set of spherical harmonics coordinates (SPHCs)

Original: X
Revised: X

Original: = ( χ 1 , . . .
Revised: = ( χ 1 , . . .

Original: , χ n ) (see below) in addition to features F and Euclidean coordinates R .
Revised: , χ n ) (see below) in addition to features F and Euclidean coordinates R .

Original: However, contrary to R , SPHCs X are refined during the message passing updates.
Revised: However, contrary to R , SPHCs X are refined during the message passing updates.

Original: Having SPHCs as part of the molecular point cloud extends the idea of current MPNNs, which learn message functions on R , only.
Revised: Having SPHCs as part of the molecular point cloud extends the idea of current MPNNs, which learn message functions on R , only.

Original: Instead, we learn a message function m (cf. eq.
Revised: Instead, we learn a message function m (cf. eq.

Original: (1)) on both, the (fixed) atomic coordinates R as well as on the SPHCs X .
Revised: (1)) on both, the (fixed) atomic coordinates R as well as on the SPHCs X .

Original: This adapted message-passing scheme allows to learn non-local geometric corrections.
Revised: This adapted message-passing scheme allows to learn non-local geometric corrections.

Original: Based on these design principles, we propose the
Revised: Based on these design principles, we propose the

Original: S O 3 KRATES architecture.
Revised: S O 3 KRATES architecture.

Original: Initialization Feature vectors are initialized from the atomic numbers z i ∈ N (denoting which chemical element an atom belongs to) by an embedding map f
Revised: Initialization Feature vectors are initialized from the atomic numbers z i ∈ N (denoting which chemical element an atom belongs to) by an embedding map f

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: (9) where f emb : N (cid:55)→ R F .
Revised: (9) where f emb : N (cid:55)→ R F .

Original: We define SPHCs χ as the concatenation of degrees L :=
Revised: We define SPHCs χ as the concatenation of degrees L :=

Original: { l min , . . .
Revised: { l min , . . .

Original: , l max } χ = [ χ ( l min ) (cid:124) (cid:123)(cid:122)
Revised: , l max } χ = [ χ ( l min ) (cid:124) (cid:123)(cid:122)

Original: (cid:125) ∈
Revised: (cid:125) ∈

Original: R 2 l min+ , . . .
Revised: R 2 l min+ , . . .

Original: , χ ( l max ) (cid:124) (cid:123)(cid:122)
Revised: , χ ( l max ) (cid:124) (cid:123)(cid:122)

Original: [Equation]
Revised: [Equation]

Original: such that their transformation under rotation can be expressed in terms of concatenated Wigner-D matrices (see appendix A.1).
Revised: such that their transformation under rotation can be expressed in terms of concatenated Wigner-D matrices (see appendix A.1).

Original: The short-hand χ ( l ) ∈ R 2 l +1 refers to the subset of SPHCs with degree l .
Revised: The short-hand χ ( l ) ∈ R 2 l +1 refers to the subset of SPHCs with degree l .

Original: They are initialized as
Revised: They are initialized as

Original: [Equation]
Revised: [Equation]

Original: where C i =
Revised: where C i =

Original: [Equation]
Revised: [Equation]

Original: ϕ
Revised: ϕ

Original: [Equation]
Revised: [Equation]

Original: ( r ij ) , ϕ
Revised: ( r ij ) , ϕ

Original: [Equation]
Revised: [Equation]

Original: : R
Revised: : R

Original: [Equation]
Revised: [Equation]

Original: R is the cosine cutoff function [6], and the sum runs over the neighborhood N ( i ) of atom i .
Revised: R is the cosine cutoff function [6], and the sum runs over the neighborhood N ( i ) of atom i .

Original: Message Passing Update Two branches of attention-weighted MP steps are defined for the feature vectors f and SPHCs χ
Revised: Message Passing Update Two branches of attention-weighted MP steps are defined for the feature vectors f and SPHCs χ

Original: (see Fig. 1a).
Revised: (see Fig. 1a).

Original: After initialization (eqs.
Revised: After initialization (eqs.

Original: (9) and (11)), the features are updated as
Revised: (9) and (11)), the features are updated as

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where α ij ∈ R are self-attention [49, 50] coefficients (see below).
Revised: where α ij ∈ R are self-attention [49, 50] coefficients (see below).

Original: In analogy to the feature vectors, it is possible to define an MP update for the SPHCs as
Revised: In analogy to the feature vectors, it is possible to define an MP update for the SPHCs as

Original: [Equation]
Revised: [Equation]

Original: 13) where individual attention coefficients α ( l ) ij ∈ R for each degree of the SPHCs are computed using multi-head attention [49].
Revised: 13) where individual attention coefficients α ( l ) ij ∈ R for each degree of the SPHCs are computed using multi-head attention [49].

Original: However, with this definition, both MP updates are limited to local neighborhoods N ( i ) .
Revised: However, with this definition, both MP updates are limited to local neighborhoods N ( i ) .

Original: To be able to model non-local effects, we introduce the SPHC distance matrix
Revised: To be able to model non-local effects, we introduce the SPHC distance matrix

Original: X ∈ R n × n with entries χ ij = ∥ χ i − χ j ∥ 2 , i.e. distances between two atoms i and j in SPHC space for all possible pair-wise combinations of n atoms.
Revised: X ∈ R n × n with entries χ ij = ∥ χ i − χ j ∥ 2 , i.e. distances between two atoms i and j in SPHC space for all possible pair-wise combinations of n atoms.

Original: To have uniform scales, we further apply the softmax along each row of X to generate a rescaled matrix ˜ X
Revised: To have uniform scales, we further apply the softmax along each row of X to generate a rescaled matrix ˜ X

Original: = softmax( X ) with entries ˜
Revised: = softmax( X ) with entries ˜

Original: χ ij .
Revised: χ ij .

Original: A polynomial cutoff function ϕ
Revised: A polynomial cutoff function ϕ

Original: ˜ χ cut [29] is then applied to ˜ X to define spherical neighborhoods N χ
Revised: ˜ χ cut [29] is then applied to ˜ X to define spherical neighborhoods N χ

Original: ( i ) (see A.2), which may include atoms that are far away in Euclidean space (see Fig. 1f).
Revised: ( i ) (see A.2), which may include atoms that are far away in Euclidean space (see Fig. 1f).

Original: The spherical cutoff distance is chosen as ˜ χ cut = 1 /n to ensure that spherical neighborhoods remain small, even when going to larger molecules.
Revised: The spherical cutoff distance is chosen as ˜ χ cut = 1 /n to ensure that spherical neighborhoods remain small, even when going to larger molecules.

Original: We then incorporate non-local geometric corrections into the MP update of the SPHCs as
Revised: We then incorporate non-local geometric corrections into the MP update of the SPHCs as

Original: [Equation]
Revised: [Equation]

Original: We will show in the first part of the experiments, how geometric corrections from SPHC space allow for modelling non-local quantum effects, inaccessible to current architectures.
Revised: We will show in the first part of the experiments, how geometric corrections from SPHC space allow for modelling non-local quantum effects, inaccessible to current architectures.

Original: In the second part, we use a S O 3 KRATES model without geometric corrections, which makes it a traditional MPNN in the sense of only localizing in R 3 .
Revised: In the second part, we use a S O 3 KRATES model without geometric corrections, which makes it a traditional MPNN in the sense of only localizing in R 3 .

Original: We find this architecture to be highly parameter, data and time efficient while capable of reaching SOTA results.
Revised: We find this architecture to be highly parameter, data and time efficient while capable of reaching SOTA results.

Original: Spherical Filter and Self-Attention
Revised: Spherical Filter and Self-Attention

Original: The self-attention coefficients in eqs.
Revised: The self-attention coefficients in eqs.

Original: (12)–(14) are calculated as
Revised: (12)–(14) are calculated as

Original: [Equation]
Revised: [Equation]

Original: where w ij ∈ R F is the output of a filter generating function and ‘ ⊙ ’ denotes the element-wise product.
Revised: where w ij ∈ R F is the output of a filter generating function and ‘ ⊙ ’ denotes the element-wise product.

Original: The filter maps the Euclidean distance r ij and per-degree SPHC distances χ ( l ) ij = ∥ χ ( l ) j − χ
Revised: The filter maps the Euclidean distance r ij and per-degree SPHC distances χ ( l ) ij = ∥ χ ( l ) j − χ

Original: ( l ) i ∥between the current SPHCs of atoms i and j into the feature space R F (as a short-hand, we write the vector containing all per-degree SPHC distances as [ χ ( l ) ij ] l ∈L ).
Revised: ( l ) i ∥between the current SPHCs of atoms i and j into the feature space R F (as a short-hand, we write the vector containing all per-degree SPHC distances as [ χ ( l ) ij ] l ∈L ).

Original: It is built as the linear combination of two filter-generating functions
Revised: It is built as the linear combination of two filter-generating functions

Original: [Equation]
Revised: [Equation]

Original: which separately act on the Euclidean and SPHC distances.
Revised: which separately act on the Euclidean and SPHC distances.

Original: We call ϕ r : R (cid:55)→ R F the radial filter function and ϕ s : R |L| (cid:55)→ R F the spherical filter function (an ablation study for ϕ s can be found in appendix A.4).
Revised: We call ϕ r : R (cid:55)→ R F the radial filter function and ϕ s : R |L| (cid:55)→ R F the spherical filter function (an ablation study for ϕ s can be found in appendix A.4).

Original: Since per-atom features f i , interatomic distances r ij , and per-degree distances χ ( l ) ij are invariant under rotations (proof in appendix A.1), so are the self-attention coefficients α ij .
Revised: Since per-atom features f i , interatomic distances r ij , and per-degree distances χ ( l ) ij are invariant under rotations (proof in appendix A.1), so are the self-attention coefficients α ij .

Original: While we choose to pass the per-degree norms directly into the filter generating function ϕ s , future work might explore the possibilities of alternative metrics (instead of the L2 norm) or an expansion in terms of basis functions as it is common practice for inter-atomic distances (see appendix A. eq.
Revised: While we choose to pass the per-degree norms directly into the filter generating function ϕ s , future work might explore the possibilities of alternative metrics (instead of the L2 norm) or an expansion in terms of basis functions as it is common practice for inter-atomic distances (see appendix A. eq.

Original: Atomwise Interaction After each MP update, features and SPHCs are coupled with each other according to
Revised: Atomwise Interaction After each MP update, features and SPHCs are coupled with each other according to

Original: [Equation]
Revised: [Equation]

Original: where ϕ 1 : R F +2 |L| (cid:55)→ R F , ϕ ( l ) 2 : R F +2 |L| (cid:55)→ R , and ϕ ( l ) 3 : R |L| (cid:55)→ R .
Revised: where ϕ 1 : R F +2 |L| (cid:55)→ R F , ϕ ( l ) 2 : R F +2 |L| (cid:55)→ R , and ϕ ( l ) 3 : R |L| (cid:55)→ R .

Original: In the inputs to ϕ 1 , 2 , 3 , degree-wise scalars χ ( l ) =
Revised: In the inputs to ϕ 1 , 2 , 3 , degree-wise scalars χ ( l ) =

Original: ∥ χ ( l ) ∥ 2 are used to preserve equivariance.
Revised: ∥ χ ( l ) ∥ 2 are used to preserve equivariance.

Original: The coupling step additionally includes cross-degree coupled SPHCs
Revised: The coupling step additionally includes cross-degree coupled SPHCs

Original: ˜ χ ( l ) i for each degree l .
Revised: ˜ χ ( l ) i for each degree l .

Original: Following [48] they are constructed as
Revised: Following [48] they are constructed as

Original: [Equation]
Revised: [Equation]

Original: where k l 1 ,l 2 ,l ∈ R are learnable coefficients for all valid combinations of l 1 , l 2 given l and the term in brackets is the contraction of degrees l 1 and l 2 into degree l (eq.
Revised: where k l 1 ,l 2 ,l ∈ R are learnable coefficients for all valid combinations of l 1 , l 2 given l and the term in brackets is the contraction of degrees l 1 and l 2 into degree l (eq.

Original: S O 3 KRATES architecture Using the design paradigm above, we build the transformer network
Revised: S O 3 KRATES architecture Using the design paradigm above, we build the transformer network

Original: S O 3 KRATES , which consists of a self-attention block on F and X (eqs.
Revised: S O 3 KRATES , which consists of a self-attention block on F and X (eqs.

Original: (12) and (13)), respectively, as well as an interaction block
Revised: (12) and (13)), respectively, as well as an interaction block

Original: eqs.
Revised: eqs.

Original: (17) and (18)) per layer.
Revised: (17) and (18)) per layer.

Original: After initialization of the features and the SPHCs according to eqs.
Revised: After initialization of the features and the SPHCs according to eqs.

Original: (9) and (11), they are updated iteratively by passing through n l layers.
Revised: (9) and (11), they are updated iteratively by passing through n l layers.

Original: Atomic energy contributions E i ∈ R are predicted from the features of the final layer using a two-layered output block.
Revised: Atomic energy contributions E i ∈ R are predicted from the features of the final layer using a two-layered output block.

Original: The individual contributions are summed to the total energy prediction
Revised: The individual contributions are summed to the total energy prediction

Original: E = (cid:80) ni E i .
Revised: E = (cid:80) ni E i .

Original: See Fig.
Revised: See Fig.

Original: 1b for an overview.
Revised: 1b for an overview.

Original: More details on the implementation, training details and network hyperparameters are given in appendix A.3 and A.13.
Revised: More details on the implementation, training details and network hyperparameters are given in appendix A.3 and A.13.

Original: [Figure]
Revised: [Figure]

Original: So3krates ( l
Revised: So3krates ( l

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: corrections
Revised: corrections

Original: [Figure]
Revised: [Figure]

Original: corrections
Revised: corrections

Original: [Figure]
Revised: [Figure]

Original: ab
Revised: ab

Original: [Figure]
Revised: [Figure]

Original: ab
Revised: ab

Original: [Figure]
Revised: [Figure]

Original: initio
Revised: initio

Original: [Figure]
Revised: [Figure]

Original: e ﬀ .
Revised: e ﬀ .

Original: cuto ﬀ
Revised: cuto ﬀ

Original: [Figure]
Revised: [Figure]

Original: Overcome the e ﬀ ective cuto ﬀ induced by message passing
Revised: Overcome the e ﬀ ective cuto ﬀ induced by message passing

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: ( r cut sGDML (global)
Revised: ( r cut sGDML (global)

Original: SpookyNet (+ non-local
Revised: SpookyNet (+ non-local

Original: [Figure]
Revised: [Figure]

Original:  Experiments
Revised:  Experiments

Original: In the first subsection, we show how non-local quantum effects can be incorporated by using nonlocal corrections from the space of SPHCs.
Revised: In the first subsection, we show how non-local quantum effects can be incorporated by using nonlocal corrections from the space of SPHCs.

Original: In the second part of the experiments, we remove the non-local part which yields a traditional, R 3 -local MPNN which reaches SOTA results on established benchmarks while requiring much less computational time and parameters than competitive models.
Revised: In the second part of the experiments, we remove the non-local part which yields a traditional, R 3 -local MPNN which reaches SOTA results on established benchmarks while requiring much less computational time and parameters than competitive models.

Original: A scaling analysis as well as an accuracy comparison for both model variants can be found in appendix
Revised: A scaling analysis as well as an accuracy comparison for both model variants can be found in appendix

Original: A.8 and A.5.
Revised: A.8 and A.5.

Original: Non-Local Geometric Interactions For efficiency reasons
Revised: Non-Local Geometric Interactions For efficiency reasons

Original: , MPNNs only consider interactions between atoms in local neighborhoods, i.e. within a cutoff radius r cut .
Revised: , MPNNs only consider interactions between atoms in local neighborhoods, i.e. within a cutoff radius r cut .

Original: Thus, information can only be propagated over a distance of r cut within a single MP step.
Revised: Thus, information can only be propagated over a distance of r cut within a single MP step.

Original: Although multiple MP updates increase the effective cutoff distance, because information can “hop” between different neighborhoods as long as they share at least one atom, each MP step is accompanied by an undesirable loss of information, which limits the accuracy that can be obtained.
Revised: Although multiple MP updates increase the effective cutoff distance, because information can “hop” between different neighborhoods as long as they share at least one atom, each MP step is accompanied by an undesirable loss of information, which limits the accuracy that can be obtained.

Original: Consequently, MPNNs are unable to describe nonlocal effects on length-scales that exceed the effective cutoff distance.
Revised: Consequently, MPNNs are unable to describe nonlocal effects on length-scales that exceed the effective cutoff distance.

Original: To illustrate this problem, we consider the challenging open task [21] of learning the potential energy of cumulene molecules with different sizes (see Fig. 2a).
Revised: To illustrate this problem, we consider the challenging open task [21] of learning the potential energy of cumulene molecules with different sizes (see Fig. 2a).

Original: Here, the relative orientation of the hydrogen rotors at the far ends of the molecule strongly influences its energy due to non-local electronic effects [21].
Revised: Here, the relative orientation of the hydrogen rotors at the far ends of the molecule strongly influences its energy due to non-local electronic effects [21].

Original: In order to be able to successfully learn the energy profile with a local model, the effective cutoff has to be large enough to allow information to propagate from one hydrogen rotor to the other.
Revised: In order to be able to successfully learn the energy profile with a local model, the effective cutoff has to be large enough to allow information to propagate from one hydrogen rotor to the other.

Original: As a representative example for MPNNs, we consider the recently proposed N EQU IP model [32], which achieves SOTA performance on several benchmarks.
Revised: As a representative example for MPNNs, we consider the recently proposed N EQU IP model [32], which achieves SOTA performance on several benchmarks.

Original: We find that even when the effective cutoff radius is large enough in principle, an MPNN with n l = 4 , r cut = 2 .
Revised: We find that even when the effective cutoff radius is large enough in principle, an MPNN with n l = 4 , r cut = 2 .

Original: 5 Å, and l max ≤ 1 fails to learn the correct energy profile.
Revised: 5 Å, and l max ≤ 1 fails to learn the correct energy profile.

Original: This is due to the fact that the relevant geometric information “cancels out” (similar to addition of vectors oriented in opposite directions) within each neighborhood,
Revised: This is due to the fact that the relevant geometric information “cancels out” (similar to addition of vectors oriented in opposite directions) within each neighborhood,

Original: [Figure]
Revised: [Figure]

Original: Θ
Revised: Θ

Original: [Figure]
Revised: [Figure]

Original: 135 ° Θ = 180 ° close close
Revised: 135 ° Θ = 180 ° close close

Original: Layer
Revised: Layer

Original: [Figure]
Revised: [Figure]

Original: Θ = 180 ° Θ = 135 °
Revised: Θ = 180 ° Θ = 135 °

Original: Before training far far f a r a p a r t (a) underlining the limited expressiveness of mean-field interactions in MPNNs.
Revised: Before training far far f a r a p a r t (a) underlining the limited expressiveness of mean-field interactions in MPNNs.

Original: Only by including higher-order geometric correlations, e.g. going to l max = 3 , the correct energy profile can be recovered (at the cost of computational efficiency).
Revised: Only by including higher-order geometric correlations, e.g. going to l max = 3 , the correct energy profile can be recovered (at the cost of computational efficiency).

Original: When going to even larger cumulene structures, however, the effective cutoff becomes too small and it is necessary to increase the number of MP layers to solve the task (again, at the cost of lower computational efficiency), which is illustrated in appendix A.11.
Revised: When going to even larger cumulene structures, however, the effective cutoff becomes too small and it is necessary to increase the number of MP layers to solve the task (again, at the cost of lower computational efficiency), which is illustrated in appendix A.11.

Original: Neither increasing the maximum degree of interactions l max , nor the number of layers n l , is a satisfactory workaround: Instead of offering a general solution to describe non-local interactions, both options decrease computational efficiency, while only shifting the problem to larger length-scales or higher-order geometric correlations.
Revised: Neither increasing the maximum degree of interactions l max , nor the number of layers n l , is a satisfactory workaround: Instead of offering a general solution to describe non-local interactions, both options decrease computational efficiency, while only shifting the problem to larger length-scales or higher-order geometric correlations.

Original: We further apply three additional models to the cumulene structure with nine carbon atoms.
Revised: We further apply three additional models to the cumulene structure with nine carbon atoms.

Original: To that end, we use an invariant S CH N ET model with varying cutoff distances (6 Å and 12 Å), an inherently global but invariant S GDML model and the S POOKY N ET architecture which explicitly includes global effects using a non-local block.
Revised: To that end, we use an invariant S CH N ET model with varying cutoff distances (6 Å and 12 Å), an inherently global but invariant S GDML model and the S POOKY N ET architecture which explicitly includes global effects using a non-local block.

Original: We find that none of the three is capable of describing the rotor energy profile of cumulene.
Revised: We find that none of the three is capable of describing the rotor energy profile of cumulene.

Original: In contrast, our proposed S O 3 KRATES architecture is able to reproduce the energy profile for cumulene molecules of all sizes independent of the effective cutoff radius.
Revised: In contrast, our proposed S O 3 KRATES architecture is able to reproduce the energy profile for cumulene molecules of all sizes independent of the effective cutoff radius.

Original: Crucially, even with l max = 1 , the predicted energy matches the ab-initio reference faithfully.
Revised: Crucially, even with l max = 1 , the predicted energy matches the ab-initio reference faithfully.

Original: We find that geometric corrections in the MP update of the SPHCs (cf. eq.
Revised: We find that geometric corrections in the MP update of the SPHCs (cf. eq.

Original: (14)) are responsible for the increased capability of describing higher-order geometric correlations, as a S O 3 KRATES model with a naive MP update (cf. eq.
Revised: (14)) are responsible for the increased capability of describing higher-order geometric correlations, as a S O 3 KRATES model with a naive MP update (cf. eq.

Original: (13)) fails to solve this task with l max
Revised: (13)) fails to solve this task with l max

Original: = 1 (see Fig. 2a).
Revised: = 1 (see Fig. 2a).

Original: We further confirm that the model picks up on the physically relevant interaction between the hydrogen rotors by analysing the attention values after training (see Fig.
Revised: We further confirm that the model picks up on the physically relevant interaction between the hydrogen rotors by analysing the attention values after training (see Fig.

Original: 8 appendix A.7).
Revised: 8 appendix A.7).

Original: To illustrate how S O 3 KRATES is able to describe non-local effects, we show a low-dimensional projection of the atomic SPHCs before and after training for the largest of the cumulene molecules (Fig. 3).
Revised: To illustrate how S O 3 KRATES is able to describe non-local effects, we show a low-dimensional projection of the atomic SPHCs before and after training for the largest of the cumulene molecules (Fig. 3).

Original: After training, the SPHCs for hydrogen atoms at opposite ends of the molecule are embedded close together in SPHC space, allowing S O 3 KRATES to efficiently model the non-local geometric dependence between the hydrogen rotors.
Revised: After training, the SPHCs for hydrogen atoms at opposite ends of the molecule are embedded close together in SPHC space, allowing S O 3 KRATES to efficiently model the non-local geometric dependence between the hydrogen rotors.

Original: Generalization to structures, larger than those in the training data are usually associated with the re-usability of the learned, local representations.
Revised: Generalization to structures, larger than those in the training data are usually associated with the re-usability of the learned, local representations.

Original: For that reason, it is unclear if this property still holds when non-local corrections are used.
Revised: For that reason, it is unclear if this property still holds when non-local corrections are used.

Original: As we show in appendix A.5 a S O 3 KRATES model with non-local corrections still generalizes well to completely unknown and larger structures.
Revised: As we show in appendix A.5 a S O 3 KRATES model with non-local corrections still generalizes well to completely unknown and larger structures.

Original: Benchmarks, Data Efficiency and Generalization As pointed out in [31] and [32], equivariant features not only increase performance, but also improve data efficiency.
Revised: Benchmarks, Data Efficiency and Generalization As pointed out in [31] and [32], equivariant features not only increase performance, but also improve data efficiency.

Original: The latter is particularly important, as ab-initio methods for reference data generation can become exceedingly expensive when high accuracy is required.
Revised: The latter is particularly important, as ab-initio methods for reference data generation can become exceedingly expensive when high accuracy is required.

Original: Here, we use a subset of the recently introduced QM7-X data set [51], which we call QM7-X250.
Revised: Here, we use a subset of the recently introduced QM7-X data set [51], which we call QM7-X250.

Original: It contains 250 different molecular structures, each with 80 data points for training, 10 data points for validation and 11-3748 data points for testing (for details, see appendix A.9).
Revised: It contains 250 different molecular structures, each with 80 data points for training, 10 data points for validation and 11-3748 data points for testing (for details, see appendix A.9).

Original: The small number of training/validation samples per molecule makes it particularly suited for evaluating model behavior in the low data regime.
Revised: The small number of training/validation samples per molecule makes it particularly suited for evaluating model behavior in the low data regime.

Original: In the following, we train (1) one model
Revised: In the following, we train (1) one model

Original: [Figure]
Revised: [Figure]

Original: a
Revised: a

Original: [Figure]
Revised: [Figure]

Original: f
Revised: f

Original: o
Revised: o

Original: [Figure]
Revised: [Figure]

Original: ce
Revised: ce

Original: [Figure]
Revised: [Figure]

Original: a
Revised: a

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: VÅ
Revised: VÅ

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: l
Revised: l

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: a v .
Revised: a v .

Original: f
Revised: f

Original: o r ce m a e
Revised: o r ce m a e

Original: ( m e
Revised: ( m e

Original: VÅ - 1 )
Revised: VÅ - 1 )

Original: [Figure]
Revised: [Figure]

Original: per structure in QM7-X250 and (2) one model for all structures in QM7-X250 ( 250 × 80 = 20 k training points), which we refer to as individual and joint models, respectively.
Revised: per structure in QM7-X250 and (2) one model for all structures in QM7-X250 ( 250 × 80 = 20 k training points), which we refer to as individual and joint models, respectively.

Original: We start by investigating the performance as a function of the maximal degree l max and find that the error strongly decreases with higher l max (Fig. 4a).
Revised: We start by investigating the performance as a function of the maximal degree l max and find that the error strongly decreases with higher l max (Fig. 4a).

Original: As kernel methods are known to perform well in the low-data regime [21], we compare our results to S GDML [38] kernel models, which only use distances as a molecular descriptor (corresponding to l max = 0 ).
Revised: As kernel methods are known to perform well in the low-data regime [21], we compare our results to S GDML [38] kernel models, which only use distances as a molecular descriptor (corresponding to l max = 0 ).

Original: For l max = 0 , we find S GDML gives competitive results, whereas for l max = 1 , S O 3 KRATES starts to outperform S GDML.
Revised: For l max = 0 , we find S GDML gives competitive results, whereas for l max = 1 , S O 3 KRATES starts to outperform S GDML.

Original: As soon as l max ≥ 2 , however, the prediction accuracy of S O 3 KRATES is greater than that of S GDML by a large margin.
Revised: As soon as l max ≥ 2 , however, the prediction accuracy of S O 3 KRATES is greater than that of S GDML by a large margin.

Original: Thus, increasing the order of geometric information in the SPHCs leads to strong improvements in the low-data regime.
Revised: Thus, increasing the order of geometric information in the SPHCs leads to strong improvements in the low-data regime.

Original: For jointly trained models, we find that S O 3 KRATES outperforms S GDML even for l max = 0 , with continuous improvement for increasing l max .
Revised: For jointly trained models, we find that S O 3 KRATES outperforms S GDML even for l max = 0 , with continuous improvement for increasing l max .

Original: In appendix A.10 we report energy and force errors across degrees and further experimental details.
Revised: In appendix A.10 we report energy and force errors across degrees and further experimental details.

Original: The generalization capability of S O 3 KRATES is tested, by applying a jointly trained model to 25 completely unknown molecules from the QM7-X data set (see, Fig. 4b, details in appendix
Revised: The generalization capability of S O 3 KRATES is tested, by applying a jointly trained model to 25 completely unknown molecules from the QM7-X data set (see, Fig. 4b, details in appendix

Original: A.10).
Revised: A.10).

Original: Again, we find that force MAEs decrease with increasing l max .
Revised: Again, we find that force MAEs decrease with increasing l max .

Original: For reference, we compare
Revised: For reference, we compare

Original: S O 3 KRATES to individually trained S GDML models and find that S O 3 KRATES performs on par, or even slightly better, for l max ≥ 2 .
Revised: S O 3 KRATES to individually trained S GDML models and find that S O 3 KRATES performs on par, or even slightly better, for l max ≥ 2 .

Original: Going beyond l max = 2 is found to only marginally improve generalization.
Revised: Going beyond l max = 2 is found to only marginally improve generalization.

Original: In addition, we report results for a model trained on the full QM7-X data set in appendix A.5, following [30].
Revised: In addition, we report results for a model trained on the full QM7-X data set in appendix A.5, following [30].

Original: For completeness, we also apply S O 3 KRATES to the popular MD17 benchmark (see Table 1).
Revised: For completeness, we also apply S O 3 KRATES to the popular MD17 benchmark (see Table 1).

Original: We find, that S O 3 KRATES outperforms networks that have the same parameter complexity by a large margin (P AI NN and N EWTON N ET ).
Revised: We find, that S O 3 KRATES outperforms networks that have the same parameter complexity by a large margin (P AI NN and N EWTON N ET ).

Original: Notably, it requires significantly less parameters than other
Revised: Notably, it requires significantly less parameters than other

Original: SH based architectures (N EQU IP and S POOKY N ET ), while performing only slightly worse or even on par with them.
Revised: SH based architectures (N EQU IP and S POOKY N ET ), while performing only slightly worse or even on par with them.

Original: Furthermore, S O 3 KRATES outperforms D IME N ET , its closest competitor in timing (cf.
Revised: Furthermore, S O 3 KRATES outperforms D IME N ET , its closest competitor in timing (cf.

Original: Fig.
Revised: Fig.

Original: 1.d), consistently by a large margin.
Revised: 1.d), consistently by a large margin.

Original: Compared to current SH based approaches,
Revised: Compared to current SH based approaches,

Original: G EM N ET Q needs less parameters (still ∼ 2 .
Revised: G EM N ET Q needs less parameters (still ∼ 2 .

Original: 5 x more than S O 3 KRATES ) to achieve competitive results.
Revised: 5 x more than S O 3 KRATES ) to achieve competitive results.

Original: However, it requires the explicit calculation of dihedral angles which scales cubically in the number of neighboring atoms.
Revised: However, it requires the explicit calculation of dihedral angles which scales cubically in the number of neighboring atoms.

Original: Due to its linear scaling (see A.8) and lightweight structure, S O 3 KRATES can significantly reduce the time for training and inference (see Fig.
Revised: Due to its linear scaling (see A.8) and lightweight structure, S O 3 KRATES can significantly reduce the time for training and inference (see Fig.

Original: 1d and A.6).
Revised: 1d and A.6).

Original:  Discussion and Conclusion
Revised:  Discussion and Conclusion

Original: Due to the locality assumption used in most MPNNs, they are unable to model non-local electronic effects, which result in global geometric dependencies between different parts of a molecule.
Revised: Due to the locality assumption used in most MPNNs, they are unable to model non-local electronic effects, which result in global geometric dependencies between different parts of a molecule.

Original: The length-scales of such interactions often greatly exceed the cutoff radius used in the MP step, and even though stacking multiple MP layers increases the effective cutoff, ultimately, MPNNs are not capable of efficiently modeling geometric dependencies on arbitrary length scales.
Revised: The length-scales of such interactions often greatly exceed the cutoff radius used in the MP step, and even though stacking multiple MP layers increases the effective cutoff, ultimately, MPNNs are not capable of efficiently modeling geometric dependencies on arbitrary length scales.

Original: In this work, we contribute conceptually by proposing an efficient and scalable solution to this problem.
Revised: In this work, we contribute conceptually by proposing an efficient and scalable solution to this problem.

Original: We suggest a set of refinable, equivariant coordinates for point clouds in Euclidean space, called spherical harmonic coordinates (SPHCs).
Revised: We suggest a set of refinable, equivariant coordinates for point clouds in Euclidean space, called spherical harmonic coordinates (SPHCs).

Original: Non-local geometric effects can then be efficiently
Revised: Non-local geometric effects can then be efficiently

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: modeled by including geometric corrections, which are localized in the space of SPHCs, but non-local in Euclidean space.
Revised: modeled by including geometric corrections, which are localized in the space of SPHCs, but non-local in Euclidean space.

Original: Further, we show that introducing spherical filter functions acting on the SPHCs increases geometric resolution and predictive accuracy.
Revised: Further, we show that introducing spherical filter functions acting on the SPHCs increases geometric resolution and predictive accuracy.

Original: We then propose the S O 3 KRATES architecture, a self-attention based MPNN, which decouples atomic features from higher-order geometric information.
Revised: We then propose the S O 3 KRATES architecture, a self-attention based MPNN, which decouples atomic features from higher-order geometric information.

Original: This allows to drastically decrease the parametric complexity while still achieving SOTA prediction accuracy.
Revised: This allows to drastically decrease the parametric complexity while still achieving SOTA prediction accuracy.

Original: We show evidence that increasing the geometric order of SPHCs greatly improves model performance in the low-data regime, as well as generalization to unknown molecules.
Revised: We show evidence that increasing the geometric order of SPHCs greatly improves model performance in the low-data regime, as well as generalization to unknown molecules.

Original: A limitation of the current implementation of S O 3 KRATES is that spherical neighborhoods N χ in eq.
Revised: A limitation of the current implementation of S O 3 KRATES is that spherical neighborhoods N χ in eq.

Original: (14) are computed from all pairwise distances in SPHC space.
Revised: (14) are computed from all pairwise distances in SPHC space.

Original: An alternative implementation could use a space partitioning scheme to find neighborhoods more efficiently.
Revised: An alternative implementation could use a space partitioning scheme to find neighborhoods more efficiently.

Original: In a broader context, our work falls into the category of approaches that can help to reduce the vast computational complexity of molecular and material simulations.
Revised: In a broader context, our work falls into the category of approaches that can help to reduce the vast computational complexity of molecular and material simulations.

Original: This can accelerate novel drug and material designs, which holds the promise of tackling societal challenges, such as climate change and sustainable energy supply [53].
Revised: This can accelerate novel drug and material designs, which holds the promise of tackling societal challenges, such as climate change and sustainable energy supply [53].

Original: Of course, our method could also be used for nefarious applications, e.g. design of chemical warfare, but this is true for all quantum chemistry methods.
Revised: Of course, our method could also be used for nefarious applications, e.g. design of chemical warfare, but this is true for all quantum chemistry methods.

Original: Future research will focus on applications of S O 3 KRATES to materials and bio-molecules, which are typical examples of chemical systems where the accurate description of non-local effects is necessary to produce novel insights.
Revised: Future research will focus on applications of S O 3 KRATES to materials and bio-molecules, which are typical examples of chemical systems where the accurate description of non-local effects is necessary to produce novel insights.

Original: Efficient treatment of non-local effects in point cloud data goes beyond the domain of quantum chemistry.
Revised: Efficient treatment of non-local effects in point cloud data goes beyond the domain of quantum chemistry.

Original: One way of representing non-local dependencies are non-local neural networks [54].
Revised: One way of representing non-local dependencies are non-local neural networks [54].

Original: In comparison to the presented approach they compute a relation in feature rather than in Euclidean space, making it incapable of capturing direct geometric relations in Euclidean space.
Revised: In comparison to the presented approach they compute a relation in feature rather than in Euclidean space, making it incapable of capturing direct geometric relations in Euclidean space.

Original: However, this might be necessary if the relative orientation of objects far apart from each other plays a role for identifying different objects.
Revised: However, this might be necessary if the relative orientation of objects far apart from each other plays a role for identifying different objects.

Original:  Acknowledgements
Revised:  Acknowledgements

Original: All authors acknowledge support by the Federal Ministry of Education and Research (BMBF) for
Revised: All authors acknowledge support by the Federal Ministry of Education and Research (BMBF) for

Original: BIFOLD (01IS18037A).
Revised: BIFOLD (01IS18037A).

Original: KRM was partly supported by the Institute of Information & Communications Technology Planning & Evaluation (IITP) grants funded by the Korea government(MSIT)
Revised: KRM was partly supported by the Institute of Information & Communications Technology Planning & Evaluation (IITP) grants funded by the Korea government(MSIT)

Original: (No. 2019-0-00079, Artificial Intelligence Graduate School Program, Korea University and No.
Revised: (No. 2019-0-00079, Artificial Intelligence Graduate School Program, Korea University and No.

Original: 2022-0-00984, Development of Artificial Intelligence Technology for Personalized Plug-and-Play
Revised: 2022-0-00984, Development of Artificial Intelligence Technology for Personalized Plug-and-Play

Original: Explanation and Verification of Explanation), and was partly supported by the German Ministry for Education and Research (BMBF) under Grants 01IS14013A-E, AIMM, 01GQ1115, 01GQ0850, 01IS18025A and 01IS18037A; the German Research Foundation (DFG).
Revised: Explanation and Verification of Explanation), and was partly supported by the German Ministry for Education and Research (BMBF) under Grants 01IS14013A-E, AIMM, 01GQ1115, 01GQ0850, 01IS18025A and 01IS18037A; the German Research Foundation (DFG).

Original: We thank Stefan Chmiela,
Revised: We thank Stefan Chmiela,

Original: Mihail Bogojeski and Nicklas Schmitz for helpful discussions and feedback on the manuscript.
Revised: Mihail Bogojeski and Nicklas Schmitz for helpful discussions and feedback on the manuscript.

Original: A Appendix A.1 Proof of Equivariance
Revised: A Appendix A.1 Proof of Equivariance

Original: SPHC Initialization
Revised: SPHC Initialization

Original: Here we give proof that certain quantities from the main text are invariant or equivariant, respectively.
Revised: Here we give proof that certain quantities from the main text are invariant or equivariant, respectively.

Original: Let us start with the spherical harmonic coordinates (SPHC) which are initialized as
Revised: Let us start with the spherical harmonic coordinates (SPHC) which are initialized as

Original: [Equation]
Revised: [Equation]

Original: where Y ( l ) : S 2 (cid:55)→ R 2 l +1 .
Revised: where Y ( l ) : S 2 (cid:55)→ R 2 l +1 .

Original: In contrast to the main text, we make the dependence of the right hand side of eq.
Revised: In contrast to the main text, we make the dependence of the right hand side of eq.

Original: (20) on the atomic positions R = { r 1 , . . .
Revised: (20) on the atomic positions R = { r 1 , . . .

Original: , r n } explicit.
Revised: , r n } explicit.

Original: Each degree-wise entry in the initialized SPHCs (eq.
Revised: Each degree-wise entry in the initialized SPHCs (eq.

Original: (20)) transforms under rotation as
Revised: (20)) transforms under rotation as

Original: [Equation]
Revised: [Equation]

Original: ) where D ( R θ ) l ∈ R (2 l +1) × (2 l +1) is the Wigner-D matrix for degree l given rotation R θ and R θ R =
Revised: ) where D ( R θ ) l ∈ R (2 l +1) × (2 l +1) is the Wigner-D matrix for degree l given rotation R θ and R θ R =

Original: { R θ r 1 , . . .
Revised: { R θ r 1 , . . .

Original: , R θ r n } .
Revised: , R θ r n } .

Original: Since the cutoff function ϕ r cut takes the inter-atomic distance r ij = ∥
Revised: Since the cutoff function ϕ r cut takes the inter-atomic distance r ij = ∥

Original: r j − r i ∥as its input argument it is always invariant under rotation.
Revised: r j − r i ∥as its input argument it is always invariant under rotation.

Original: Thus, each degree-wise entry χ ( l ) is equivariant after the initialization.
Revised: Thus, each degree-wise entry χ ( l ) is equivariant after the initialization.

Original: SPHC Message Passing Update After initialization, each per-degree entry in χ is updated as
Revised: SPHC Message Passing Update After initialization, each per-degree entry in χ is updated as

Original: [Equation]
Revised: [Equation]

Original: where α ( l ) ij are rotationally invariant, per-degree self-attention coefficients.
Revised: where α ( l ) ij are rotationally invariant, per-degree self-attention coefficients.

Original: In the first layer, the first part in eq.
Revised: In the first layer, the first part in eq.

Original: (22) corresponds to the initialized SPHC entries, which have been already shown to be equivariant (see above).
Revised: (22) corresponds to the initialized SPHC entries, which have been already shown to be equivariant (see above).

Original: The second part of the equation has the same structural form as the initialization, with the additional self-attention value, which is a rotationally invariant scalar.
Revised: The second part of the equation has the same structural form as the initialization, with the additional self-attention value, which is a rotationally invariant scalar.

Original: Thus we can write
Revised: Thus we can write

Original: [Equation]
Revised: [Equation]

Original: which shows that the updated χ ′ ( l ) is also equivariant.
Revised: which shows that the updated χ ′ ( l ) is also equivariant.

Original: The proof of equivariance for later layers follows analogously.
Revised: The proof of equivariance for later layers follows analogously.

Original: Transformation of the SPHCs After having shown that each per-degree entry of χ transforms under rotation according to the corresponding Wigner-D matrix D l ( R θ ) ∈ R (2 l +1) ×
Revised: Transformation of the SPHCs After having shown that each per-degree entry of χ transforms under rotation according to the corresponding Wigner-D matrix D l ( R θ ) ∈ R (2 l +1) ×

Original: (2 l +1) , one can write the direct sum of all Wigner-D matrices as concatenation of matrices along the diagonal
Revised: (2 l +1) , one can write the direct sum of all Wigner-D matrices as concatenation of matrices along the diagonal

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: such that D ( R θ ) ∈ R ( l max − l min +1) ×
Revised: such that D ( R θ ) ∈ R ( l max − l min +1) ×

Original: ( l max − l min +1) has a block diagonal structure.
Revised: ( l max − l min +1) has a block diagonal structure.

Original: The, full SPHC vectors transform under rotation as
Revised: The, full SPHC vectors transform under rotation as

Original: [Equation]
Revised: [Equation]

Original: As each of the blocks along the diagonal of D ( R θ ) is an orthogonal Wigner-D matrix, D ( R θ ) itself is also orthogonal.
Revised: As each of the blocks along the diagonal of D ( R θ ) is an orthogonal Wigner-D matrix, D ( R θ ) itself is also orthogonal.

Original: A.1.1 Invariance of the (per-degree) Norm
Revised: A.1.1 Invariance of the (per-degree) Norm

Original: The per degree norm is used as input to the spherical filter function ϕ s .
Revised: The per degree norm is used as input to the spherical filter function ϕ s .

Original: As shown above, each of the per-degree entries in χ transforms under rotation as
Revised: As shown above, each of the per-degree entries in χ transforms under rotation as

Original: [Equation]
Revised: [Equation]

Original: The squared norm can be expressed in terms of an inner product
Revised: The squared norm can be expressed in terms of an inner product

Original: [Equation]
Revised: [Equation]

Original: ) where we use the orthogonality of Wigner-D matrices to show that the inner product is rotationally invariant.
Revised: ) where we use the orthogonality of Wigner-D matrices to show that the inner product is rotationally invariant.

Original: If ∥·∥ 22 is invariant, so is ∥·∥ 2 , which completes the proof of equivariance for the degree-wise norm.
Revised: If ∥·∥ 22 is invariant, so is ∥·∥ 2 , which completes the proof of equivariance for the degree-wise norm.

Original: The squared norm of the full SPHC vector χ transforms under rotation as
Revised: The squared norm of the full SPHC vector χ transforms under rotation as

Original: [Equation]
Revised: [Equation]

Original: where we used that the orthogonality of D l results in orthogonality of D (as argued above).
Revised: where we used that the orthogonality of D l results in orthogonality of D (as argued above).

Original: A.2 Spherical Neighborhood
Revised: A.2 Spherical Neighborhood

Original: Starting point for the construction of the spherical neighborhood are the SPHCs X in a given layer of
Revised: Starting point for the construction of the spherical neighborhood are the SPHCs X in a given layer of

Original: S O 3 KRATES .
Revised: S O 3 KRATES .

Original: Consequently, the distance matrix in SPHC for all atomic pairs is given as an n × n matrix X with entries χ ij = ∥
Revised: Consequently, the distance matrix in SPHC for all atomic pairs is given as an n × n matrix X with entries χ ij = ∥

Original: χ i − χ j ∥ 2 .
Revised: χ i − χ j ∥ 2 .

Original: The idea of a spherical neighborhood is to only consider atoms that lie within a certain distance w. r. t. each other in SPHC space, meaning only those for which χ ij ≤ χ cut holds.
Revised: The idea of a spherical neighborhood is to only consider atoms that lie within a certain distance w. r. t. each other in SPHC space, meaning only those for which χ ij ≤ χ cut holds.

Original: However, compared to the inter-atomic distances in Euclidean space, for which specific knowledge e.g. about bond lengths or interaction length scales exits, this is not the case for the rather abstract space of SPHCs.
Revised: However, compared to the inter-atomic distances in Euclidean space, for which specific knowledge e.g. about bond lengths or interaction length scales exits, this is not the case for the rather abstract space of SPHCs.

Original: Thus, we apply a SOFTMAX function along each row (neighborhood) of X , which gives a rescaled version of the spherical distance matrix ˜ X with rescaled entries ˜
Revised: Thus, we apply a SOFTMAX function along each row (neighborhood) of X , which gives a rescaled version of the spherical distance matrix ˜ X with rescaled entries ˜

Original: χ ij ∈ (0 , 1) .
Revised: χ ij ∈ (0 , 1) .

Original: Neighborhoods for each atom are then selected based on a polynomial cutoff function [29]
Revised: Neighborhoods for each atom are then selected based on a polynomial cutoff function [29]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where the input is given as x = ˜ χ ij /χ cut .
Revised: where the input is given as x = ˜ χ ij /χ cut .

Original: Here we chose χ cut
Revised: Here we chose χ cut

Original: = κ/n , where n is the number of atoms in the system and κ allows to increase or decrease the size of the cutoff radius.
Revised: = κ/n , where n is the number of atoms in the system and κ allows to increase or decrease the size of the cutoff radius.

Original: By scaling with the inverse of the number of atoms in the system, we ensure that the relative number of neighboring atoms remains approximately constant, even when going to larger molecules.
Revised: By scaling with the inverse of the number of atoms in the system, we ensure that the relative number of neighboring atoms remains approximately constant, even when going to larger molecules.

Original: The parameter κ allows to fine tune the relative number of atoms and can depended on the size and type of the molecule under investigation.
Revised: The parameter κ allows to fine tune the relative number of atoms and can depended on the size and type of the molecule under investigation.

Original: For our cumulene experiments we chose κ = 1 and p = 6 , which reduces the number of considered interactions by ∼ 40% compared to a global model.
Revised: For our cumulene experiments we chose κ = 1 and p = 6 , which reduces the number of considered interactions by ∼ 40% compared to a global model.

Original: n l l a y e r s
Revised: n l l a y e r s

Original: E
Revised: E

Original: Z
Revised: Z

Original: R
Revised: R

Original: So3krates Layer
Revised: So3krates Layer

Original: Feature initialization
Revised: Feature initialization

Original: SPHC initialization
Revised: SPHC initialization

Original: So3krates Layer So3krates Layer
Revised: So3krates Layer So3krates Layer

Original: MLP Σ j
Revised: MLP Σ j

Original: Feature self-attention +
Revised: Feature self-attention +

Original: Atomwise interaction + + +
Revised: Atomwise interaction + + +

Original: SPHC self-attention
Revised: SPHC self-attention

Original: S o3 k r a t e s l a y e r
Revised: S o3 k r a t e s l a y e r

Original: R F
Revised: R F

Original: F' χ ' χ χ i Δ f i Δ χ
Revised: F' χ ' χ χ i Δ f i Δ χ

Original: i f i
Revised: i f i

Original: Selfmix stack
Revised: Selfmix stack

Original: W MLP split ⊙ + [||· ( l ) || 2 ] l ∈ ℒ ⊙ l + l [||· ( l ) || 2 ] l ∈ ℒ ⊙ l | ℒ |
Revised: W MLP split ⊙ + [||· ( l ) || 2 ] l ∈ ℒ ⊙ l + l [||· ( l ) || 2 ] l ∈ ℒ ⊙ l | ℒ |

Original: F+| ℒ |
Revised: F+| ℒ |

Original: F | ℒ |
Revised: F | ℒ |

Original: F 2F+2| ℒ |
Revised: F 2F+2| ℒ |

Original: F
Revised: F

Original: F | ℒ | | ℒ | A t o m w i s e i n t e r a c t i o n
Revised: F | ℒ | | ℒ | A t o m w i s e i n t e r a c t i o n

Original: Filtermodule ⊙ ⟨ · i ,· j ⟩ ϕ cut Σ
Revised: Filtermodule ⊙ ⟨ · i ,· j ⟩ ϕ cut Σ

Original: j
Revised: j

Original: W q
Revised: W q

Original: W k ⊙
Revised: W k ⊙

Original: Y lm for all j ∈ 𝒩 i f i f
Revised: Y lm for all j ∈ 𝒩 i f i f

Original: j r i , r j χ i , χ j Δ χ
Revised: j r i , r j χ i , χ j Δ χ

Original: i
Revised: i

Original: SPHC self-attention
Revised: SPHC self-attention

Original: Filtermodule ⊙ ⟨ · i ,· j ⟩ ϕ cut Σ
Revised: Filtermodule ⊙ ⟨ · i ,· j ⟩ ϕ cut Σ

Original: j
Revised: j

Original: W q
Revised: W q

Original: W v W k f i f j r i , r j χ i , χ j for all j ∈ 𝒩 i Δ f
Revised: W v W k f i f j r i , r j χ i , χ j for all j ∈ 𝒩 i Δ f

Original: i
Revised: i

Original: Feature self-attention || χ i(l) - χ j(l) ||
Revised: Feature self-attention || χ i(l) - χ j(l) ||

Original: MLP stack l || r i - r j ||
Revised: MLP stack l || r i - r j ||

Original: RBF ϕ cut
Revised: RBF ϕ cut

Original: MLP + r i , r j χ i , χ j w ij
Revised: MLP + r i , r j χ i , χ j w ij

Original: Filter module interaction module) illustrates that each degree is multiplied with one entry of the |L| -sized output vector.
Revised: Filter module interaction module) illustrates that each degree is multiplied with one entry of the |L| -sized output vector.

Original: A.3 Network Architecture
Revised: A.3 Network Architecture

Original: The S O 3 KRATES model is available at https://github.com/thorben-frank/mlff .
Revised: The S O 3 KRATES model is available at https://github.com/thorben-frank/mlff .

Original: It makes extensive use of the libraries F LAX [55] for building the networks, O PTAX [56] for training and
Revised: It makes extensive use of the libraries F LAX [55] for building the networks, O PTAX [56] for training and

Original: N UMPY [57] and J AX [58] for additional processing steps.
Revised: N UMPY [57] and J AX [58] for additional processing steps.

Original: All code and data that is necessary to produce the results step by step presented in the main body of the paper can be downloaded from https://zenodo.org/record/6584855 (newest version).
Revised: All code and data that is necessary to produce the results step by step presented in the main body of the paper can be downloaded from https://zenodo.org/record/6584855 (newest version).

Original: Detailed Network Design We now describe the specific implementation of the building blocks that we have formally introduced in the main text.
Revised: Detailed Network Design We now describe the specific implementation of the building blocks that we have formally introduced in the main text.

Original: The computational flow of the architecture is shown in
Revised: The computational flow of the architecture is shown in

Original: Fig.
Revised: Fig.

Original: Radial Filter For the radial filter function ϕ
Revised: Radial Filter For the radial filter function ϕ

Original: r : R (cid:55)→ R F (first summand in eq.
Revised: r : R (cid:55)→ R F (first summand in eq.

Original: the interatomic distances are expanded in terms of K radial basis functions as proposed in P HYS N ET [19], which are given as
Revised: the interatomic distances are expanded in terms of K radial basis functions as proposed in P HYS N ET [19], which are given as

Original: [Equation]
Revised: [Equation]

Original: where µ k is the center of the k -th basis function ( k = 1 . . .
Revised: where µ k is the center of the k -th basis function ( k = 1 . . .

Original: K ).
Revised: K ).

Original: The cutoff function
Revised: The cutoff function

Original: [Equation]
Revised: [Equation]

Original: guarantees that θ k
Revised: guarantees that θ k

Original: ( r ij )
Revised: ( r ij )

Original: = 0 smoothly goes to zero when interatomic distances exceed the cutoff radius r cut .
Revised: = 0 smoothly goes to zero when interatomic distances exceed the cutoff radius r cut .

Original: Here we chose K = 32 basis functions in total.
Revised: Here we chose K = 32 basis functions in total.

Original: The expanded distance is passed in a two-layered MLP with SILU non-linearity in-between and 128 hidden and F output neurons, respectively.
Revised: The expanded distance is passed in a two-layered MLP with SILU non-linearity in-between and 128 hidden and F output neurons, respectively.

Original: Spherical Filter
Revised: Spherical Filter

Original: The spherical filter function ϕ s (cid:16) [ χ ( l ) ij ] l ∈L (cid:17) : R |L| (cid:55)→ R F (second summand in eq.
Revised: The spherical filter function ϕ s (cid:16) [ χ ( l ) ij ] l ∈L (cid:17) : R |L| (cid:55)→ R F (second summand in eq.

Original: (16)) is also modelled as a two-layered MLP with the same non-linearity.
Revised: (16)) is also modelled as a two-layered MLP with the same non-linearity.

Original: Its input is given as the stacked per-degree distances χ ( l ) ij , such that the input to the MLP is of low dimension (number of spherical degrees l ).
Revised: Its input is given as the stacked per-degree distances χ ( l ) ij , such that the input to the MLP is of low dimension (number of spherical degrees l ).

Original: For that reason, we only use 32 neurons in the first and again F neurons in the second layer.
Revised: For that reason, we only use 32 neurons in the first and again F neurons in the second layer.

Original: Self-Attention
Revised: Self-Attention

Original: The self-attention matrix (cf. eq.
Revised: The self-attention matrix (cf. eq.

Original: (15)) is calculated from the filter vector w ij as well as from a pair of feature vectors f i ∈ R F and f j ∈ R F .
Revised: (15)) is calculated from the filter vector w ij as well as from a pair of feature vectors f i ∈ R F and f j ∈ R F .

Original: Before being passed to the inner product, each of the feature vectors is refined using a linear layer which we denote as W q and W k in the fashion of the key and query matrices usually appearing in the calculation of self-attention coefficients [49].
Revised: Before being passed to the inner product, each of the feature vectors is refined using a linear layer which we denote as W q and W k in the fashion of the key and query matrices usually appearing in the calculation of self-attention coefficients [49].

Original: Thus, the self-attention coefficients are calculated as
Revised: Thus, the self-attention coefficients are calculated as

Original: [Equation]
Revised: [Equation]

Original: (32) where q i = W q f i ∈ R F , k j = W k f j ∈ R F and w ij ∈ R F comes from the filter module.
Revised: (32) where q i = W q f i ∈ R F , k j = W k f j ∈ R F and w ij ∈ R F comes from the filter module.

Original: The features f j of the neighboring atoms are also passed through an additional linear layer W v .
Revised: The features f j of the neighboring atoms are also passed through an additional linear layer W v .

Original: The updated features are then given as
Revised: The updated features are then given as

Original: [Equation]
Revised: [Equation]

Original: and the updates to the SPHCs as
Revised: and the updates to the SPHCs as

Original: [Equation]
Revised: [Equation]

Original: Different parameters are used for the feature and SPHC updates, respectively.
Revised: Different parameters are used for the feature and SPHC updates, respectively.

Original: For the feature update, we use a predefined number of heads whereas the number of heads in the SPHC update equals the number of degrees in the SPHC vector.
Revised: For the feature update, we use a predefined number of heads whereas the number of heads in the SPHC update equals the number of degrees in the SPHC vector.

Original: In order to ensure permutation invariance, all parameters of the linear layers are shared across atoms.
Revised: In order to ensure permutation invariance, all parameters of the linear layers are shared across atoms.

Original: Atomwise Interaction After the update MP step, we update the features as well as the SPHCs per atom.
Revised: Atomwise Interaction After the update MP step, we update the features as well as the SPHCs per atom.

Original: In this step, we do not only include cross-degree coupling in χ but also allow for information exchange between the feature and the SPHC branch.
Revised: In this step, we do not only include cross-degree coupling in χ but also allow for information exchange between the feature and the SPHC branch.

Original: The functions ϕ 1 (cid:16) f i , [ χ ( l ) i ] l ∈L , [˜ χ ( l ) i ] l ∈L (cid:17) and ϕ ( l ) 2 (cid:16) f i , [ χ ( l ) i ] l ∈L , [˜ χ ( l ) i ] l ∈L (cid:17) are implemented by a shared MLP.
Revised: The functions ϕ 1 (cid:16) f i , [ χ ( l ) i ] l ∈L , [˜ χ ( l ) i ] l ∈L (cid:17) and ϕ ( l ) 2 (cid:16) f i , [ χ ( l ) i ] l ∈L , [˜ χ ( l ) i ] l ∈L (cid:17) are implemented by a shared MLP.

Original: The function ϕ ( l ) 3 (cid:16) [˜ χ ( l ) i ] l ∈L (cid:17) is implemented by a single linear layer, without bias term.
Revised: The function ϕ ( l ) 3 (cid:16) [˜ χ ( l ) i ] l ∈L (cid:17) is implemented by a single linear layer, without bias term.

Original: The full computational flow is shown in
Revised: The full computational flow is shown in

Original: Fig.
Revised: Fig.

Original: A.4 Ablation Study Spherical Filter
Revised: A.4 Ablation Study Spherical Filter

Original: To illustrate the importance of the spherical filter, we examine its effect in an ablation study on the
Revised: To illustrate the importance of the spherical filter, we examine its effect in an ablation study on the

Original: MD17 benchmark for varying maximal degree l max .
Revised: MD17 benchmark for varying maximal degree l max .

Original: As can be seen in Fig.
Revised: As can be seen in Fig.

Original: 6, using spherical filters (see eq.
Revised: 6, using spherical filters (see eq.

Original: (16)) improves performance compared to a S O 3 KRATES model without them (solid vs. dotted lines), where the difference becomes even more pronounced for force predictions.
Revised: (16)) improves performance compared to a S O 3 KRATES model without them (solid vs. dotted lines), where the difference becomes even more pronounced for force predictions.

Original: Further, the effect of spherical filters becomes stronger for smaller l max .
Revised: Further, the effect of spherical filters becomes stronger for smaller l max .

Original: Since many equivariant MPNNs carry features including additional channels for higher order geometric information, spherical filters can be straight forwardly integrated into current architectures, offering the potential of increased accuracy.
Revised: Since many equivariant MPNNs carry features including additional channels for higher order geometric information, spherical filters can be straight forwardly integrated into current architectures, offering the potential of increased accuracy.

Original: [Figure]
Revised: [Figure]

Original: Aspirin
Revised: Aspirin

Original: Ethanol
Revised: Ethanol

Original: Malondialdehyde
Revised: Malondialdehyde

Original: Naphthalene
Revised: Naphthalene

Original: Salicylic Acid
Revised: Salicylic Acid

Original: Toluene
Revised: Toluene

Original: Uracil with spherical ﬁ lter no spherical ﬁ lter
Revised: Uracil with spherical ﬁ lter no spherical ﬁ lter

Original: [Table]
Revised: [Table]

Original: A.5 QM7-X Experiments
Revised: A.5 QM7-X Experiments

Original: As an additional benchmark
Revised: As an additional benchmark

Original: [Figure]
Revised: [Figure]

Original: we train S O 3 KRATES on the full QM7-X data set and compare our results to the reported errors for known and unknown molecules in table 2.
Revised: we train S O 3 KRATES on the full QM7-X data set and compare our results to the reported errors for known and unknown molecules in table 2.

Original: The known molecules correspond to the test samples that have not been seen during training.
Revised: The known molecules correspond to the test samples that have not been seen during training.

Original: Following [30], we use the keys ( IDMOL in the QM7-X data base) [1771, 1805, 1824, 2020, 2085, 2117, 3019, 3108, 3190, 3217, 3257, 3329, 3531, 4010, 4181, 4319, 4713, 5174, 5370, 5580, 5891, 6315, 6583, 6809, 7020] for testing the generalization of our model to unknown structures, which have been excluded from the data set prior to generating train, test and validation splits.
Revised: Following [30], we use the keys ( IDMOL in the QM7-X data base) [1771, 1805, 1824, 2020, 2085, 2117, 3019, 3108, 3190, 3217, 3257, 3329, 3531, 4010, 4181, 4319, 4713, 5174, 5370, 5580, 5891, 6315, 6583, 6809, 7020] for testing the generalization of our model to unknown structures, which have been excluded from the data set prior to generating train, test and validation splits.

Original: Accuracy Comparison Local vs. Non-Local Model
Revised: Accuracy Comparison Local vs. Non-Local Model

Original: By comparing the results reported in tablewe find, that the non-local model shows the same accuracy as the local model for both known and unknown molecules.
Revised: By comparing the results reported in tablewe find, that the non-local model shows the same accuracy as the local model for both known and unknown molecules.

Original: This underlines the applicability of the presented, non-local corrections to a large variety of molecular structures including transfer learning to unknown molecules.
Revised: This underlines the applicability of the presented, non-local corrections to a large variety of molecular structures including transfer learning to unknown molecules.

Original: Accuracy Improvement by Up-Scaling Despite the parametric leightweight structure of
Revised: Accuracy Improvement by Up-Scaling Despite the parametric leightweight structure of

Original: S O 3 KRATES , it is important that results can be imrproved by upscaling the model.
Revised: S O 3 KRATES , it is important that results can be imrproved by upscaling the model.

Original: Here we take one of the most straight forward paths and upscale the model by simply increasing the feature dimension from F = 132 to F = 164 .
Revised: Here we take one of the most straight forward paths and upscale the model by simply increasing the feature dimension from F = 132 to F = 164 .

Original: As it can be seen in table 2, this already allows to increase the accuracy on the QM7-X dataset compared to the base line model.
Revised: As it can be seen in table 2, this already allows to increase the accuracy on the QM7-X dataset compared to the base line model.

Original: It should be noted, however, that one could follow additional/other directions, such as increasing the maximal degree l max .
Revised: It should be noted, however, that one could follow additional/other directions, such as increasing the maximal degree l max .

Original: Generalization to Larger Molecules with a Non-Local Model As seen in table 2, a model with non-local correction is capable of generalizing to structures not seen during training.
Revised: Generalization to Larger Molecules with a Non-Local Model As seen in table 2, a model with non-local correction is capable of generalizing to structures not seen during training.

Original: However, the re-usability of local information is one of the key assumptions for building models that can be trained on small molecules and afterwards be applied to larger, completely unknown structures [30].
Revised: However, the re-usability of local information is one of the key assumptions for building models that can be trained on small molecules and afterwards be applied to larger, completely unknown structures [30].

Original: Thus, testing the non-local model on unknwon structures from the QM7-X data set is an insufficient test to investigate this property of transfer learning.
Revised: Thus, testing the non-local model on unknwon structures from the QM7-X data set is an insufficient test to investigate this property of transfer learning.

Original: We therefore use the non-local model for geometry optimization of molecules that range from 47 (riboflavin) up to 109 atoms (Ala10) in size (see 7).
Revised: We therefore use the non-local model for geometry optimization of molecules that range from 47 (riboflavin) up to 109 atoms (Ala10) in size (see 7).

Original: We find that the non-local model is robust for geometry optimization when being applied to much larger, unknown molecules.
Revised: We find that the non-local model is robust for geometry optimization when being applied to much larger, unknown molecules.

Original: The largest molecule in the QM7-X data set has 23 atoms.
Revised: The largest molecule in the QM7-X data set has 23 atoms.

Original: Ribo ﬂ avin
Revised: Ribo ﬂ avin

Original: Cholesterin
Revised: Cholesterin

Original: Ala # atoms 

Revised: Ala # atoms

Original: RMSD: 0.154Å
Revised: RMSD: 0.154Å

Original: RMSD:
Revised: RMSD:

Original: 0.144Å
Revised: 0.144Å

Original: RMSD:
Revised: RMSD:

Original: 0.343Å
Revised: 0.343Å

Original: S o3 k r a t e s o p t
Revised: S o3 k r a t e s o p t

Original: i m i z e d s t r u c t u r e s
Revised: i m i z e d s t r u c t u r e s

Original: R e f e r e n c e
Revised: R e f e r e n c e

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: A.6 Time Analysis
Revised: A.6 Time Analysis

Original: In order to determine the training and inference times we follow [34] and evaluate our model on the toluene molecule from the MD17 benchmark with a batch size of 4.
Revised: In order to determine the training and inference times we follow [34] and evaluate our model on the toluene molecule from the MD17 benchmark with a batch size of 4.

Original: We trained a N EQU IP model with the same hyperparameters as reported in [32] as well as a S O 3 KRATES model on a Tesla P with 12GB.
Revised: We trained a N EQU IP model with the same hyperparameters as reported in [32] as well as a S O 3 KRATES model on a Tesla P with 12GB.

Original: The reported training time corresponds to the wall time it took each model to evaluate a single gradient update (without time for validation).
Revised: The reported training time corresponds to the wall time it took each model to evaluate a single gradient update (without time for validation).

Original: We compare these runtimes to the runtimes that have been reported for D IME N ET [29] and G EM N ET Q [34] in [34].
Revised: We compare these runtimes to the runtimes that have been reported for D IME N ET [29] and G EM N ET Q [34] in [34].

Original: However, these reported times have been measured on a GeForce GTX 1080Ti (a GPU we did not have access to).
Revised: However, these reported times have been measured on a GeForce GTX 1080Ti (a GPU we did not have access to).

Original: In [37] it has been found that a Tesla P100 gives a speedup factor of ∼ 1 . 3 , such that we downscale the reported runtimes accordingly.
Revised: In [37] it has been found that a Tesla P100 gives a speedup factor of ∼ 1 . 3 , such that we downscale the reported runtimes accordingly.

Original: The resulting times are shown in table 3, which are the values plotted in 1d.
Revised: The resulting times are shown in table 3, which are the values plotted in 1d.

Original: It should be noted, that our implementation did not focus on the runtime, such that it is likely to be possible to further reduce the computational cost that is required for training and inference.
Revised: It should be noted, that our implementation did not focus on the runtime, such that it is likely to be possible to further reduce the computational cost that is required for training and inference.

Original: A.7 Analysis of Attention Coefficients
Revised: A.7 Analysis of Attention Coefficients

Original: In figure 8 we plot the attention coefficients from the MP update of the SPHCs (cf. eq.
Revised: In figure 8 we plot the attention coefficients from the MP update of the SPHCs (cf. eq.

Original: (14)) after training for different dihedral angles between the rotors.
Revised: (14)) after training for different dihedral angles between the rotors.

Original: Attention values are calculated as the average over all attention values obtained for a given atomic pair throughout all SPHC update steps.
Revised: Attention values are calculated as the average over all attention values obtained for a given atomic pair throughout all SPHC update steps.

Original: As it can be verified visually, the model picks up physically important interactions.
Revised: As it can be verified visually, the model picks up physically important interactions.

Original: Care should be taken,
Revised: Care should be taken,

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: however, since it has been pointed out in [59] that looking at the bare attention values only has limited expressiveness.
Revised: however, since it has been pointed out in [59] that looking at the bare attention values only has limited expressiveness.

Original: A.8 Scaling Analysis
Revised: A.8 Scaling Analysis

Original: We further compare the scaling of three different versions of the S O 3 KRATES model.
Revised: We further compare the scaling of three different versions of the S O 3 KRATES model.

Original: Namely, we compare a local S O 3 KRATES model, which only localizes in R 3 (used in 4.1), a non-local S O 3 KRATES model with non-local corrections from SPHC space (used in 4.2) and a fully global model for which the cutoff radius is chosen such that all atoms are in each others neighborhood.
Revised: Namely, we compare a local S O 3 KRATES model, which only localizes in R 3 (used in 4.1), a non-local S O 3 KRATES model with non-local corrections from SPHC space (used in 4.2) and a fully global model for which the cutoff radius is chosen such that all atoms are in each others neighborhood.

Original: The inference times for energy predictions for molecules ranging from 9 up to 370 atoms are shown in Fig.
Revised: The inference times for energy predictions for molecules ranging from 9 up to 370 atoms are shown in Fig.

Original: As expected, we find linear scaling for the local S O 3 KRATES version with a remarkable inference time of only ∼ 196 ms for 370 atoms (batch size is 25).
Revised: As expected, we find linear scaling for the local S O 3 KRATES version with a remarkable inference time of only ∼ 196 ms for 370 atoms (batch size is 25).

Original: For the fully global mode, we see quadratic scaling in the number of atoms.
Revised: For the fully global mode, we see quadratic scaling in the number of atoms.

Original: The model with non-local corrections can be found somewhere in-between the local and the global model.
Revised: The model with non-local corrections can be found somewhere in-between the local and the global model.

Original: To that end, we want to stretch the fact that we did not focus on an efficient implementation for the non-local corrections.
Revised: To that end, we want to stretch the fact that we did not focus on an efficient implementation for the non-local corrections.

Original: The experiments were run on a Tesla P with 12 Gb and a batch size of 25.
Revised: The experiments were run on a Tesla P with 12 Gb and a batch size of 25.

Original: A.9 QM7-X 
Revised: A.9 QM7-X

Original: As starting point for the recently introduced QM7-X dataset [51] serve ∼ 7 k molecular graphs with a maximum of 7 non-hydrogen atoms (C, N, O, S, Cl).
Revised: As starting point for the recently introduced QM7-X dataset [51] serve ∼ 7 k molecular graphs with a maximum of 7 non-hydrogen atoms (C, N, O, S, Cl).

Original: By sampling and optimizing structural and constitutional isomers for each graph, ∼ 42 k equilibrium structures are generated.
Revised: By sampling and optimizing structural and constitutional isomers for each graph, ∼ 42 k equilibrium structures are generated.

Original: Using normal mode sampling at 1500 K, 100 out-of-equilibrium points are generated for each structure resulting in 101 data points per structure and ∼ 4 .
Revised: Using normal mode sampling at 1500 K, 100 out-of-equilibrium points are generated for each structure resulting in 101 data points per structure and ∼ 4 .

Original: 2 M geometries in total.
Revised: 2 M geometries in total.

Original: In order to make the data set well suited for both, kernel and neural network models, we group the geometries by structural isomers which gives ∼ 13 k individual data sets, each consisting of #stereo-isomers × 101 geometries.
Revised: In order to make the data set well suited for both, kernel and neural network models, we group the geometries by structural isomers which gives ∼ 13 k individual data sets, each consisting of #stereo-isomers × 101 geometries.

Original: For each of the data sets we choose 80 points for training,points for validation and the remaining points for testing after the training.
Revised: For each of the data sets we choose 80 points for training,points for validation and the remaining points for testing after the training.

Original: Afterwards, we randomly sample 250 data sets out of the 13k data sets.
Revised: Afterwards, we randomly sample 250 data sets out of the 13k data sets.

Original: The comparison of the probabilities of drawing a molecule with a given number of atoms, number of symmetries and number of stereo-isomers from the original QM7-X and the QM7-X250 data sets are shown in Fig.
Revised: The comparison of the probabilities of drawing a molecule with a given number of atoms, number of symmetries and number of stereo-isomers from the original QM7-X and the QM7-X250 data sets are shown in Fig.

Original: Since we ensure that each of the structural subsets present in the original data set is also present in the 250 drawn samples at least once, it can be seen that these structures are over represented in the QM7-X250 data set, even though
Revised: Since we ensure that each of the structural subsets present in the original data set is also present in the 250 drawn samples at least once, it can be seen that these structures are over represented in the QM7-X250 data set, even though

Original: [Figure]
Revised: [Figure]

Original: the distribution of the original dataset.
Revised: the distribution of the original dataset.

Original: As for each property value at least once structure is present in QM7-X250 (per design), values with a low count in the original data have higher relative importance, which can be seen e.g. for structures with only a few
Revised: As for each property value at least once structure is present in QM7-X250 (per design), values with a low count in the original data have higher relative importance, which can be seen e.g. for structures with only a few

Original: , or decreases, since in the SPHC step, many used as there are degrees (one degree for l max
Revised: , or decreases, since in the SPHC step, many used as there are degrees (one degree for l max

Original: = 0 and l max = 1 ).
Revised: = 0 and l max = 1 ).

Original: As a consequence, the size of the matrices Q and K , that are applied to the feature vectors that make up the inner product, become smaller in size.
Revised: As a consequence, the size of the matrices Q and K , that are applied to the feature vectors that make up the inner product, become smaller in size.

Original: The feature dimension F = 132 is the same for all models.
Revised: The feature dimension F = 132 is the same for all models.

Original: This leads to a decrease of the total number of network parameters, which is larger than the increase in network parameters, due to the additional degrees.
Revised: This leads to a decrease of the total number of network parameters, which is larger than the increase in network parameters, due to the additional degrees.

Original: Note that it is important to take the average over per structure MAEs, since the different structures have differently many test samples (see section A.9).
Revised: Note that it is important to take the average over per structure MAEs, since the different structures have differently many test samples (see section A.9).

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: they only make up a single structure.
Revised: they only make up a single structure.

Original: Apart from these special structures, we see that the sampled dataset correctly reproduces the distribution of the original data set.
Revised: Apart from these special structures, we see that the sampled dataset correctly reproduces the distribution of the original data set.

Original: From the 250 sampled structures, we build 250 per structure data sets with 80 training, 10 validation and 10–3748 (depending on the number of stereo-isomers) test points, which we referred to as individual dataset in the main text.
Revised: From the 250 sampled structures, we build 250 per structure data sets with 80 training, 10 validation and 10–3748 (depending on the number of stereo-isomers) test points, which we referred to as individual dataset in the main text.

Original: We can further build a joint dataset, by merging all individual structures into a single data set, which gives 20 000 ( 250 × 80 ) training, 2 500 ( 250 × 10 ) validation, and 108 800 testing points.
Revised: We can further build a joint dataset, by merging all individual structures into a single data set, which gives 20 000 ( 250 × 80 ) training, 2 500 ( 250 × 10 ) validation, and 108 800 testing points.

Original: A.10 QM7-X250 Experiments
Revised: A.10 QM7-X250 Experiments

Original: In the main text we show the force MAE as a function of maximal order l max in Fig. 4a.
Revised: In the main text we show the force MAE as a function of maximal order l max in Fig. 4a.

Original: In table 4, we further show the exact errors for energy and forces as a function of maximal order l max , as well as the number of parameters per S O 3 KRATES model.
Revised: In table 4, we further show the exact errors for energy and forces as a function of maximal order l max , as well as the number of parameters per S O 3 KRATES model.

Original: For all models with l max ≥ 1 , we do not include l max = 0 within the SPHCs, since the zeroth degree evaluates to a constant one and thus does not contain any additional geometric information.
Revised: For all models with l max ≥ 1 , we do not include l max = 0 within the SPHCs, since the zeroth degree evaluates to a constant one and thus does not contain any additional geometric information.

Original: The generalization to unknown molecules is tested on the same structure keys as in the S POOKY N ET paper [30] and in the QM7-X experiment from above A.5.
Revised: The generalization to unknown molecules is tested on the same structure keys as in the S POOKY N ET paper [30] and in the QM7-X experiment from above A.5.

Original: As we only trained on forces and generalize to completely unknown molecules, we can not fit the energy integration constant (as we assume to have no reference data for the unknown molecules).
Revised: As we only trained on forces and generalize to completely unknown molecules, we can not fit the energy integration constant (as we assume to have no reference data for the unknown molecules).

Original: For that reason we only report force errors (see 5).
Revised: For that reason we only report force errors (see 5).

Original: However, one could train a S O 3 KRATES model on both, energy and forces to obtain meaningful predictions for both.
Revised: However, one could train a S O 3 KRATES model on both, energy and forces to obtain meaningful predictions for both.

Original: In that case, atomization energies would need to be included to obtain equal energy scales across different molecules as done for the full QM7-X data set.
Revised: In that case, atomization energies would need to be included to obtain equal energy scales across different molecules as done for the full QM7-X data set.

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

Original: A.11 Additional Experiments for Non-Local Effects
Revised: A.11 Additional Experiments for Non-Local Effects

Original: Here we demonstrate, that increasing the number of local MP steps in N EQU IP allows to model the non-local effects in cumulene, as an increase of the number of steps increases the effective cutoff of the model (see Fig. 11).
Revised: Here we demonstrate, that increasing the number of local MP steps in N EQU IP allows to model the non-local effects in cumulene, as an increase of the number of steps increases the effective cutoff of the model (see Fig. 11).

Original: However, this just shifts the problem to larger distances; increasing the length of the cumulene molecule again leads to a scenario where the local model fails.
Revised: However, this just shifts the problem to larger distances; increasing the length of the cumulene molecule again leads to a scenario where the local model fails.

Original: A.12 Benchmarks for Non-Local Effects
Revised: A.12 Benchmarks for Non-Local Effects

Original: We further apply S O 3 KRATES to the benchmark presented in [60], which explicitly introduces structures which exhibit non-local effects.
Revised: We further apply S O 3 KRATES to the benchmark presented in [60], which explicitly introduces structures which exhibit non-local effects.

Original: We find, that for the carbon chain with non-local charge transfer, S O 3 KRATES with non-local corrections improves by a factor ∼ 2 compared to a local model.
Revised: We find, that for the carbon chain with non-local charge transfer, S O 3 KRATES with non-local corrections improves by a factor ∼ 2 compared to a local model.

Original: For the remaining structures no such difference can be found.
Revised: For the remaining structures no such difference can be found.

Original: Notably, the non-local effects from the presented benchmark do not result in global geometric relations but rather change the overall behavior by adding or removing certain atoms (called doping).
Revised: Notably, the non-local effects from the presented benchmark do not result in global geometric relations but rather change the overall behavior by adding or removing certain atoms (called doping).

Original: Thus, a combination of the S O 3 KRATES mechanism (for geometric relations) and the mechanism presented in S POOKY N ET might be a valuable direction for future work.
Revised: Thus, a combination of the S O 3 KRATES mechanism (for geometric relations) and the mechanism presented in S POOKY N ET might be a valuable direction for future work.

Original: Full results are reported in table 6.
Revised: Full results are reported in table 6.

Original: Note that S O 3 KRATES could be easily extended to predict partial charges as well following e.g. [30].
Revised: Note that S O 3 KRATES could be easily extended to predict partial charges as well following e.g. [30].

Original: A.13 Training Details
Revised: A.13 Training Details

Original: We train S O 3 KRATES by minimizing a combined loss of energy and forces
Revised: We train S O 3 KRATES by minimizing a combined loss of energy and forces

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: True 4k 1k 1k True 4k 1k 1k
Revised: True 4k 1k 1k True 4k 1k 1k

Original: Fig
Revised: Fig

Original: 4a   
Revised: 4a

Original: False
Revised: False

Original: True 1.5k
Revised: True 1.5k

Original: Fig
Revised: Fig

Original: False
Revised: False

Original: True 6k 20k 2.5k
Revised: True 6k 20k 2.5k

Original: Tab.
Revised: Tab.

Original: False
Revised: False

Original: True 0.
Revised: True 0.

Original: 4k 1k 1k
Revised: 4k 1k 1k

Original: Fig
Revised: Fig

Original: False
Revised: False

Original: True / False 0.
Revised: True / False 0.

Original: 4k 1k 1k
Revised: 4k 1k 1k

Original: Fig
Revised: Fig

Original: True
Revised: True

Original: True 0.
Revised: True 0.

Original: 1k 100 3.6M 360k
Revised: 1k 100 3.6M 360k

Original: Tab.
Revised: Tab.

Original: True/False
Revised: True/False

Original: True 0.
Revised: True 0.

Original: 1k 100 3.6M 360k
Revised: 1k 100 3.6M 360k

Original: Tab
Revised: Tab

Original: False
Revised: False

Original: True 0.
Revised: True 0.

Original: 1k 100 3.6M 360k
Revised: 1k 100 3.6M 360k

Original: Tab.
Revised: Tab.

Original: True/False
Revised: True/False

Original: True 0.
Revised: True 0.

Original: 1k * * where ˜ E and ˜ F are the ground truth and E and F are the predictions of the model.
Revised: 1k * * where ˜ E and ˜ F are the ground truth and E and F are the predictions of the model.

Original: The loss is evaluated on mini batches with a batch size given in table Tab. 7.
Revised: The loss is evaluated on mini batches with a batch size given in table Tab. 7.

Original: The parameter β is used to control the trade-off between energy and forces and additionally accounts for different energy and force scales.
Revised: The parameter β is used to control the trade-off between energy and forces and additionally accounts for different energy and force scales.

Original: We train our models with the ADAM optimizer [61] and an initial learning rate of µ = 1 × 10 − 3 .
Revised: We train our models with the ADAM optimizer [61] and an initial learning rate of µ = 1 × 10 − 3 .

Original: We use exponential learning rate decay where the learning rate is decreased by a factor of 0 .
Revised: We use exponential learning rate decay where the learning rate is decreased by a factor of 0 .

Original: 5 every 1k epochs for the MD17 benchmark and the joint QM7-X250 dataset and every 300 epochs for individual models on the QM7-X250 dataset.
Revised: 5 every 1k epochs for the MD17 benchmark and the joint QM7-X250 dataset and every 300 epochs for individual models on the QM7-X250 dataset.

Original: For training on the full QM7-X data set we reduced the leraning rate every 250 epochs by a factor of 0.7.
Revised: For training on the full QM7-X data set we reduced the leraning rate every 250 epochs by a factor of 0.7.

Original: We further applied gradient clipping to a maximal norm of 1.
Revised: We further applied gradient clipping to a maximal norm of 1.

Original: Additional hyperparameters that have been used to produce the tables and figures in this work are given Tab. 7.
Revised: Additional hyperparameters that have been used to produce the tables and figures in this work are given Tab. 7.

Original: Whenever β = 1 is reported, no energy contribution did enter the loss function.
Revised: Whenever β = 1 is reported, no energy contribution did enter the loss function.

Original: In that case, we calculated the integration constant for energy according to section A.14.
Revised: In that case, we calculated the integration constant for energy according to section A.14.

Original: A.14 Energy Integration Constant
Revised: A.14 Energy Integration Constant

Original: When only training on forces, the resulting energy predictions are likely to be shifted w. r. t. the correct energy values, due to vanishing constants when taking the derivative.
Revised: When only training on forces, the resulting energy predictions are likely to be shifted w. r. t. the correct energy values, due to vanishing constants when taking the derivative.

Original: Since force fields are conservative vector fields, one can define the following loss for the constant c as
Revised: Since force fields are conservative vector fields, one can define the following loss for the constant c as

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where index i runs over the M data points, R i ∈ R n × 3 are the atomic coordinates and E
Revised: where index i runs over the M data points, R i ∈ R n × 3 are the atomic coordinates and E

Original: i is the reference value of the PES.
Revised: i is the reference value of the PES.

Original: The functions f F : R n × 3 (cid:55)→ R n × 3 and f E : R n × 3 (cid:55)→ R are the force and energy function, respectively.
Revised: The functions f F : R n × 3 (cid:55)→ R n × 3 and f E : R n × 3 (cid:55)→ R are the force and energy function, respectively.

Original: Minimization w. r. t. to c then gives
Revised: Minimization w. r. t. to c then gives

Original: [Equation]
Revised: [Equation]

Original: Thus, the shifted energy function is given as ˜ f E ( R i ) =
Revised: Thus, the shifted energy function is given as ˜ f E ( R i ) =

Original: c + f E ( R i ) .
Revised: c + f E ( R i ) .