Smooth, exact rotational symmetrization for deep learning on point clouds
Authors: Sergey Pozdnyakov
Michele Ceriotti
Venue: NeurIPS 2023 poster (2023)
Keywords: geometric deep learning
point clouds
equivariance
machine learning potentials
GNN
transformer
atomic-scale modeling
TL;DR: We introduce a general method to enforce rotational symmetry in existing deep-learning models for point clouds, and show with a model we name PET how this simplifies the development of accurate models for problems with exact symmetry constraints.

Abstract: Point clouds are versatile representations of 3D objects and have found widespread application in science and engineering. Many successful deep-learning models have been proposed that use them as input. The domain of chemical and materials modeling is especially challenging because exact compliance with physical constraints is highly desirable for a model to be usable in practice. These constraints include smoothness and invariance with respect to translations, rotations, and permutations of identical atoms. If these requirements are not rigorously fulfilled, atomistic simulations might lead to absurd outcomes even if the model has excellent accuracy. Consequently, dedicated architectures, which achieve invariance by restricting their design space, have been developed. General-purpose point-cloud models are more varied but often disregard rotational symmetry. We propose a general symmetrization method that adds rotational equivariance to any given model while preserving all the other requirements.
Our approach simplifies the development of better atomic-scale machine-learning schemes by relaxing the constraints on the design space and making it possible to incorporate ideas that proved effective in other domains.
We demonstrate this idea by introducing the Point Edge Transformer (PET) architecture, which is not intrinsically equivariant but achieves state-of-the-art performance on several benchmark datasets of molecules and solids. A-posteriori application of our general protocol makes PET exactly equivariant, with minimal changes to its accuracy.

Smooth, exact rotational symmetrization for deep learning on point clouds
Sergey N. Pozdnyakov and Michele Ceriotti
Laboratory of Computational Science and Modelling,
Institute of Materials, Ecole Polytechnique Fédérale de Lausanne, Lausanne 1015, Switzerland sergey. wowyyyoaooooeepp....... ch, michele.. wew.eleeeerittt@eelll..... ch
Abstract
Point clouds are versatile representations of 3D objects and have found widespread application in science and engineering. Many successful deep-learning models have been proposed that use them as input. The domain of chemical and materials modeling is especially challenging because exact compliance with physical con- straints is highly desirable for a model to be usable in practice. These constraints include smoothness and invariance with respect translations, rotations, and per- mutations identical atoms. these requirements are not rigorously fulfilled, atomistic simulations might lead absurd outcomes even if the model has excel- lent accuracy. Consequently, dedicated architectures, which achieve invariance by restricting their design space, have been developed. General-purposs point-cloud models are more varied but often disregard rotational symmetry. We propose a general symmetrization method that adds rotational equivariance any given model while preserving all the other requirements. Our approach simplifies the development better atomic-scale machine-leamae learning schemes by relaxing the con- straints on the design space and making it possible to incorporate ideas that proved effective in other domains. We demonstrate this idea by introducing the Point Edge Transformer (PET) architecture, which is not intrinsically equivariant but achieves atate-of-the-art performance on several benchmark datasets molecules and solids. A-posteriori application of our general protocol makes PET exactly equivariant, with minimal changes its accuracy.
1 Introduction
Contrary to 2D images that are well-described by regular and dense pixel grids, 3D objects usually have non-uniform resolution and are represented more naturally by an irregular arrangement of points in 3D. The resulting point clouds are widely used in many domains, including autonomous driving, augmented reality, and robotics, well in chemistry and materials modeling, and are the input of a variety of dedicated deep-learing techniques[[]. Whenever they are used for applications to the physical sciences, it is desirable to ensure that the model is consistent with fundamental physical constraints: invariance to translations, rotations, and permutation of identical particles, as well as smoothness with respect to geometric deformations ( In the application domain we focus on here, exact compliance with these requirements is highly sought after. Lack of smoothness or symmetry breaking are common problems of conventional atomistic modeling techniques | 6 These have been shown in the past to lead to artifacts ranging from numerical instabilities 7 (which are particularly critical in the context of tigh-tthooughut calculations g 9) to qualitatively incorrect and even absurd 10 simulation outcomes.
of atomistic modeling, the
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

These concerns have led to the development of dedicated models for atomistic simulations that rigorously incorporate all these constraints [5] by using only symmetry preserving operations. These restrictions limit the design space of these models, for example, leading to the lack of universal- approximation property in many popular methods, as we discuss in Section 2 Furthermore, symmetry requirements prevented the application of models developed in other domains to atomistic simulations. As an illustrative example, one can mention PointNet++[I11 an iconic architecture for generic point clouds. This model is not rotationally invariant. Therefore, even though it might have excellent accuracy, it has never been applied to atomistic simulations to the best of our knowledge. It is essential to distinguish between exact, rigorous equivariance and an approximate one, which any model can learn by rotational augmentations. Due to the subtle nature of possible artifacts, an approximate equivariance is considered insufficient.
Out of the mentioned symmetry constraints, the only challenging one rotational equivariance. As we discuss in Section 4 all the other requirements are either already fulfilled in most functional forms used by existing models for generie point clouds or can be enforced with trivial modifications.
In this paper, we introduce a general symmetrization protocol that enforces exact rotational invariance a posteriori, for an arbitrary backbone architecture, without affecting its behavior with respect to all the other requirements. It eliminates the wall between communities, making most of the developments for generic point clouds applicable to atomistic simulations. Furthermore, simplifies the development of new, more efficient architectures by removing the burden of incorporating the exact rotational invariance into their functional form. As an illustration, we design a model named Point Edge Transformer (PET) that achieves state-of-the-an performance on several benchmarks ranging from high-energy CH₃ configurations to a diverse collection of small organic molecules to the challenging case of periodic crystals with dozens of different atomic species. Being not intrinsically rotationally equivariant, PET has benefited from the enlarged design space. Our symmetrization method posteriori makes it rigorously equivariant and, thus, applicable to atomistic simulations.
2 Equivariant models and atomic-scale applications
Models used in chemical and materials modeling achieve rotational invariance by two main mech- anisms. The first involves using only invariant internal coordinates such as interatomic distances, angles, or even dihedral angles[[]]. The second involves using equivariant hidden representations that transform in predictable manner under symmetry operations. For example, some intermediate activations in a neural network can be expressed as vectors that rotate together with the input point cloud [8]. This approach restricts the design space to only such functional forms that preserves the equivariance.
Local decomposition. Models in atomistic machine learning often rely on the prediction of local properties associated with atom-centered environments A,, either because they are physical observables (e.g., NMR chemical shieldings [19 (20) or because they provide a decomposition of a global extensive property, such as the energy, into a sum of atomic contributions, y E, (A)) Here, ((A)) indicates learnable function of the environment of the i-th -h atom, defined as all the neighbors within the cutoff radius R., This local decomposition rooted in physical considerations 21 and is usually beneficial for the transferability of the model.
Local invariant descriptors. The classical approach to construct approximants for v(A,) is to first compute smooth, invariant descriptors for each environment 14, 22-27 and then feed them to conventional machine learning model. Such models can be 1) linear regression 27). 2) kernel regression 28. or 3) feed-forward neural network with a smooth activation function. 22
Distance-based messsae--pssingg Another popular method involves constructing molecular graph and feeding it to a Graph Neural Network (GNN). Early, and very popular, models rely on invariant two-body messages based only interatomic distances 30 31] In this case, molecular graph is constructed by representing all the atoms by nodes, drawing edges between all the atoms within a certain cutoff radius, and decorating edges with the Euclidean distance between the corresponding atoms.
Recently, it was discovered that such models are not universal approximators. Specifically, they cannot distinguish between certain atomic configurations(32) This issue can be directly attributed to the restricted design space of rotationally equivariant models. In an enlarged design space, one
2

---

could decorate the edges of molecular graph with x, y, and components of the corresponding displacement vectors. This would immediately ensure the universal approximality of the model. However, the necessity to enforce rotational equívariance of the predictions dictates decorating the edges of a molecular graph only with rotationally invariant Euclidean distances. Such a restriction severely limits the expressiveness of the corresponding models.
Subsequently, more expressive models were developed. For example, some of them use angles 3x] or even dihedrals 17, in addition to distances between the atoms. Others, such as Tensor Field Network [34, Nequip 35 MACE SE(3) transformers
Higher-order gressage-passing
employ SO(3) algebra to maintain equivariance of hidden representations. These models solve the incompleteness issue of simple atom-centered geometric descriptors and distance-based GNNs. The quality of a model, however, doesn't 't reduce to the simple presence of universal approximation behavior. An extended design space can still be beneficial to obtain more accurate or efficient models.
3 Models for generic point clouds
Point clouds have found many applications beyond atomistic modeling. For example, detectors such LiDARs represent the scans of the surrounding world as collections of points. Multiple methods have been developed for such domains. Contrary to atomistic machine learning, there are no such strict symmetry requirements for practical applications. Consequently, most models do not exactly incorporate rotational equivariance and rely instead on rotational augmentations.
Many successful models based on 2D projections of the cloud have been proposed in computer vision 38 41.
Here, we focus on explictly 3D models, that are most relevant for chemical applications.
Voxel-based methods. The complete 3D geometric information for structure can be encoded in
cnnpuutaaioonnnvariinn manner by projecting the point cloud onto a regular 3D voxel grid, and further manipulated by applying three-dimensional convolutions. The computational cost of naive approach, however, scales cubically with the resolution, leading to the development of several schemes to exploit the sparsity of the voxel population 42 44), which is especially pronounced for high resolutions.
Point-based methods. By extending the convolution operator to an irregular grid, one can avoid the definition of a fixed grid of voxels. For instance 46), one can evaluate an expression such as
Nin
(F) >> 9mn (r, r)...
where, fú is the n-th feature of the point i, and is a collection of Nout learnable three- dimensional functions, where Ni. and Nout are the numbers of input and output features respectively, The output features (F * (r)) can be evaluated at an arbitrary point r, allowing for complete freedom in the construction of the grid. PointNet [7] is paradigmatic example of architecture that eliminates the distinction between point features and positions. The core idea of these approaches is that an expression such as f(11, F2.., In) (maxx ((h(rr)))) where " and h are learnable functions, can approximate to arbitrary precision any continuous function of the point set {rx)- PointNet+ and many other models 44 48 (4) apply similar functional form in hierarchical manner to extract features of local neighborhoods. Many of these methods have graph-convolution or message- paaiiggg-aannggggge forms, similar to those discussed in Section [ even though they do not enforce invariance or equivariance of the representation. Many transformer models for point clouds have also been proposed (30)53) that incorporate attention mechanisms at different points of their architecture.
4 Everything but rotational equivariance
The generic point-cloud models discussed in the previous section do not incorporate all of the requirements (permutation, translation and rotation symmetry as well as smoothness) that are required by applications to atomistic simulations. Most models, however, do incorporate everything but rotational invariance, or can be made to with relatively small modifications. For instance, translational invariance can be enforced by defining the grids or the position of the points in a coordinate system that is relative reference point that is rigidly attached to the object.
Another common problem of many of these methods is the lack of smoothness. However, most architectures can be made differentiable with relative ease, with the exception of very few operations
3

---

such as downsampling via farthest point sampling(16, For example, derivative discontinuities associated with non-smooth activation functions. such as ReLU[54), can be eliminated simply by using a smooth activation functions instead 55 58 A slightly less trivial problem arises for models using the convolutional operator of Eq. T that introduces discontinuities related to (disaappearance of new points at the cutoff sphere, even for smooth functions Imn- These discontinuities can be removed by modifying the convolutional operator as:
Nan
(F * g)))
Σ>
Imn (r+ ffflff((l -r1] Rr,, AR.)
where the cutoff function fa. illustrated in Fig- 5t. can be chosen as an infinitely differentiable switching function that smoothly zeroes out contributions from the points approaching the cutoff' sphere, and the parameter AR. controls how fast converges to for r < Rg. A similar strategy can be applied to PointNet-like methods. 3D CNNs on regular grids that use smooth activation and pooling layers such as sum or average operations are also smooth functions of the voxel features. Including also smooth projection of the point cloud on the voxel grid suffices to make the overall methods smooth. We further discuss in the Appendix G examples of such smooth projections, along with other modifications that can be applied to make the general point-cloud architectures discussed in Sec. [differrntiabl.. Finally, most but not all [59] - the generic models for point clouds are invariant to permutations.
In summary, there a wealth of point-cloud architectures that have proven to be very successful in geometric regression tasks, and are, or can be easily made, smooth, permutation and translation invariant. The main obstacle that hinders their application to tasks that require fully-symmetric behavior, such those that are common in atomistic machine learning, is invariance covariance under rigid rotations. The ECSE protocol presented in the next section eliminates this barrier.
5 Equivariant Coordinate System Ensemble
Our construction accepts a smooth, permutationally. translationally, but not necessarily rotationally invariant backbone architecture along with point that is rigidly attached to the point cloud, and produces an architecture satisfying all four requirements. While in principle, there are many choices of the reference point, we will focus for simplicity on the case in which is one of the nodes in the point cloud. More specifically, (possibly tensorial) target (AA = Σ, y(A)) where the reference point is given by the position of the central atom r, and the model estimates an atomic contribution yA,) given the local neighborhood. We name our approach Equivariant Coordinate System Ensemble (ECSE, pron. efe).
formulate our method in the context of local decomposition of the
Local coordinate systems. A simple idea to obtain a rotationally equivariant model would be to define coordinate system rigidly attached to an atomic enviroument. Next, one can express Cartesian coordinates of all the displacement vectors from the central atom to all the neighbors and feed them to any, possibly not rotationally invariant, backbone architecture. Since the reference axes rotate together with the atomic environment, the corresponding projections of the displacement vectors are invariant with respect to rotations, and so is the final prediction of the model.
This approach can be applied easily to rigid molecules [60 61). but in the general case it is very difficult to define coordinate system in a way that preserves smoothness to atomic deformations. For example, the earliest version of the DeepMD framework (62 used a coordinate system defined by the central atom and the two closest neighbors. Smooth distortions of the environment can change the selected neighbors, leading to a discontinuous change of the coordinate systems and therefore to discontinuities in the predictions of the model. For this reason, later versions of DeepMD[63) switched to descriptors that can be seen as close relative of guarantee smoothness and invariance.
symmetry functions [2 which
Ensemble of coordinate systems. The basic formulation of our symmetrization protocol, illustrated in Fig. 55, is simple: using all the possible coordinate systems defined by all pairs of neighbors instead of singling one out, and averaging the predictions of a non-equivariant model over this ensemble of reference frames:
ys(A,) = Σ
Why Ryg (vo( (A.DI/ Σ je'
Wij',
(3)

---

Ryy
R7f A.]
yo(Ar,31A.))
stomne cavironment 4.
Backbone architecture
Weighted average Vs(A.) with weights 10y!"
Backhone architertnure
d Ri. Fune(A))
Backhone archiitecturre
for
" - Selre)l R As. Ba)eetrr 1a x Pay Plw, A.)
sparity-wwile
Mir R. Ap.)
9dallw 3.)
OJ
R2'
R La, K Ti
Fi X F
Figure I: (a) Equivariant coordinate-system ensemble: Each ordered pair of neighbors defines a local coordinate system. Next, an atomic environment is projected on all of them (which is equivalent to rotation) and used input for backbone architecture. If outputs covariant, such vectors, they are rotated back to the initial coordinate system. Finally, predictions are averaged over. (b) Discontinuities related plain average. The weighted average with weights wzy' resolves these problems. (c) Cutoff functions fc and &. used to define weights Wyy (d) To reduce the computational cost, adaptive cutoff Rin is used, which adjusts to given geometry instead of being global user-specified constant.
where ys indicates the symmetrized model, yo the initial (non-equivariant) backbone architecture, and Rg indicates the rotation operator for the coordinate system defined by the neighbors j and j' within the i-centered environment A,. other words, Rn [A,] is an atomic environment expressed in the coordinate system defined by neighbors j and j' The summation is performed over all the ordered pairs of neighbors within a cutoff sphere with some cutoff radius R..
Given pair of neighbors with displacement vectors from central atom ri and rij' the corresponding coordinate system consists of the vectors fij. Fij X fay', and Pij x [Faj Fiyy]. If the model predicts a vectorial (or tensorial) output, it is rotated back into the original coordinate system by an outer application of operator Rg Thus, our symmetrization scheme allows for getting not only invariant predictions, but also covariant ones, such vectorial dipole moments. It is worth to note that the idea of using all the possible coordinate systems has also been used to define a aooomillly-comptaall invariant metric measure the similarity between crystal structres (64 and to construct provably complete invariant density-correlatte descriptors [67].
Fig. 5p depicts the necessity to use the weighted average with weights wyy instead of plain one. The use of plain average in Eq, 3 would lead to the lack smoothness. Indeed, a new atom enters the cutoff sphere, it immediately yields new terms into the summation in eq. [] which would lead to a discontinuous gap in predictions. Furthermore, if two neighbors and a central atom appear to be collinear, the associated coordinate system is ill-defined.
Both issues can be solved by introducing a mechanism in which each coordinate system assigned different weight, depending on the positions (rij, (iy)) of the neighbors that define the reference frame:
Way w/ (rij. riy:) fe(riy) Re, A (n) Re, An. (nFr, x Ptye Aw),
where the cutoff functons fa and ga are illustrated Fig. Sp. Thanks to the presence Sc, the terms in Eq. (3) that are associated with atoms entering or exiting the cutoff sphere have zero weights. Similarly, 4: ensures that pairs of neighbors that are nearly collinear do not contribute to the symmetrized prediction. One last potential issue is the behavior when all pairs of neighbors
5

---

@ ambiguity. We discuss in the (nearly) collinear, which would make Eq. (3) ill ill-defined, falling to o Appendix two possible solutions for this corner case.
Adaptive cutoff. To make the ECSE protocol practically feasible, it is necessary to limit the number of evaluations of the non-equivariant model, given that the naive number of evaluations grows quadratically with the number of neighbors. An obvious consideration is that there is no reason why the cutoff radius used by ECSE should be the same as the one used by the backbone architecture. Thus, one can achieve significant computational savings by simply defining a smaller cutoff' for symmetrization. However, particularly for inhomogeneous point cloud distributions, a large cutoff may still be needed to ensure that all environments have at least one well-defined coordinate system. For this reason, we introduce an adaptive inner cutoff Rm, which is determined separately for each atomic environment A, instead of being a global, user-specifin eser-specified constant. Eq. (3) requires that at least one pair of neighbors inside the cutoff sphere. Simultaneously, it is desirable to make it as small as possible for computational efficiency. Thus, it makes sense to define Rm,A;) along the lines of the distance from the central atom to the second closest neighbor. should be larger, but not much larger, than the rerondd-earest-neighbbe distance.
Our definition of Rin (A:), given in the Appendix F.4 is inspired by this simple consideration, and has the following properties: (1) Contrary to the naive second-neighbor distance, Ra, (A,) depends smoothly on the positions of all the atoms. (2) It encompasses at least one non-collinear pair of neighbors, which yields at least one well-defined coordinate system: the functional form is chosen so that if & nearest neighbors are collinear, Ru, (A,) expanded to enclose the k + 1-th for any k.
With such an adaptive cutoff, only a few pairs of neighbors are used to construct a coordinate system. Thus, this construction greatly reduce the cost of evaluating the ECSB-symmetrized model.
Training and symmetrization. The most straightforward way to train a model using the ECSE protocol is to apply it from the very beginning and train model that is overall rotationally equivariant. This approach, however, increases the computational cost of training, since applying ECSE entails multiple evaluations of the backbone architecture. An alternative approach, which we follow in this work, is to train the backbone architecture with rotational augmentations, and apply ECSE only for inference. Given that ECSE evaluates the weighted average of the predictions of the backbone architecture, this increases the cost of inference, but may also increase the accuracy of the backhone architecture, playing the role of test augmentation.
Message-passing, gressage-passingg schemes can be regarded as local models with cutoff radius given by the receptive field of the GNN, and therefore ECSE can be applied transparently to them. In a naive implementation, however, the same message would have to be computed for the coordinate systems of all atoms within the receptive field. This implies very large overhead, even though the asymptotic cost remains linear with system size. An alternative approach involves symmetrizing the outgoing messages with the ECSE protocol. This requires deeper integration within the model, that also changes the nature of the message-passe passing step, given that one has to specify a transformation rule to be applied to messages computed in different coordinate systems. We discuss in the Appendix F.10 the implications for the training strategy. and some possible workarounds.
Related work. There are several approaches, such as the earliest version of the DeepMD framework(62) discussed above, that use local coordinate systems to enforce rotational equiv- ariance for any backbone architecture. The frame averaging (FA) framework (68 69 proposes to use the eigenvectors of the centered covariance matrix to define local coordinate systems, which leads to discontinuities every time the eigenvalues coincide with each other. To the best of our knowledge, ECSE is the first general symmetrization method allowing to enforce rotational equivariance while preserving smoothness, which is a highly desirable property for atomistic simulations.
Finally, completely different approach is introduced by the vector neurons framework 70 that con- verts arbitrary backbone architecture into an equivariant one by enforcing all hidden representations in the model to transform as vectors with respect to rotations.
6 Point Edge Transformer
An obvious application of the ECSE protocol would be to enforce equivariance on some of the point- cloud architectures discussed in Section 3 making them directly-applicalle to atomistic modeling and to any other domain with strict symmetry requirements. Here we want instead to focus less-

---

a
point cloud decorated with cheenical
Sugiet
displncenennt vonter
claanical spocae of the neighbor
((h,
Fy
MLP/Lines
MLP/ Lineat 3 + dppr
Embeeding
dosr
MP blodl
MP block
MP block
MP blonk
Concatetate
NCNN
3d
3dy dygr
b
atomic environment
predictions
Josr
Input messaage
Message- Passing (MP) block
Thtpuut massinge
Ne neighbors
PET
Eacoder
dest
Head,
predictions
Transformer layers
72
Impat musssage
ontpout messingp dper
Fim
permut covariant
Tina
dPsr
Encader
depp
Ieadd
messages
oubput
imput
Embeddiag
drer
dyser
Head,
messages
Figure 2: Architecture the Doint-Edge Transformer (PET). White and colored boxes represent layers; gray boxes and lines represent data. (a) PET is message- passiggge-aingggggo architecture. At each of the VIGNN pessage-passing (MP) interactions, messages are communicated between all the pairs of atoms closer than certain cutoff distance R. At each stage, the corresponding MP block computes output messages and predictions the target property given the input messages and geometry of define atom-centered environment A, collection of all the neighbors within the cutoff distance Re The MP block is applied to cach such atomic environment. Given the geometry of the atomic environment, 2) the chemical species of the atoms, and 3) input messages from all the neighbors to the central atom it produces output messages from the central atom to all the neighbors and contribution to the prediction of the target property. The first step is to encode all the information associated with each neighbor to an abstract token of dimensionality dpEr- Next, the collection of such tokens (with the one associated with the central atom) is fed into the transformer with R'H self-attention layers. The transformer does permutationally covariant transformation. Thus, the association between the tokens and neighbors preserved. Therefore, we can simply treat output tokens output messages to the corresponding neighbors. (c) The Encoder layer first maps all the sources of information into dimensionality dpET- Next, all 3 tokens are concatenated and compressed into single one of the desired size.
the point cloud. (b) For each atom in the system,
obvious, but perhaps more significant, application: demonstrating that lifting the design constraint of rotational equivariance makes it possible to construct better models. To this end, we introduce new architecture that we name Point Edge Transformer (PET), which built around a transformer that processes edge features and achieves state-of-the-art accuracy on several datasets that cover multiple subdomains of atomistic ML. Due to the focus on edge features, PET shares some superficial similarities with the Edge Transformer 71 developed for natural language processing, Allegro (a strictly local invariant interatomic potential (22) that loosely resembles a single message- block of PET) well with early attempts by Behler et al. applying MLPs to invariant edge features (3)
The PET architecture illustrated in Fig. More details are given Appendix A component of each message- message-passing block is permutationally--eee -equivvaraaa transformer that takes a set of tokens of size dpET associated with the central atom and all its neighbors, and generates new tokens that are used both to make predictions and output messages. The transformer we use in PET straightforward implementation the classical one (4), with a modification to the attention mechanism that ensures smoothness with respect to (dis)appearance of the neighbors the
The core
7

---

cutoff radius. The attention coefficients azy determining the contribution from token j to token i are modified as aij - foy R2, AnJ, and then renormalized.
One of the key features of PET is that it operates with features associated with each edge, at variance with other deep learning architectures for point clouds that mostly operate with vertex features. Whereas the calculation of new vertex features by typical GNNs involves aggregation over the neighbors, the use of edge features allows for the construction of an aggregation-feee message-passing scheme, avoiding the risk of introducing an information bottleneck.
The transformer itself can be treated as a GNN[75], thus our architecture can be seen as performing localized message passing, increasing the complexity of local interactions it can describe, while avoiding an over-inerease of the receptive field. The latter is undesirable because it makes parallel computation less efficient, and thus hinders the application of such models to large-scale molecular dynamics simulations (7..
Since transformers are known to be universal approximators (76), so is each of the message-passing blocks used in PET.
Limitations. As discussed in Section 5 the most efficient application of the ECSE protocol on top of a GNN requires re-designing the message- paasgggg-agngggggg mechanism. For the current moment, our proof-of-principee
implementation follows a simpler approach (see more details in Appendix [F.1I) and favors case of implementation instead of computational efficiency (c.g., all weights and most of the intermediate values used in the ECSE scheme are stored in separate pero-dimensional Py PyTorch[77] tensors). This leads to significant (about 3 orders of magnitude) overhead over the inference of the backbone architecture. The base PET model, however, gains efficiency by operating directly on the Cartesian coordinates of neighbors. Even with this inefficient implementation of ECSE, exactly cquivariant PET inference is much cheaper than the reference first-principles calculations.
Benchmarks
We benchmark PET and the ECSE scheme over six different datasets, which have been previously used in the literature and which allow to showcase the performance of our framework, and the ease with which can be adapted to different use cases. The main results, are compared with the state of the art in Figure and
and Table while in-depth analyses can be found in the Appendix C
As a first benchmark, we conduct several experiments with the liquid-wate water configurations from Ref. [85] This dataset is representative of those used in the construction of interatomic potentials, and presents interesting challenges in that it contains distorted structures from path integral molccular dynamics and involves long-range contributions from dipolar electrostatic interactions. One of the key features of PET is the possibility to increase the expressive power by either adding MP blocks or by making transformers in each block deeper. Panel (a) in Fig. 3 shows that increasing the number of GNN blocks improves the accuracy of PET, and that for a given number of blocks, a shallow transformer with single layer performs considerably worse than one with two more layers. Given that stacking multiple GNN blocks increases both the receptive field and the flexibility in describing local interactions, it is interesting (Fig. Bb), that shows that a 6 x 2 model outperforms a large 12 x stack of shallow transformers, even though the latter has additional flexibility because of the larger number of heads and pre-processing units. With the exception of the shallower models, PET reduces the сттоr over the state-of of titte-of-theaate t-the-rrt equivariant model NEQUIP 335 by 30%. In problems that are less dependent on long-range physics, it may be beneficial to increase the depth of transformers and reduce the number of MP blocks. The accuracy of the model can also be further improved by extending Re beyond 3.7 À (a value we chose to ensure approximately 20 neighbors on average), reaching a force MAF of 14.4 me V/Á when using a cutoff of 4.25A.
look the trend for a fixed total number of transformer layers
We then move to the realm of small molecules with the COLL 86] that contains distorted configurations of molecules undergoing a collision. COLL has been used extensively as a benchmark for chemical ML models, in particular for GNNs that employ information on angles and dihedrals to concile rotaional invariance and universal approximation [17 87). PET improves by 13% the emror on forces (the main optimization target in previous studies) while reducing by factor of 4 the error on atomization energies relative to the respective state of the art. In order to assess the accuracy of PET for extreme distortions, and very wide energy range, consider the database of CH2
8

---

water
water
water
Rrae
MAE meVi
10
fru
NEQUIP
NEQUIP
xx
NEQUI
2x10
12 3.0
3.5
4,0
4,5
CH4
f
QM9 dipales
REANMB4]
10
10a
de. 256
MAE (M /
10a
WXE
ROMSE
10a
S3AP
RECE/75)
[CHL((83]
GPRI82]
LE ACE/79)
WK2811
PET 120
10 10a
10a
Figure 3: (a-c) Accuracy of PET potentials (16) of liquid water, compared with NEQUIP[3). (a) Accuracy for different numbers pessage-passingg blocks AGNN and transformer layers nm. (b) Accuracy function of VIGNN, for constant AINNN Y'. 12.; (c) Accuracy as function of cutoff. (d-f) Learning curves for different molecular data sets, comparing symmetrized PET models (ys) with several previous works 23 (8). including the current state of the art. (d) Random CH₃ dataset, training only on energies: (c) Random CH₃ dataset, training on energies and forces; (f) Vectorial dipole moments in the QM9 [82]
configurations first introduced in Ref. 80 This dataset contain random arrangements of one C and H atoms, that are only filtered to eliminate close contacts. The presence of structures close to those defying C-centered angle-distances descriptors adds to the challenge of this dataset. The learning curves for PET Figure 3e and demonstrate the steep, monotonic improvement of performance for increasing train set size, surpassing after few thousand training points the best existing models (a linear many-body potential for the energy-only training 81] and REANN, an angle-based based GNN 84 for combined energy/gradient training).
Table 1: Comparison of the accuracy of PET and current state-of-the-ee models for the COLL, MnO, HM21 and HEA data sets. A more comprehensive comparison with leading models provided in the Appendix Energy nearly equivariant, and the difference between the accuracy of yo and ys is minuscule. For HME21, report error bars 5 random seeds, and results for ensemble of symmetrized models.
are given in me V/atom, force errors in me V/À. In all cases the PET model
dataset metric
COLL
MnO
HME21
HEA
MAE f
MAE E
RMSE f RMSE E/at.
MAE [fl MAE E/at. MAE f MAE E/at
SOTA model
16
47
125 88
1.11
88
89
15.7
190
10 90
GemNet
DimeNet++
mHDNNP
MACE
HEA25-4-NN
PET (yo)
23.1
12.0
22.7
0.312
140.5 = 2.0 17.8
+ 0.1
60.2
1.87
PET (ys)
23.1
11.9
22.7
0.304
141.6 = 1.9 17.8 - 0.1
60.1
1.87
PET (ens.)
128.5 16.8
The PET performs well even for condensed-phase datasets that contain dozens of different atomic types. The HEA dataset contains distorted crystalline structures with up to 25 transition metals simultaneously [00] Ref. 90 performs an explicit compression of chemical space, leading to a model that is both interpretable and very stable. The tokenization of the atomic species in the encoder layers of PET can also be seen as a form compression. However, compared to HEA25-4-NN allows for a 3(5) reduction of force(energy) hold-out errors. The model, however, loses somewhat in transferability: in the extrapolative test on high-temperature MD trajectories suggested in Ref. 90) PET performs less well than HEA25-4-NN (152 vs 48 me V/at.
more flexible form of the model

---

MAE at 5000 K) even though room room-tempentte temperature MD trajectories are better described by PET. The case of the HME21 dataset, which contains high-temperature molecular dynamiess configurations for structures with a diverse composition, is also very interesting. PET outperforms most of the existing equivariant models, except for MACE(89), which incorporates a carefully designed set of physical priors inheriting much of the robustness of shallow models. PET, however, comes very close: the simple regularizing effect of 5-models PET ensemble is sufficient to tip the balance, bringing the force error to 128.5 V/at Another case in which PET performs well, but not as well as the state of the art, is for the prediction of the atomization energy of molecules in the QM9 dataset PET achieves MAE of 6.7 me V/molecule, in line with the aceuracy of DimeNet++ [8] but not as good as the 4.3 me Vmolecule MAE achieved by Wigner kemels(81]].
Finally, we consider two cases that allow us to showcase the extension of PET beyond the prediction of the cohesive energy of molecules and solids. The МnO dataset of Eckhoff and Behler[88) includes information on the colinear spin of individual atomic sites, and demonstrates the inclusion of information beyond the chemical nature of the elements the description of the atomic environments. The improvement with respect to the original model is quite dramatic (Tab. [. Finally, the QM9 dipole dataset [82] allows us to demonstrate how to extend the ECSE to targets that are covariant, rather than invariant. This is as simple as predicting the Cartesian components of the dipole in each local coordinate system, and then applying to the prediction the inverse of the transformation that is applied to align the local coordinate system (cf. Eq. (3) The accuracy of the model matches that of recently-des deeeely-dee
bady kernel regression scheme [8]] for small dataset size, and outperforms it by up to 30% at the largest train set size.
8 Discussion
In this work, we introduce ECSE, a general method that enforces rotational equivariance for any backbone architecture while preserving smoothness and invariance with respect to translations and permutations. To demonstrate its usage, also develop the PET model, deep-Icarning architecture that is not intrinsically rotationally invariant but achieves state-of-the-ee results on several datasets across multiple domains of atomistic machine learning. The application of ECSE makes the model comply with all the physical constraints required for atomistic modeling.
We believe our findings to be important for several reasons. On the one hand, they facilitate the application of existing models from geometric deep learning to domains where exact equivariance is necessary condition for practical applications. On the other, they challenge the notion that equivariance is a necessary ingredient for an effective ML model of atomistic properties. The state-of- the-art performance of PET on several benchmarks reinforces early indications that non-equivariant models trained with rotational augmentation can outperform rigorously equivariant ones. Similar observations were also made independently by other groups. Spherical Channel Network [92 and ForceNet [93] achieve excellent performance on the Open Catalyst dataset 94] exact equivariance of the model.
while relaxing the
is worthwhile to mention that in the other domains involving point clouds, the application of not invariant models fitted with rotational augmentations is a predominant approach. As example, one can examine the models proposed for, e.g., the ModelNet40 dataset a popular benchmark for point cloud classification. Even though the target is rotationally invariant, most of the developed architectures are not rigorously equivariant, in the exact sense, as discussed in Section. 2 This provides considerable empirical evidence that the use of not rigorously equivariant architectures with rotational augmentations might be more efficient.
The ECSE method we propose in our work allows to symmetrize exactly, and a-posteriori, any point-cloud model, making them suitable for all applications, such as atomistic modeling, which have traditionally relied on symmetry to avoid qualitative artifacts in simulations. This shall facilitate the translation of methodological advances from other domains to atomistic modeling and the implementation into generic point-cloud models of physically-inspires inductive biases (such as smoothness, range, and body order of interactions) that have this far been conflated only with specific class of equivariant architectures.
10

---

Acknowledgements
We thank Marco Eckhoff and Jörg Behler for sharing the MnO dataset with the collinear spins and Artem Malakhov for useful discussions.
This work was supported by the Swiss Platform for Advanced Scientific Computing (PASC) and the NCCR MARVEL, funded by the Swiss National Science Foundation (SNSF, grant number 182892).
References
[1] Yulan Guo, Hanyun Wang. Qingyong Hu, Hao Liu, Liu, and Mohammed Bennamoun. Deep learning for 3d point clouds: A survey. IEEE transactions pattern analysis and machine intelligence, 43(12):4338-4333
2020.
[2] Michael M. Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geometric Deep Learning: Going beyond Euclidean data. IEEE Signal Process. Mag., 34(4):18-42, July 2017.
[3] Giuseppe Carleo, Ignacio Cirac, Kyle Cranmer, Laurent Daudet, Maria Schuld, Naftali Tishby, Leslie Vogt-Maranto, and Lenka Zdeborová. Machine learning and the physical sciences. Rev. Mod. Phys., 91(4):045002, December 2019.
[4] George Em Kariiadakis, Ioannis G. Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning, Nat Rev Phys, 3()):422-440, May 2021.
[5] Jirasak Wong-ekkabut and Mikko Karttunen. The good, the bad and the user in soft matter simulations. Biochimica Biophysica Acta (BBA) Biomembranes, (558(10):559-2558) October 2016.
[6] Danny E. P. Vanpoucke, Kurt Lejaeghere, Veronique Van Speybroeck, Michel Waroquier, and An Ghysels. Mechanical Properties from Periodic Plane Wave Quantum Mechanical Codes: The Challenge of the Flexible Nanoporous MIL-47(V) Framework. J. Phys. Chem. C, 1 119(41):23752-2333 237666,
October 2015.
[7] Michael F Herbst and Antoine Levitt. Black-box inhomogeneous preconditioning for self- consistent field iterations in density functional theory. Phys.: Condens. Matter, 33(8)0085503, February 202..
[8] Wahyu Setyawan and Stefano Curtarolo. High-throughput electronic band structure calcu- lations: Challenges and tools. Computational Materials Science, 49(2):299 312, August 2010.
[9] Valerio Vitale, Giovanni Pizzi, Antimo Marrazzo, Jonathan R. Yates, Nicola Marzari, and Arash A. Mostofi. Automated high-throughhur Wannierisation. npj Comput Mater, 6(1):66, December 2020.
[10] Xiaojing Gong, Jingyuan Li, Hangjun Lu, Rongzheng Wan, Jichen Li, Jun Hu, and Haiping Fang. charge-driven molecular water pump. Nature Nanotech, 2(11):709-712, November 2007.
[11] Jirasak Wong-ekkabut, Markus S. Miettinen, Cristiano Dias, and Mikko Karttunen. Static charges cannot drive continuous flow of water molecules through carbon nanotube. Nature Nanotech, 5(8):555-557, August 2010.
[12] Jörg Behler and Michele Parrinello. Generalized Neural-Network Representation of High- Dimensional Potential-Enery Energy Surfaces. Phys. Rev. Lett., 98(14): April 2007.
[13] Albert P. Bartók, Mike C. Payne, Risi Kondor, and Gábor Csányi. Gaussian Approximation Potentials: The Accuracy of Quantum Mechanics, without the Flectrons. Phys. Rev. Lett,, 104(13):136403 April 2010.

---

[14] O. Anatole von Lilienfeld, Raghunathan Ramakrishnan, Matthias Rupp, and Aaron Knoll. Fourier series of atomic radial distribution functions: molecular fingerprint for machine learning models of quantum chemical properties. Int. J. Quantum Chem., 1 11(16)10044-109, August 2015.
[15] Stefan Klus, Patrick GelB, Feliks Nüske, and Frank Noé. Symmetric and antisymmetrric kernels for machine learning problems in quantum physics and chemistry. Mach. Learn.: Sci. Technol., 2(4):045016, December 2021.
[16] Charles Ruizhongtai Qi, Li Yi, Hao Su, and Leonidas J Guibas. Pointnet++: Deep hierarchical feature learning on point sets in a metric space. Advances in neural information processing systems, 30, 2017.
[17] Johannes Gasteiger, Florian Becker, and Stephan Günnemann, Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790-6802, 2021.
[18] Kristof Schüitt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In Int. Conf. Mach. Learn., pages 9377-9388. PMLR, 2021.
[19] Jérôme Cuny, Yu Xie, Chris J. Pickard, and Ali A. Hassanali. Ab Initio Quality NMR Parame- ters in Stlid-State Materials Using a High-Dimensional Neural-Network Representation. J. Chem. Theory Comput., 12(2):765-773, February 2016.
[20] Federico M. Paruzzo, Albert Hofstetter, Félix Musil, Sandip De, Michele Ceriotti, and Lyndon Emsley. Chemical shifts in molecular solids by machine learning. Nat. Commun., 9(1)4500, December 2018.
[21] E Prodan and W Kohn. Nearsightedness of electronic matter. Proc. Natl. Acad. Sci., (62533): August 2005.
[22] Jörg Behler. Atom-centered symmetry functions for constructing high-dimensional neural network potentials. The Journal of Chemical Physics, 134(7):074106, February 2011.
[23] Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Phys. Rev. B, 87(18):184115, May 2013.
[24] Mitchell A. Wood and Aidan P. Thompson. Extending the accuracy of the SNAP interatomic potential form. The Journal of Chemical Physics, 148(24) June 2018.
[25] Haoyan Huo and Matthias Rupp. Unified representation for machine learning of molecules and crystals. ArXiv Prepr. ArX1v170006439, 13754, 2017.
[26] Alexander V. Shapeev. Moment Tensor Potentials: A Class of Systematically Improvable Interatomic Potentials. Multiscale Model. Simul., 14(3):1153-1173 January 2016.
[27] Ralf Drautz. Atomic cluster expansion for accurate and transferable interatomic potentials. Phys. Rev. B, 99(1):014104, January 2019.
[28] Volker L. Deringer, Albert P. Bartók, Noam Bernstein, David M. Wilkins, Michele Ceriotti and Gābor Csányi. Gaussian Process Regression for Materials and Molecules. Chem. Rev., 121(16):10033-0144, August 2021.
[29] Jörg Behler. Four Generations of figh-Dimensional Neural Network Potentials. Chem. Rev., 1000770000037--007777 August 2021.
[30] Justin Gilmer. Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In Int. Conf. Mach. Learn., pages 1263-1272 2017.
[31] Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoe Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Mlaus-Robers Müller. SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. In NIPS, 2017.
12

---

[32] Sergey N Pozdnyakov and Michele Ceriotti. Incompleteness of graph neural networks for points clouds in three dimensions. Mach. Learn.: Sci. Technol. 344::0450 December 2022.
[33] Yaolong Zhang, Junfan Xia, and Bin Jiang. Reann: A pytorch-based based end-to-end multi- functional deep neural network package for molecular, reactive, and periodic systems. The Journal of Chemical Physics, 156(11). 1), 2022.
[34] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-oquivariann neural networks for 3D point clouds.
2018.
[35] Simon Batzner, Albert Musaclian, Lixin Sun, Mario Geiger, Jonathan Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nat Commun, 13(1):2453, May 2022.
[36] llyes Batatia, Simon Batzner, Dávid Péter Kovács, Albert Musaelian, Gregor N. C. Simm, Ralf Drautz, Christoph Ortner, Boris Kozinsky, and Gábor Csányi. The design space of E(3)-equivarian) atom-centered interatomic potentials. arxiv:2205.0664/ 2022.
[37] Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto-translation equivariant attention networks. Advances in Neural Information Processing Systems, 33:1970-1981, 2020.
[38] Hang Su, Subhransu Maji, Evangelos Kalogerakis, and Erik Learned-Miller. Multi-view con- volutional neural networks for 3d shape recognition. In Proceedings of the IEEE international conference computer vision, pages 945 943-953, 2015.
[39] Seyed Saber Mohammadi, Yiming Wang. and Alessio Del Bue. Pointview- gen:: 3d shape classification with multi-view point clouds. 202/ IEEE International Conference on Image Processing (ICIP), pages 3103-3107 IEEE, 2021.
[40] Xiaozhi Chen, Huimin Ma, Wan, Bo Li, and Tian Xia. Multi-view 3d object detection network for autonomous driving. In Proceedings of the IEEE conference Computer Vision and Pattern Recognition, pages 1907-1915, 2017.
[41] Asako Kanezaki, Yasuyuki Matsushita, and Yoshifumi Nishida. Rotationnet: Joint object categorization and pose estimation using multiviews from unsupervised viewpoints. In Pro- ceedings of the IEEE conference on computer vision and pattern recognition, pages 5010-5019, 2018.
[42] Peng-Shual Wang, Yang Liu, Yu-Xiao Guo, Chun Yu Sun, and Xin Tong. O-enn: Octree-based convolutional neural networks for 3d shape analysis. ACM Transactions On Graphics (TOG), 36(4):1-11, 2017.
[43] Gernot Riegler, Ali Osman Ulusoy, and Andreas Geiger. Octnet: Learning deep represen- tations at high resolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3577-3586. 2017.
[44] Roman Klokov and Victor Lempitsky, Escape from cells: Deep kd-networks for the recognition 3d point cloud models. Proceedings of the IEEE international conference on computer vision, pages 863-872, 2017.
[45] Shenlong Wang, Simon Suo, Wei-Chiu Ma, Andrei Pokrovsky, and Raquel Urtasun. Deep parametric continuous convolutional neural networks. In Proceedings of the IEEE conference on compuler vision and pattern recognition, pages 2589-2597, 2018
[46] Hugues Thomas, Charles R Qi, Jean-Emmanuel Deschaud, Beatriz Marcotegui, François Goulette, and Leonidas J Guibas. Kpconv: Flexible and deformable convolution for point clouds. In Proceedings of the IEFF/CVF international conference on computer vision, pages 6411-6420, 2019.
13

---

[47] Charles R Qi, Hao Su, Kaichun Mo, and Leonidas Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 652-660, 2017.
[48] Xu Ma, Can Qin, Haoxuan You, Haoxi Ran, and Yun Fu. Rethinking network design and local geometry in point cloud: simple residual mlp framework. arXiv preprint (rXiv:2202.07123, 2022.
[49] Mutian Xu, Runyu Ding, Hengshuang Zhao, and Xiaojuan Qi. Paconv: Position adaptive convolution with dynamic kernel assembling on point clouds. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3173-3182, 2021.
[50] Hengshuang Zhao, Li Jiang. Jiaya Jia, Philip HS Torr, and Vladlen Koltun. Point transformer. In Proceedings of the IEEE/CVF international conference on computer vision, pages 16259-
16268, 2021.
[51] Cheng Zhang, Haocheng Wan, Xinyi Shen, and Zizhao Wu. Pvt: Point-voxel transformer for point cloud learning. International Journal of Intelligent Systems, 37(12):11855-1200 2022.
[52] Jiageng Mao, Yujing Xue, Minzhe Niu, Haoyue Bai, Jiashi Feng, Xiaodan Liang, Hang Xu, and Chunjing Xu. Voxel transformer for 3d object detection. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 3164-3173, 2021.
[53] Meng-Han Hao Guo, Jun-Xiong Cai, Zheng-Ning Liu, Tai-Jiang Mu, Ralph Martin, and Shi-Min Hu. Pct: Point cloud transformer. Computational Visual Media, 7:187-199, 2021.
[54] Abien Fred Agarap. Deep learning using rectified linear units (relu). arXiv preprint arXiv:/803.0883 08375, 2018.
[55] Arnrk-Aréé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint 1oXix:/5/1.07899 2015.
[56] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.00 06441 2016.
[57] Diganta Misra. Mish: self regularized non-monotonic activation function. arXiv preprint arXiv:/908.0860 2019.
[58] Jonathan T Barron. Continuously differentiable exponential linear units. arXiv preprint urXiv:/704.074450 2017.
[59] Yangyan Li, Rui Bu, Mingchao Sun, Wei Wu, Xinhan Di, and Baoquan Chen. Pointenn: Convolution x-transformed points. arXiv preprint arXiv:1801.0779: 2018.
[60] Tristan Bereau, Denis Andrienko, and O. Anatole Von Lilienfeld. Transferable Atomic Multipole Machine Learning Models for Small Organic Molecules. Chem. Theory Comput., 11(7):225-32333 July 2015.
[61] Chungwen I. Liang, Gabriele Tocci, David M. Wilkins, Andrea Grisafi, Sylvie Roke, and Michele Ceriotti. Solvent fluctuations and nuclear quantum effects modulate the molecular hyperpolarizability of water. Phys. Rev. B, 964))044407 July 2017.
[62] Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and Weinan E. Deep Potential Molecular Dynamics: A Scalable Model with the Accuracy of Quantum Mechanics. Phys. Rev. Lett., 120(14):14300) April 2018.
[63] Han Wang. Linfeng Zhang, Jiequn Han, and Weinan E. DeePMD-kit: A deep learning package for body potential energy representation and molecular dynamics. Computer Physics Communications, 228: 18--184, July 2018
[64] Vitaliy Kurlin. Exactly computable and continuous metrics on isometry classes of finite and 1-periodic sequences. arXiv preprint Stti:2205500388. 2022.
[65] Vitaliy Kurlin. Computable complete invariants for finite clouds of unlabeled points under euclidean isometry. arXiv preprint 2rXiv:2277.08502 2022.
14

---

[66] Daniel Widdowson and Vitaliy Kurlin. Resolving the data ambiguity for periodic crystals. Advances in Neural Information Processing Systems (Proceedings of NeurIPS 2022). 35, 2022.
[67] Jigyasa Nigam, Sergey N. Pozdnyakov, Kevin K. Huguenin-Dumittann and Michele Ceriotti. Completeness of atomic structure representations. arxiv:2302. 2oxxi:2302.147770, 2023.
[68] Omri Puny, Matan Atzmon, Heli Ben- Ben-Hamu, Edward J Smith, Ishan Misra, Aditya Grover, and Yaron Lipman. Frame averaging for invariant and equivariant network design. arXiv preprint arXiv:21 2021.
[69] Alexandre Agm Duval, Victor Schmidt, Alex Hernández-Garent Santiago Miret, Fragkiskos Malliaros, Yoshua Bengio, and David Rolnick. Faenet: Frame averaging equivariant gnn for materials modeling. In International Conference on Machine Learning, pages 9013-9033. PMLR, 2023.
[70] Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas Guibas. Vector neurons: A general framework for so (3)-cqivvariamt networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 12200-12209, 2021.
[71] Leon Bergen, Timothy O' O'Donnell, and Dzmitry Bahdanau. Systematic generalization with edge transformers. Advances in Neural Information Processing Systems, 34:1390-1402, 2021.
[72] Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nat Commun, 14(1):579, February 2023.
[73] K. V.Jovnn Jose, Nongnuch Artrith, and Jörg Behler. Construction of high-dimensional neural network potentials using enviromment-dependenn atom pairs. J. Chem. Phys., 136(19):194111 2012.
[74] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017
[75] Chaitanya Joshi. Transformers are graph neural networks. The Gradient, 12, 2020.
[76] Chulhee Yun, Srinadh Bhojanapalli,, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077 2019.
[77] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vetooo, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
[78] Jigyasa Nigam, Sergey Pozdnyakov, and Michele Ceriotti. Recursive evaluation and iterative contraction of n-body equivariant features. The Journal of Chemical Physics, 153(12):121101, 2020.
[79] Filippo Bigi, Kevin K Srgguenin-uumitaan Michele Ceriotti, and David E Manolopou- los. A smooth basis for atomistic machine learning. The Journal of Chemical Physics, 157(23):234101, 2022.
[80] Sergey N Pozdnyakov, Michael Willatt, Albert P Baróók, Christoph Ortner, Gábor Csányi, and Michele Ceriotti. Incompleteness of Atomic Structure Representations. Phys. Rev. Lett., 125: 166006 2020.
[81] Filippo Bigi, Sergey N. Pozdnyakov, and Michele Ceriotti. Wigner kernels: body-ordered equivariant machine learning without basis. arxiv:2303.0m 04144, 2023.
[82] Max Veit, David M. Wilkins, Yang Yang, Robert A. DiStasio, and Michele Ceriotti. Predicting molecular dipole moments by combining atomie partial charges and atomic dipoles. Chem. Phys., 153(2):024113, July 2020.
15

---

[83] Felix A. Faber, Anders S. Christensen, Bing Huang, and O. Anatole Von Lilienfeld. Alchemical and structural distribution based representation for universal quantum machine learning. Chem. Phys., 148(24):241717, June 2018.
[84] Yaolong Zhang, Junfan Xia, and Bin Jiang. REANN: A PyTorch-based end-to-end multi- functional deep neural network package for molecular, reactive, and periodic systems. J. Chem. Phys., 156(1 1): 14801, 2022.
[85] Bingging Cheng, Edgar A. Engel, Jörg Behler, Christoph Dellago, and Michele Ceriotti. Ab initio thermodynamics of liquid and solid water. Proc. Natl. Acad. Sei. U. S. A., 116(4):1110- 11 1115, January 2019.
[86] Johannes Gasteiger, Shankari Giri, Johannes T Margraf, and Stephan Günnnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14115, 2020.
[87] Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional Message Passing for Molecular Graphs. 2020.
[88] Marco Eckhoff and Jörg Behler. High- dighhdimmnsional neural network potentials for magnetic sys- tems using spin-dependent atom-centered symmetry functions. npj Computational Materials, 7(1):770, 2021.
[89] James P Darby, Dávid P Kováes, Ilyes Batatia, Miguel A Caro, Gus LW Hart, Christoph Ortner, and Gábor Csányi. Tensor-reduced atomic density representations. arXiv preprint axXiv:2210.01705, 2022.
[90] Nataliya Lopanitsyna, Guillaume Fraux, Maximilian A. Springer, Sandip De, and Michele Ceriotti. Modeling high-entropy transition metal alloys with alchemical compression. Phys. Rev. Materials, 7(4)0045802, April 2023.
[91] Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Sei. Data, 1:1-7, August 2014.
[92] Larry Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sri- ram, Zachary Ulissi, and Brandon Wood. Spherical channels for modeling atomic interactions. Advances in Neural Information Processing Systems, 35:8054-8067. 2022.
[93] Weihua Hu, Muhammed Shuaibi, Abhishek Das, Siddharth Goyal, Anuroop Sriram, Jure Leskovec, Devi Parikh, and C Lawrence Zitnick. Forcenet: A graph neural network for large-scale quantum calculations. arXiv preprint arXiv:2103.01436 2021.
[94] Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo Caleb Ho, Weihua Hu, et al. Open catalyst 2020 (oc20) datasct and community challenges. Acs Catalysis, 11(10):6559-6072 2021.
[95] Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao. 3d shapenets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1912-1920, 2015.
[96] Stefan Elfwing, Biji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
[97] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[98] Scipy generator of random rotations. httpp:://doot http:///s...oo.orrrr/////////eee generated/ sctpyy/cpppp..y. www.yss.paaaa....a..a........... transf orm. wwt.tttn.nnnn.....o.....ot....tt
www.assiioomttttioiooooommmmo wwwpaatiiinrrnnno....mmm html
[99] Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv
preprint arXiv:1412.6941 2014.
16

---

[100] So Takamoto, Chikashi Shinagawa, Daisuke Motoki, Kosuke Nakago, Wenwen Li, Iori Kurata, Taku Watanabe, Yoshihiro Yayama, Hiroki Iriguchi, Yusuke Asano, al. Towards universal neural network potential for material discovery applicable to arbitrary combination of 45 elements. Nature Communications, (3(1):2991, 2022.
[101] So Takamoto, Satoshi Izumi, and Ju Li. Teanet: Universal neural network interatomic potential inspired by iterative electronic relaxations. Computational Materials Science, 207: 1e211112000 1280, 2022.
[102] Nataliya Lopanitsyna, Guillaume Fraux, Maximilian A. Springer, Sandip De, and Michele Ceriotti. Dataset: Modeling high-entropy transition-metal alloys with alchemical compression: Dataset HEA25, April 2023.
[103] Bingqing Cheng, Edgar Engel, Jörg Behler, Christoph Dellago, and Michele Ceriotti. Dataset: Ab initio thermodynamiis of liquid and solid water, 2018.
[104] Sergey Pozdnyakov, Michael Willatt, and Michele Ceriotti. Dataset: Randomly-displaced methane configurations, 2020.
[105] Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
[106] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages 9323-9332. PMLR, 2021.
[107] Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez- Gonzalez, Yulia Rubanova, Petar Veličković, James Kirkpatrick, and Peter Battaglia, Simple gnn regularisation for 3d molecular property prediction & beyond. arXiv preprint arXiv:2 07971, 2021.
[108] Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Spherical message passing for 3d graph networks. arXiv preprint arXiv:2102.05011 2021.
[109] Philipp Thölke and Gianni De Fabritui.. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations, 2022.
[110] Andrea Grisafi and Michele Ceriotti. Incorporating long-range physics in atomic-scale machine learning. Chem. Phys., (51(20)2004055 November 2019.
[111] Tsz Wai Ko, Jonas A Finkler, Stefan Goedecker, and Jörg Behler. fourth fourth-generatee generation high- dimensional neural network potential with accurate electrostatics including non-local charge transfer. Nature communications, 12(1):398, 2021.
[112] Sergey Pozdnyakov, Artem Oganov, Efim Mazhnik, Arslan Mazitov, and Ivan Kruglov, Fast general two-and three-body interatomic potential. Physical Review B, 107(12): 2023.
[113] Stephen R Xie, Matthias Rupp, and Richard G Hennig. Ultra-fast interpretable machine- learning potentials. arXiv preprint arXiv:21 10.00624, 2021.
[114] Cardinal b-splines. https: www.s.iic....... www.:iii.....ddeeooooo dilsc.eeddddebbor/888777 8wtp:///w8//cc1ee..... .pdf
[115] Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E Sarma, Michael M Bronstein, and Justin M Solomon. Dynamic graph enn for learning on point clouds. Acm Transactions On Graphics (tog), 38 2019.
[116] Emil Prodan and Walter Kohn. Nearsightedness of electronic matter. Proceedings of the National Academy of Sciences, 102633):1635--1163 2005.
17

---

Appendices
A Point Edge Transformer
In this section we provide additional details on the specific implementation of the PET model that use in this work. As discussed in Section D we provide a reference implementation that can provide further clarification on specific details of the architecture.
The transformer that the central component of PET is complemented by several pre- and post- processing units, Two embedding layers, (*( and Ei are used to encode the chemical species of the central atom and all the neighbors, respectively. A 3 - dPEr linear layer ( is used map the 3-dimensional Cartesian coordinates of displacement vectors Fij to feature vectors of dimensionality dpEr, which are further transformed by an activation function a (we use [96
for all our experiments) to form position embedding layer Ex() a (L (ray). An MLP with one hidden layer, denoted as M(²), with dimensions 3dPET dPEr, is used to compress the encoding of the displacement vector, chemical species, and the corresponding input message into a single feature vector containing all this information. The input token fed the transformer associated with neighbor j is computed as follows:
( concatenate(#.) (2(f/t)))) (k--1) E (s;) E (ri;))
(5)
where z!f (((k-1) denotes input message from atom j to the central atom i, 8j is the chemical species of neighbor j. For the very first message-passing block, there are input messages, and M performs a 2dPET - dPET transformation. The token associated with the central atom is given simply E
In addition, each message- peessagepaassing passing block has two post-processing units: two heads H(A) and H(()) to compute atomic and bond contributions to the total energy or any other target at each message
passing block. Atomic contributions are computed as H1( (x((((=) ( where x7(5) is the output token associated with central atom after application of transformer. Bond (or edge) contributions are (R) (1)))))
given H (1i)s (e(riiy/Res AR.), where TJj indicates the output token associated with the j-th neighbor. We modulate them with fe(riy RR,, AR.) in order to ensure smoothness with respect to (dis)appearance of atoms at the cutoff sphere. The output tokens are also used to build outgoing messages. We employ the idea of residual connections (97] to update the messages. Namely, the output tokens are summed with the previous messages.
(k)
To extend the model to describe atomic properties, such as the collinear atomic spins, we simply concatenate the value associated with neighbor j to the three-dimensional vector with Cartesian components of rij. Thus, the linear layer (k) detines 4 + dpEt transformation, in contrast to the 3 + dPET of standard model. Additionally, encode the property associated with the central atom i into the token associated with the central atom. This is done in a manner similar to Eq.
There is great flexibility in the pre-and post- pest-processing. steps before and after the application of the transformer. Our implementation supports several modifications of the procedure discussed above, which are controlled by set of microarchitectural hyperparameters. For example, one can average bond contributions to the target property instead summing them. this case, the total contribution from all the bonds attached to central atom i is given as:
EjEA, H (a (t )f((Te) Re, ARC)
bond (k)
Vi
Ejer, Se(Tij) Re. An.)
(6)
rather than
hend (A)
Σ H (4 (1) (e(ry DR.)
V/
JEA;
Another possibility is to not explicitly concatenate embeddings of neighbor species in Eq, (5). and instead define the input messages for the very first pessage-passiigg block these neighbor species embeddings. This change ensures that all (K) have an input dimensionality of 2dpET across all the message-passing blocks. brief description of these hyper-parameters accompanies the reference PET implementation that discuss in Section D
18

---

B Details of the training protocol
In this section, we provide only the most important details related to benchmarks reported in the main text. The complete set of settings is provided in electronic form, together with the code we used to train and validate models. See Appendix D for details.
Self-contributions. We pre-process all our datasets by subtracting atomic self-contributionss from the targets. These atomic terms often constitute large (if trivial) fraction of the variability of the targets, and removing them stabilizes the fitting procedure. Self-contributions are obtained by fitting a linear model on the train dataset using "bag of atoms" features. For example, the COLL dataset contains atomic species H, C, and O. The "bag of atoms" features for an HzO molecule are [2, 0, 1J. During inference, ceff-contribution are computed using the weights of the linear regression and are added to the predictions of the main model.
Loss. For a multi-target regression, which arises when fitting simultancously on energies and forces, we use the following loss function:
F) I. Eia(-
aF Jr.a
Fin)
L WE MSEE
+
3N
(8)
MSEP
where E' and F are the ground-truth energies and forces, respectively. E' is the energy prediction given by the model, N the number of atoms in the sample, a runs over a, y. 2, and MSEg and MSEp are the exponential moving averages of the Mean Squared Errors in energies and forces, respectively, on the validation dataset. We update MSEE and MSEF once per epoch. We found that for such loss definition the choice of the dimensionless parameter wy is crelatively robust, with the best results achieved approximately for w R 0.03 0.1.
General details. In most cases, we use the StepL R learning rate scheduler, sometimes with a linear warmup. For the cases of the HME21 and water datasets, we use slightly more complicated schemes that are described in the corresponding paragraphs. Unless otherwise specified, we use dpEr = 128, multi-head attention with heads, SiLU activation, and the dimension of the feed- forward network model in the transformer is set to 512. We use a rotational augmentation strategy during fitting, which involves randomly rotating all samples each epoch. This is accomplished using Scipy's implementation (98] of a generator that provides uniformly distributed random rotations. We employ Adam (99] optimizer for all cases. For all models with dPET 128, we set the initial learning rate at 10-1. However, some models with dpEr = 256 were unstable at this learning rate. Consequently, we adjusted the initial learning rate to 5 10-5 5 for two cases: 1) CH2 E+F, 100k samples, and 2) the COLL dataset. We never use dropout or weight decay, but rather avoid overfitting using an early-stopping criterion on a validation set.
VGNN r
C Detailed benchmark results
C.1 COLL
The COLL dataset contains both "energy" and "attoizatiooneeerry" targets, that can be used together with the force data as gradients, and that are entirely equivalent after subtracting celf-contributions. We follow the general training procedure summarized above. The most noteworthy detail is that COLL is the only dataset for which we observed molecular geometries for which all neighbor pairs lead to degenerate, collinear coordinate systems. We tackle this problem by using as fallback modified version of PET as an internally rotationally invariant model, see Appendix F.3 for details. In order to enforce rotational invariance in PET we modify the encoding procedure of the geometric information into input tokens. We encode only information about the length of the displacement vector Tiy. opposed to the entire displacement vector F.y. as in the standard PET. As a result, the model becomes intrinsically rotationally invariant but loses some of its expressive power: the rotationally- invariant -nnvariallyyinvariaatt PET belongs to the class of 2-body GNNs, that have been shown to be incapable of discriminating between specific types of point clouds, see Ref. 32 Nevertheless, the accuracy of the 2-body PET remains surprisingly high. Furthermore, this auxiliary model was active only for 85 out of 9480 test molecules, while all other structures were handled completely by the main PET-256 model. Therefore, the accuracy the auxiliary model has a negligible impact on the accuracy of the overall model. We summarize the performance of several models in Table 2
19

---

Table 2: Comparison of the accuracy of different versions of PET on the COLL dataset, including also the models reported in Ref. [77] PET-256 is the main model with d d 256. 1ET--22--2body is the auxiliary internally invariant model with dpgr dPET 128 and all the other settings match the ones of PET- 1-8-2-2-dddy. PET-ECSE is an overall rotationally invariant model constructed with PET-256 as the main model and PET- 28-2-body as the
128. PET-128 is a normal PET model
auxiliary one. The accuracy SchNet is reported Ref. [7
model
MAE f, meV
MAE E, meV/molecule
SchNet
172
198
DimeNet" ++ 86
40
17
26.4
53
PET- 11T-128-2-body
56.7
27.4
PET- 128
29.8
15.3
PET-256
23.1
12.0
PET-ECSE
23.1
11.9
C.2 HME21
For the HME21 dataset employ learning rate scheduler that differs from the standard StepL.. As shown in Fig. C.2 when StepLR reduces the learning rate, one observes a rapid decrease of validation error, which is however followed by clear signs of overfitting. To avoid this, we implement an alternative scheduler that decreases the learning rate rapidly to almost zero, As shown in Fig. C.2 (orange lines) this alternative scheduler leads noticeable improvement in validation accuracy.
HME21
Water 6 2 4.25
150
30
standard StepLR fast Ir decay
val mae f/ meVIA
val meviA
140
25
130
20
120
standard StepLR
fast decay
15-
110
1.0
1.0
learning rate /le44
learning rate /le-4
0.6
0.6
standard StepLR
fast Ir decay
0.4
0.4
0.2
standard StepLR fast decay
0.2
0.0
50
100
150
200
250
300
350
500
1000
1500
2000
2500
epoch num
epoch num
Figure 4: Comparison of standard Stepl.R R. learning rate schedule and the strategy we use for the HME21 1 and water datasets.
In addition, we slightly modify the loss function for this dataset. The error metric for energies used in prior work is the Mean Absolute Error (MAE) per atom, computed first by determining energies per atom for both the ground truth energies and predictions, then by calculating the MAE for these normalized values. Therefore, it is logical to adjust the part of the loss function related to energies to reflect this metric. We re-define the loss as the mean squared error between the ground truth energies and predictions per atom. This is equivalent to the weighted loss function where weights are given by the inverse of the number of atoms in the structure. We observe that the performance of the model fitted with such modified loss is nearly identical to that obtained when targeting total energies.
20

---

Table C.2
compares the results of PET models with those of recent approaches from the literature. Besides the ECSE-symmetrized PET model, and the non-symmetrized results obtained with single evaluation (which are also computed for other datasets) here we also investigate the impact of test rotational augmentation, which appears to slighlty improve the test accuracy.
Table 3: A comparison between the accuracy of PET and that of SchNet, TeaNet, PaiNN and NequlP (values reproduced from Ref. (100)) and MACE [89]. We compute different PET models, and report both the mean errors (including also the standard deviation between the models) as well as the error on the ensemble models obtained by averaging the predictions of the models. The regularizing effect of ensemble averaging leads to noticeable improvement in accuracy. We also compare the results from a the non-symmetrized model (PET-1), those obtained by averaging 10 random orientations (PET-10) 10) and those with ECSE applied upon inference.
model
MAE E, me
MAE [fl meV/À
MAE f. V/Â
MACE 89
15.7
138
TeaNet 101
19.6
174
153
SchNet
33.6
283
247
PaiNN
22.9
237
208
NequlP 33
47.8
199
175
PET-1
17.8 + 0.1
140.5 + 2.0
124.0 = 1.6
PET-1-ens
16.8
128.1
113.6
PET-10
17.7 + 0.1
139.9 1.9
123.5 1.6
PET-10-ens
16.8
127.9
113.5
PET-ECSE
17.8 + 0.1
141.6 1.9
125.1 1.6
PET-ECSE-ens
16.8
128.5
114
C,3 MnO
We randomly split the dataset into a training subset of 2608 structures, a validation subset of 200 structures, and a testing subset of 293 structures. Thus, the total size of our training and validation subsets matches the size of the training subset in Ref. 88]. One should note that we use all forces in each epoch, while Ref. [88] only uses fraction of the forces typically those with the largest error. Given that the selected forces can change at each epoch, however, all forces can be, at one point, included in the training.
In addition to the position of each atom in 3D space and a label with atomic species, this dataset contains real values of collinear magnetic moments associated with each atom. The continuous values of the atomic spins are obtained by self-consistent, spin-polarized density-functional theory calculation. Even though using this real-valued atomic spins as atomic labels improves substantially the accuracy of the PET model, these are quantities that cannot be inferred from the atomic positions, use them for inference. Similar to what is done in Ref. 88
and therefore it is impractical discretize the local spin values, rounding them to the according to the rule (-oc, -.255] - -1.
(-0.25, 0.25) + 0, [0.25, o0) + 1. These discrete values can be sampled, e.g. by Monte Carlo. to find the most stable spin configuration, or to average over multiple states in a calculation. We also report the error energies of model that doesn't use spins at all, which agrees within two significant digits with the analogous results reported in Ref. [8). This suggests that both models were able to extract possible energetic information from an incomplete description that disregards magnetism.
C.4 High-entropy alloys
The main dataset of High- fathh-Etttpy--lllo configurations, which we used for training, was generated by randomly assigning atom types chosen between 25 transition metals to lattice sites of fee or bec crystal structures, followed by random distortions of their positions [102]. The resulting configurations were then randomly shuffled into training, validation, and testing datasets, following a protocol similar to that discussed in Ref. 990 We used 24630 samples for training, 500 for validation, and 500 for testing. While Ref. 90 uses the energies available for the training dataset, we acknowledge that uses only part of the forces. This, however, can be attributed primarily to the considerable Random
21

---

Table 4: Comparison between the accuracy of the high-dimensional neural network model of Ref. 88 and that of several symmetrsmm)r) PET models. spins-discretized models approximate the atomic spins to integer values, PET-ECSE-spin uses the real-valued spins as atomic labels, while PET-ECSE-no-spile treats all Mn atoms if they had no magnetization.
model
MAE E, meV
MAE f. meV/À
mHDNNP PROONNP-spiiig -eins-ppsiesseeee discretized
1.11
125
PET-ECSE-sppe SeeeE--pinsdiiie apins-discretized
0.30
22.7
PET-ECSE-spins
0.14
8.5
PET-ECSEEnn--pins
11.6
94.8
Access Memory (RAM) demands imposed by the model introduced in Ref(9) The efficiency of PET allowed us to use all the available forces on a standard node our GPU cluster.
In addition, Ref. 90 proposes additional "out-of-sample" test scenarios. In particular, it provides molecular-dynamies trajectories of a configuration containing all 25 atom types, performed at temperatures of T = 300 and T 5000/. For the trajectory at TI' 300K, the errors of the PET model (BCSE-symmelrized) are 9.5 m/a/m and 0.13 eV/A for energies and forces, respectively, which compares favorably with the values of 14 me V/tto and 0.23 eV/Á obtained with the HEA25-4-NN model. For the T 5000 trajectory, which undergoes melting and therefore differs substantially from the crystalline structures the model is trained on, PET yields errors 152 me V/atom and 0.28 ev/Å. HEA25-4-NN yelds lower energy error and comparable force error, 48 me V/atom and 0.29 eV/A respectively.
C.5 Water
We randomly split the dataset [03] into a training subset of size 1303, validation subset of size 100, and testing subset size 190. For the water dataset, we apply a similar trick with fast learning rate decay as we discussed for HME21 abpve. We start fast decaying of the learning rate at the epoch with the best validation error. The right subplot of Fig. C.2 illustrates the improvement achieved by this strategy compared to the constant learning rate for the model with AGNNN 6, nTL = 2 and
R. 4.25A. We use this trick for the ablation studies illustrated in panels b and c of Fig. in the main text and not for those depicted in panel
C.6 CH2
This dataset contains more than 7 million configurations of a CH2 molecule, whose atomic positions are randomized within 3 sphere, discarding configurations where two atoms would be closer than 0.5 [04]. Our testing subset consists of configurations with indices from 3050000 to 3130000 in the original XYZ file. For the validation set, use indices from 3000000 to 3005000. When train a model with Uurain training configurations, we use the ones with indices from 0 to 7 These full validation dataset for fitting the dyer 256 models and 5k subset for all the dPEr = 128 ones in order to reduce the computational cost.
settings match those of Ref. [80] We
For methane, the only metric of interest is the Root Mean Square Error (RMSF) in energies, even when the fitting involves both energies and forces. The loss function we use defined equation (8) This definition includes the dimensionless hyperparameter wg. which determines the relative importance of energies. We have observed that, surprisingly, the best accuracy energies is achieved when wg < 1. For all our training runs, used wg 0.125.
C.7 QM9 energies
We randomly split the dataset into training subset size 110000, a validation subset of size 10000, and testing subset of size 10831. The reported MAE of 6.7 me m//mollcule corresponds to the base model, without application the ECSE protocol. other models, including the current state-of-the-arr model, Wigner Kernels(81].
the accuracy of PET with several
22

---

Table 5: A comparison of the accuracy of several recent models for predicting the atomization energy
of molecules in the QM9 dataset.
Model
Test MAE (meV)
Cormorant 105
22
Schnet 31
14
EGNN 106
NoisyNodes 107
7.3
SphereNet[108
6.3
DimeNet++ [86
6.3
ET1
6.2
PaiNN[88
5.9
Allegro 72
4.7 + 0.2
WK
4.3 + 0.1
PET
6.7
C.8 QM9 dipoles
The dataset from Ref. 82
contains vectorial dipole moments computed for 21000 molecules from the QM9 dataset. These are not to be confused with the sealar values of the norms of dipole moments contained in the original QM9 dataset. We 1000 configurations WK[8]] Our validation subset contains 500 molecules, and we use to 19500 samples for training.
testing subset, similarly
D Reproducibility Statements
In order facilitate the reproducibility of the results present, we have released the following assets: 1) the source code for the PET model, 2) the source code for our prrpicil-p implementation of the ECSE protocol, 3) complete set of hyperparameters for each training procedure, organized as a collection YAML files, 4) similarly organized hyperparameters for the ECSE, 5) the Singularity container used for most numerical experiments 6) all the checkpoints, including those obtained at intermediate stages of the training procedure, and 7) the datasets we used. This should suffice for the reproducibility of all experiments reported in this manuscript. All these files available at:
https://doi.oon 10. 5281/ 7967/ 7967079
In addition, the most important details of the different models trained are discussed in Appendix
E Computational resources
Training times substantially depend on the dataset used for fitting. For instance, fitting the PET model with dpEr = 128 on the CH2 dataset using 1000 energy-only training samples (the very first point of the learning curve in Fig. 3d)) takes about 40 minutes on V100 GPU. In contrast, fitting the PET model with dPET 256 energies and forces using 300,000 samples is estimated take about 26 GPU-days on a V100 (in practice the model was fitted partially on a V100 and partially on an RTX-4090), A similar computational effort is associated with training the dpEr 256 model on the COLL dataset. worth mentioning, however, that the model already improves upon the previous state-of-the- model in forces after approximately one-third of the total fitting time the validation dataset. We used dPET = 256 models only for cases: 1) CH, E+F, 100k samples, 2) CH2 E+F, 300k samples, and 3) COLL dataset. All the other calculations were substantially cheaper. For example, using a V100 GPU it takes about 30h to achieve the best validation error on the HME21 dataset for single dpEr 128 model.
F ECSE
The overall idea behind the ECSE symmetrization protocol is relatively simple, but an efficient implementation requires rather complicated construction, that discuss here in detail.
23

---

fc(r|Rc = 4, AR. 2)
k a. =
d2z|w= 1, A.
1.0
0.5
1
2
3
0
1
3
0
2
Figure 5: Smooth cutoff functions fee d and o.
F1 Preliminary definitions
We begin by defining several mathematical functions that in what follows. In order obtain smooth predictions, PET uses several cutoff functions. While there is a degree of flexibility in defining them, found the following functional forms, that are also illustrated in Fig. [ to be effective:
if r S R no
fe An
- + 1); x 2 R.t0..AR (tanh 1 a+1
if R. An. <T < Re
ARc
if 2 R
(9)
0,
i w if < (10)
(=lw,, Aw) =
2 tanh
x+1 + x-1 + 1);x 2 22-5-0.54
Aw
1,
if= Au
0,
if
(2 Aw)
(-
tanh
x+1 + z-1) 1): 2--------00
0.5A A..
<fw<=<w+Ai (11)
2,
if=
We also need smooth functions prune the ensemble of coordinate systems. To do so, several variations on a well-known smooth approximation for the max function. Given a finite set numbers { and smoothness parameter one can define:
Smoothhaxx(::)1333 E, exp((rr)a
(12)
Ee exp(Br.)
This function satisfies the following properties:
SmoothMax < max(( (2)))))
8 (axothMaa(((2)1333) = max((x;)))
lim
(13)
For set of just two numbers (21.22) one can show that SmoothMax is bounded from below:
SmoothMax( (21, x2)) + T(3) 2 max max((211 x2)) lim T(3) 0
(14)
for T(3) = ((exp(-1))/3 where W is the Lambert W function.
For the case when the numbers æ; associated with weights p; 2 0 it is possible extend the definition of SmoothMax as:
SmoothMax WeenhaMMeWWigge((((( (x Σ, (xxp(3x:)pa:;
(15)
24

---

This function satisfies the properties in Eq. 13, where maximum is taken only of the subset of {z;} where p; O. useful consequence of the definition of SmootiiMaweightee is that its value not changed when including point with zero weight, (z: 0):
MaxoethhaaWeighhaaee (x), U {(x, 0)} = SmoothMax WeeighheedeWeeeeeeeeeeee (t).)) (16)
Corresponding functions that provide a smooth approximation to the min operation, SmoothMin and Amoothhineeighhee can be defined the same way as SmoothMax and SmoothMaxWeightce using (-3) instead of B.
Finally, given an ordered pair of noncollinear unit vectors 01 and و we define the associated coordinate system by computing un = új x d2 ún = un / u| and ú = 6 x ūjj The reference frame is given by the vectors é1, úi and 02. This procedure generates only right-handed coordinate systems, and thus, the transformation between the absolute coordinate system to the atom- one can always be expressed the coordinate system given by a central atom and two neighbors j and j", we mean that the vectors oi and di are the ones pointing from the central atom to neighbors, j and j' respectively.
proper rotation, never requiring inversion. When we later refer to
F.2 Assumptions
We assume that (1) there is finite number of neighbors within any finite cutoff sphere and (2) there is lower boundary dmin to the distance between pairs of points. Neither of these assumptions is restrictive in chemical applications, because molecules and materials have a finite atom density, and because atoms cannot overlap when working on relevant energy scale.
F.3 Fully-collinear problem
The simplest version of ECSE protocol was defined in the main text in the following way:
(s.))) = Σ
way R Rer.rryyooooooooo (o HADI
Σ
Wiy's
(17)
DE'EA:
where
W3y' (r fc(r,) Re, AR)fc(ry RR, AR) x fy Aw). The de function can be chosen either dt or 2 while w and Aw are user-specified constant parameters.
(18)
This simple formulation doesn't ' handle correctly the corner case of when all pairs of neighbors form nearly collinear triplets with the central atom. In this scenario, for all pairs of j and j'. |f fy is smaller than w which implies that all the angular cutoff values g=()i f Aw) zeros, and so all the weights Wjy'- As result, the ECSE weighted average is undefined, falling ambiguity. We propose two solutions to treat this case.
First solution to the fully-collinear problem. The first solution we propose is to incorporate predictions of some internally equivariant model in the functional form of ECSE:
/
(s(A,) = Wanx. aux (Ai) +
Σ W5y Roy (vo
( (2AA))))))))
Waux +
Σ
Wis'
(19)
J
J''EA:
In this case, (A,) can fall to yanx ) all the weights Wjj' are zeros.
Under the assumption that the (typically simpler) auxiliary model will be less accurate than the main backbone architecture, it is desirable to ensure that it is only used make predictions for corner cases. For this purpose we define Waux as follows:
Wanx
MaSmoothMaxxWigliini
(Wizr, Wys' (8w) ttuxs Amxx).
(20)
With this expression, Waux takes a non-zero weight only when the smooth maximum of the coordinate- system weights Wjjr is below t in other words, as long as at least one pair of neighbors is non-collinear, the auxiliary model will be ignored. Since SmoohhMsseighted does not depend on quantities with zero weight (see Eq. (16), one can compute Eq. (20) efficiently, without explicit iterations over the coordinate systems with zero weight. Bwr laux and A are user-specified constant parameters.
25

---

Second solution to the fully-collinear problem. Even though the use of an auxiliary model is an effective and relatively simple solution, it would be however desirable to use consistently the non-equivariant backbone architecture. In this section, we discuss a sketch of how one could tackle this corner-case without the usage of an auxiliary model.
To eliminate the 0/0 singularity, we need to use an adaptive definition for the angular cutoff parameters and A. that enter the definition (18) of the ECSE weights:
SmoothMax(((177 11,1)))))1))
(21)
With this definition, the weights are never zero except for the exact fully-collinear singularity, which makes the weighted average of the ECSE protocol well-defined everywhere. For environments approaching the exact fully-collineds singularity, the coordinate systems used by ECSE are not stable, However, given the definition of the coordinate systems discussed in Appendix [.I the axes of all the coordinate systems are aligned along the collinear direction. Given that the atomic environment itself is nearly collinear, the input to the -onevauiva models is always the same, irrespective of the y and & axes, that will be arbitrarily oriented in the orthogonal plane. Thus, all evaluations of the backbone architecture Yo lead to approximately the same predictions. Note that this is only true if all the neighbors that are used for the backbone architecture are also considered to build coordinate systems, which requires that the cutoff radius employed by ECSE is larger than that of the backbone architecture. When atoms approach smoothly a fully-collinear configuration, the coordinates that are input to the backbone architecture converge smoothly to being fully aligned along æ for all coordinate systems, making the ECSE average smooth.
One small detail is that, within this protocol, the orientation of the axis can be arbitrary. This problem is easily solved by always including an additional coordinate system in which 01 is oriented along r,j rather than along Fsj. Both coordinate systems should be used with the same weight Wyz'. One should note, however, that this approach has a number of limitations. In particular, it doesn't 't support covariant inputs, such as vectorial spins associated with each atom, and covariant outputs, such as dipole moments.
F.4 Adaptive inner cutoff radius
The minimal inner cutoff radius Rin should be chosen to ensure that it includes at least one pair of neighbors that define a "good" coordinate system. To determine an inner cutoff that contains at least a pair of neighbors, we proceed as follows
1. For each pair of neighbors, we use SmoothMax((ssr4113) to select the farthest distance in each pair
2. Applying a SmoothMin function to the collection of neighbor pairs
T(8))r13))
(22)
selects distance that contains at least pair of neighbors. That this function is always greater or equal to the second-closest pair distance follows from the inequality (14) and the SmoothMin analogue of Eq. (T3).
3. We ensure to pick at least one non-collinear pair (if there is one within Row) by defining the adaptive inner cutoff as
Rm MinothMin Weighted ShnntttiiiiiiggglleefSfoeaaoaaar
+ hir U ((Roat, 1.0))11)
(23)
where Pij' are defined as:
Pos fe(ry|Rout AR. fe(y) Rou- AR. 34. (Pty x fy. + Aw).
(24)
j and j' run over all the neighbors inside the outer cutoff sphere.
This definition maintains smoothness with respect to the (dis)appearance of atoms at the outer cutoff sphere. It ensures that if there's 's at least one good pair of neighbors inside the outer cutoff sphere, there's at least one within the inner cutoff as well. no such good pair exists within the outer cutoff sphere, the definition falls back to Ran- Note that the expression includes also some "tolerance" parameters AR. and Ag. These are introduced to ensure that at least one of the weights computed to define the equivariant ensemble in Eq. 24 is of the order of I.
26

---

F,5
Weights pruning
Given that the cost of applying ECSE depends on the number of times the inner model needs to be evaluated, it is important to minimize the number of coordinate systems that have to be considered. To do so, one can e.g. disregard all coordinate systems with weights below a threshold, e.g. half of the maximum weight. A smooth version of this idea can be implemented as follows:
Constant parameters:
3w Bw >
Ty, AT, T > 0, Ar, > T, <
Te, Ar. To 0, Ar. > 0r < Te (T stands for Threshold).
Smooth weights pruning:
1: WMAX + SmoothMax Weighted {{ (way, (ay))
I 2: fir + | WMAX
3: e +- Se(wMAx| Te. AT.)
4: W'y' + ewoy + (1 E1
First, one computes the maximum weight present, using a SmoothhMaWWeighted function. Then, pruning factors fay are computed using the cutoff g.' function, aiming to zero out the smallest weights with W'z' <- W/y' Joy An additional modification is needed to ensure smooth behavior around fully- collinear configurations. where all weights become zero. If this occurs, 9% (wy [uMAxT', WMAX Ar,) ) converges step function. Therefore, makes sense activate weight pruning only if wMAX is larger than certain threshold Te, which is implemented using smooth activation factor e. We found this pruning to be more efficient the de function in Eq. (4) is selected to be q² rather than g!- For multi-component systems, an alternative approach (which can be used on top of the discussed strategy) to increase efficiency involves using only a subset of all the atoms of specific atomic species to define coordinate systems (e.g., only O atoms in liquid water).
F.6 Tradeoff between smoothness and computational efficiency
For any set of selected parameters, the ECSE protocol yiclds a smooth approximation of the target property in a rigorous mathematical sense, meaning the resulting approximation of the target property is continuously differentiable. The important question, however, pertains to the smoothness of this approximation in a physical sense. That is, how quickly it oscillates or, equivalently, how large the associated derivatives are. In this section, we discuss the tradeoff between "physical" smoothness and the computational efficiency of the ECSE protocol. This tradeoff is controlled by the user-specified parameters discussed in the previous paragraphs.
an example of such parameters, with a clear impact on the regularity of equivariant predictions. Among other things, it is used to determine the inner cutoff radius, Rin. which is defined along the lines with a smooth approximation of the second minimum of all the distances from the central atom to all its neighbors. Assigning relatively low value to 3 generates a very smooth approximation of this second minimum, but it significantly overestimates the exact value, and therefore R., potentially includes more neighbor pairs than needed. In contrast, for large value of [3, the second minimum approximation reflects the exact value of the second minimum more precisely, at the cost of decreasing the smoothness of the approximation when the distances between atoms close to the center vary. Even though the ensemble average remains continuously differentiable for any 3 > O, it might develop unphysical spikes in the derivatives, reducing its accuracy. To see more precisely the impact of consider the limit of large value (considering also that other adjustment parameters such as AR. are set to a value that does not entail an additional smoothening effect). In this limit, the inner cutoff selected by the ECSE protocol encompasses only one pair of neighbors for a general position, so most of the time the ECSE protocol uses only one coordinate system defined by the two closest neighbors. Only when the nearest pair of atoms changes, within a narrow transition region, two coordinate systems are used simultaneously. The value of /3 determines the characteristic size of this transition region. Thus, for excessively large values of B, there is a sharp (but still continuously differentiable) transition from one coordinate system to another. On the other hand, 3 is set to a very low value, R., will always be significantly larger than the exact distance to the second neighbor. As result, the ECSE protocol will constantly utilize multiple coordinate systems, leading to a smoother symmetrized model but with higher computational cost.
27

---

In summary, the parameters used by the ECSE dictate the balance between smoothness and compu- tational efficiency. We name the selection of such parameters that leads a very smooth model as "loose". The opposite selection that minimizes computational cost is labeled as "tight". The proper selection of these parameters directly impacts the accuracy of the symmetrized model in forces (and in higher derivatives), which we discuss in the next section.
E.7 Accuracy of e-postefoor--ECSE-yymmetttt
models
The most computationally efficient method for fitting model is to initially train non-equivariant backbone architecture with random rotational augmentations and then apply the ECSE protocol posteriori. There are general considerations symmetrized model relative to the initial, non-equivariant one.
can make the expected accuracy of the
The functional form of ECSE is weighted average of predictions generated by backbone archi- tecture for several coordinate systems. Therefore, it can be seen as implementing the rotational augmentation strategy during inference. This implies that the accuracy of ECSE in energies (and in any other direct targets) is expected to be bounded between the accuracy of standard single inference of the backbone architecture and that evaluated with fully-converged rotational augmentations.
Forces, however, present more complicated situation. The forces produced by ECSE are not just linear combinations of the predictions given by the backbone architecture. This complexity stems from the fact that the coordinate systems and weights used by ECSE depend on atomic positions themselves. Since forces are defined by the negative gradient of energy, it is necessary to take the derivatives of Eq. (17) with respect to atomic positions to obtain the predictions of ECSE. Additional terms related BRyg to "Pk and Bw'zz 85% are not associated with the forces produced by backbone architecture @y.((A.)) Consequently, the error in forces the final, symmetrized model may significantly exceed that of the backbone architecture. This issue is more pronounced for tight selection of ECSE parameters and less pronounced for a loose one. In principle, this can be mitigated by fine-tuning the overall symmetrized model. In this work, we have selected loose parameters of the ECSE protocol sufficient to ensure that the difference in accuracy between the backbone architectures and symmetrized models is minuscule, even without fine-tuning, which results however in more pronounced computational overhead.
dr,
F88 Additional rotational augmentations
In some cases, such as for the CH4 configurations, the total number of all possible coordinate systems very limited. However, using a larger number of rotational augmentations may still beneficial during inference. For such cases, it's possible to incorporate additional augmentations into the ECSE protocol as follows:
Σ Σ
Wiy Rig Raug Noo (Rmug R.g Ac.)]
JA R
(s,),)
(25)
Naug EA
Wig
where Saug is a predefined constant set of random rotations Rang.
F9 Application of the ECSE protocol to long-range models
Several models capable of capturing long-range interactions have been developed within the atomistic machine-lcarming community. These models cannot be directly cast into the local energy decomposi- tion discussed in Sec. 2 However, they are constructed as either pre- post-processing operations for the local models. Thus, these schemes can also benefit from the ECSE protocol by making possible to use some models developed for generic point clouds local engine.
An example of such models is [IO. As a first step, it constructs a representation that is inspired by the electrostatic potential. This three- diree-dimeiiional function, or collection of three- dimensional functions for each atomic species, captures the long-range information. This global potential is then used to define local descriptors, that are combined with other short-range features to model a local energy decomposition. Thus, the ECSE protocol can applied LODE, or similar schemes, without any modifications.
28

---

Another example of Longerange model is given in Ref. [. This scheme applies local models to compute partial charges associated with each atom and later uses them to calculate non-local electrostatic energy contribution. The ECSE protocol can ensure that these partial charges remain invariant with respect to rotations.
E.10 Application of the ECSE protocol to message-passing schemes
As already discussed in the main text, message-passing schemes can also be considered within the framework of local models, with their receptive field playing the role of a cutoff. Thus, the ECSE protocol is directly applicable to them. Whereas this approach has linear scaling, it is very computationally expensive. As an example, consider a GNN with 2 message- paassigg layers. At the first layer, atom A sends some message, denoted m AB. atom B. At the second, atom B sends messages mBe and m2p to atoms C and D, respectively. The second layer additionaly computes predictions associated with atoms C and D, given all the input messages. If one naively applies the ECSE protocol, it is first employed to get a rotationally invariant prediction associated with atom C and then to get prediction associated with atom D. Both calculations rely on the message m computed at the first layer of the GNN. The problem is that, in contrast to normal GNNs, this message cannot be reused compute predictions associated with atoms C and D. The coordinate systems defined by the ECSE protocol for atoms C' and D do not match each other, and thus, the message m' should be recomputed for each of them. The problem becomes progressively more severe as the depth of the GNN increases.
If the size of a finite system the size of unit cell is smaller than the receptive field of GNN, one can accelerate and simplify the discussed approach by defining a global pool of coordinate systems. This can be achieved by applying the ECSE protocol all atoms, subsequently unifying all the coordinate systems and associated weights. If one employs pruning techniques similar to the one described in Appendix F.5 this approach can become highly efficient, outperforming the naive one even if the size of the system exceeds the receptive field of the model. For finite systems, an even more efficient scheme can be designed by applying the ECSE protocol single point associated with the whole system, such as the center of mass.
The most computationally efficient way, however, is to apply the ECSE protocol layerwise, sym- metrizing all the output messages at each layer. In this case, one must be especially careful with the fitting scheme. The standard scheme of first training a non-invariant model with random rotational augmentations and then applying the layerwise ECSE protocol a posteriori could lead to a significant drop in accuracy. The strategy of random rotational augmentations during training corresponds to applying random rotation to an entire atomic configuration. Thus, during training, a model can learn to encode some information about relative orientations into the messages. However, if one then applies the ECSE protocol layer by layer, the coordinate systems of the atom that sends a message do not match those of the receiving one. Consequently, the information about relative orientations can no longer be propagated.
This problem can be solved by using individual augmentation strategy that stands for rotating all atomic environments independently. In this case, the model can encode only rotationally invariant information into the messages. Thus, can expect that similar accuracy will be achieved if one enables a layerwise ECSE protocol a posteriori. Furthermore, it is possible to extend this approach to covariant messages. One can specify how the messages should transform with respect rotations, for example, as vectors or tensors of rank L. Then, when performing individual rotation augmentations, one should explicitly transform the messages from the coordinate system of the sending atom that of the receiving one. The layerwise ECSE protocol should also be modified accordingly.
To sum up, the efficient layerwise ECSE protocol supports covariant messages with any transformation rule. However, this transformation rule should be specified as part of the architecture. The global ECSE protocol is more powerful, as it allows the model to learn behavior of the messages with respect to rotations on its own.
An interesting question that we put out of the scope of this paper is the behavior of the messages with respect to rotations for the layerwise scheme if the ECSE protocol turned on immediately and symmetrized model is trained from the very beginning.
29

---

F.11 Current implementation
Our current implementation uses global pool of coordinate systems, as discussed in Appendix F.10 In addition, it implements most of the logic of the ECSE protocol using zero- pero-dimensional Py Torch tensors, each containing only one scalar. Thus, it has the potential to substantially accelerated. We want to reiterate that we have selected loose parameters for the ECSE protocol corresponding to the best aceuracy and substantial computational cost (see more details in Appendices |.66 and F.7 Therefore, even with the current coof-of-prrnncpl. implementation, ECSE can be substantially accelerated at the cost of a minuscule drop in accuracy.
F.12 Numerical verification of smoothness
To verify that implementation of the ECSE protocol indeed preserves the smoothness of the backbone architecture, we numerically analyzed alterations the predictions with respect to random perturbations of the imput. For this purpose, we selected the CH2 dataset (discussed in the Appendix C since it represents random configurations, thus covering most of the configurational space. For 100 molecules from this dataset, applied 50 random Gaussian perturbations all the positions of all the atoms with different amplitudes. For each perturbation, the corresponding change in the prediction of the ECSE symmetrized model was measured. As backbone architecture, used
PET model with hyperparameters similar to the one benchmarked in Appendix The parameters of ECSE were selected be loose in sense discussed in Appendix [F66
Results are shown in Fig. 6 One can see that a small amplitude the applied noise entails a minor alteration of the predictions. No instances were observed where a minor noise led to significant change the predictions of the ECSE symmetrized model.
ECSE symmetrized model Non-symmetrized model
10
10-3
Difference in predictions
10
10-7
10
10-11
10 13
10-9
10-7
10-5
10
10-1
L2 norm of the noise
Figure 6: Numerical verification of preserving smoothness. Each point represents one random perturbation of one CH, molecule. The horizontal coordinate is an amplitude of Gaussian noise applied to all the positions the atoms. Vertical coordinate represents the difference in the predictions. The experiment conducted for both symmetrized and non-symmetrizeed PET models.
30

---

a
B
(x)
b
2x + 6
B
B
B2
0.9
x² 2
6x 9
0.6
2
0.3
0.0
2
3
X
X
Figure 7: (a) B-spline of 2nd order. (b) Equidistant collection of B-splines. The figure is adapted from Ref. 112
G
Architectures for generic point clouds made smooth
In the main text, stated that most architectures for generic point clouds can be made smooth (continuously differentiable) with relatively small modifications. Here we provide a few practical examples of how non-smooth point cloud architectures can be modified.
G.1 Voxel and projection based methods
For naive voxelization point cloud, voxel features are estimated based on the points inside the corresponding voxel. This means that this approach leads to discontinuities whenever point moves from one voxel to another. The construction used in GTTP [12] and [F] is an example of smooth projection of point cloud onto the voxel grid. We first describe the smooth projection onto one-dimensional pixel/voxel grid for simplicity. relies on smooth, piecewise polynomial, bell- like functions known as B-splines.
Fig- 7 demonstrates the shape of second-order B-spline. It spreads three one-dimensional voxels, on each of them it is defined by second-order polynomial, and is continuously differentiable. For any p. is possible construct an analogous function 8 (z) which spreads over + 1 voxels, is piecewise polynomial of order and is p 1 times continuously differentiable. Next, an equidistant collection of such B-splines is constructed, as shown in Fig. (b). The projection onto the one-dimensional voxel grid is defined as follows:
G Σ ((()..
(26)
where my is the position of a k-th point. The coefficients & can be treated as features associated with the i-th voxel.
Three-dimensional B-splines can be built tensor product of these functions, i.e.:
B igis (x. y. 2) = B,, Ba) (v) B. (z).
(27)
and the corresponding projection of a point cloud onto a 3D voxel grid is:
Cigigis
Σ Bi (rx).
(28)
An alternative method define smooth projection is to define smooth point density, and to compute its integrals for each voxel, as demonstrated in Fig. for the one-dimensional case. The point density is defined follows:
p(r) Σ B(r rt).
(29)
31

---

Point positions
splines centered at point positions
point density
Point density
x
Figure 8: Integral projection. The integral of the red area determines the features of the third left voxel.
where any bell-like function. contrast to the previous example, here, bell-like functions are associated with each point rather than with each voxel.
The integral can be rewritten into the form of the first projection method:
/ p(r)dr= =
Σ
Bin irizig (Ix),
(30)
voxel f13923
where
/
Bint
i (Fk)
B(r-rs)dr.
(31)
voxel 111213
If B is given by B-spline of the p-th order, then Bint is given by a B-spline of + 1-th order 114
Since the B-spline functions have compact support, the discussed projections preserve the sparsity of the voxel features.
Within paradigm of local energy decomposition, one center the voxel grid the position of the central atom. When building a projection of an atomic environment onto a local voxel grid, one last source of discontinuity is associated with the (dis) (dis)appearance of new atoms at the cutoff sphere. This problem can be avoided by modifying Eq. (28) as:
Gi = Σ fe(ra Re, AR. B. Biii (rx),
(32)
and Eq. 22 as:
p(r) Se(rrlRR Brr rk).
(33)
(rk here denotes the displacement vector from central atom the k-th neighbor). Finally, the 3D convolutional NN applied on top of the voxel projection should be smooth. One can ensure this by selecting smooth activation function and using smooth pooling layers, such as sum or average, or using a CNN without pooling layers. The ECSE protocol applied on top of this construction adds the only missing ingredient rotational equivariance which makes voxel-based methods applicable for atomistic modeling. The paradigm of local energy decomposition is not the only way to use methods. For instance, for finite systems of moderate sizes, such as molecules and nanoparticles, one can center the voxel grid at the center of mass and apply ECSE protocol to the same point. Although intuitively, projection- based 2D methods are not expected to perform well given that the 3D information is crucial for estimating the energy of an atomic configuration, in principle, they can made smooth and applied using the 2D analogue of the methods discussed above.
32

---

G.2
Point models
common building block of various point methods is the following functional form
48 115
x = ([(((((( E A,)),
(34)
where is feature vector associated with central point i, X, are feature vectors associated with neighbors, 7 and h are learnable functions, is permutationally invariant aggregation function, and the neighhorhood A, is defined either as the & nearest neighbors as all the points within certain cutoff radius R.- The distance used to define local neighborhoods can be computed either in the initial space or between the feature vectors at the current layer of the neural network.
The construction above contains several sources of discontinuities even for smooth functions y and h. First, the definition of the local neighborhood as the k nearest neighbors is not smooth with respect to the change of a set of neighbors; that is, when one point is leaving the local neighborhood, and another is entering. Thus, one should instead use the definition of a local neighborhood with fixed cutoff radius. In principle, one can make adaptive, meaning could be a smooth function of neighborhood, embracing k neighbors on average.
I is given by summation, then, as it was already outlined in the main text for point convolutions, the smooth form can be constructed as:
x
E fe(dey Re. M(x,,
(35)
where di is the distance between points and
I is given by maximum, then one can ensure smoothness using an expression such as:
x 7(Smooth ((. (day Re, AR. Mx xz)li E weighborhood((x)))
(36)
where fa (d.y Re, AR.) is some smooth function which converges to -co at Rg. Minimum and be adapted analogously.
average aggregations
We are unaware of any simple modification to ensure smoothness in the downsampling operations that applied in PointNett+-like methods, and are typically performed by the Farthest Point Sampling algorithm. However, this operation is probably unnecessary for architectures designed for atomistic modeling due the locality [116) of quantum mechanical interactions and the fact that many successful domain-specific models do not use anything resembling a downsampling operation.
To sum up, many neural networks developed for generic point clouds can be straightforwardly made smooth and, thus, can be made applicable to materials modeling if using the ECSE protocol. For another large group of architectures, we are currently unaware of how to ensure smoothness for entire models because of operations such as downsampling. However, in many such cases, core building blocks, neural network layers, be made smooth and used to construet PointNet-like architectures for atomistic applications.
H Broader impact
Machine learning force fields can significantly accelerate progress in areas such as drug discovery and materials modeling by alleviating the substantial computational costs associated with ab-initio simulations. Even though training models involves substantial energy cost, the energy cost of inference is minuscule in comparison with that of the reference electronic structure calculations, which leads to overall energy savings. Although these models could theoretically be used for malicious intentions, we argue that the likelihood is minimal. This stems from the fact that machine learning force fields are not direct agents of impact. Instead, they function primarily as scientific tools, typically deployed within the framework of organized research institutions rather than being exploited by single individuals with malevolent objectives.
33


Reviews:
[{"review_id": "RKWH9UDtwG", "reviewer": "Reviewer_r8ta", "summary": "The paper suggests a network design tackling chemical and materials modeling. The paper identifies rotational symmetries as lacking in existing solutions. The suggested approach is based on symmetrization, i.e., projecting a general function to be equivariant by averaging it over aligning rotations. In addition, a transformer-based architecture is proposed.", "questions": "I would appreciate any response from the authors regarding the weakness stated above.", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The idea to enforce smoothness with other symmetries such as rotational symmetries seems to be valuable, and indeed missing from existing solutions.", "weaknesses": "The main weaknesses of the paper are related to its level of readability, the quality of its exposition, and its mathematical formulation. In what follows, I will share some relevant examples. \n\nIt is understood that the selected approach to incorporate rotation equivariance is symmetrization. However, the text does not make it clear what are the challenges in doing so and why the starting point chosen is using frames from pairs. If that is the starting point, it is not entirely clear what is the next point. In addition, it is not shown why this symmetrization is indeed equivariant. In addition, there so no clear and simple mathematical formulation of how the aligning rotations for the symmetrization are selected. Furthermore, all the input and output dimensions are missing from all the functions stated in the paper. Figure 2 is hard to decipher.\n\nIt is hard to follow the experiment section. I would expect to see some ablation regarding the symmetrization technique. I would expect to see the symmetrization tested with other architectures than the transformer as well. Furthermore, it seems that symmetrization has a minor effect on the model’s generalization capacity, making the benefit from it, given its computational budget increase, questionable.\n\nLastly, the contribution of the proposed transformer architecture seems to be orthogonal to the challenge the paper set to tackle.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "A large part of the perceived weaknesses is, in our opinion, connected to the fact that were not sufficiently clear in explaining that the issue of achieving exact equivariance is mainly related to the very applicability of the models for atomistic simulations. A slight improvement in accuracy is just a fortunate side effect. We discuss this point in the global response, but we would like to address specifically some of the points raised here. \n\n > However, the text does not make it clear what are the challenges in doing so and why the starting point chosen is using frames from pairs.\n\nThe challenge is to achieve exact rotational equivariance while preserving permutation equivariance and smoothness of the model, as we discuss in the introduction and in Sections 4 and 5. Defining a reference system that is rigidly attached to a group of atoms is a straightforward way to obtain equivariant predictions. It is however very hard to do so while preserving smoothness. We review two approaches that define a single coordinate system to ensure rotational equivariance. The first uses a single pair of two closest neighbors, while the second utilizes eigenvectors of the tensor of inertia. Both methods lead to discontinuous predictions. Thus, our method, which sums over an ensemble of local coordinate systems, is the first one which does rotational symmetrization while preserving smoothness. This idea is the “starting point” of our construction.\n\n> If that is the starting point, it is not entirely clear what is the next point.\n\nThe next point is a discussion of the modifications that have to be applied to make this idea computationally efficient, by reducing the number of coordinate systems without affecting differentiability. They are described in the paragraphs \"Coordinate weighting\" and \"Adaptive cutoff and limiting cases.\" Details of implementation are given in Appendix F. \n\n> In addition, it is not shown why this symmetrization is indeed equivariant.\n\nThe prediction is computed as a weighted average of predictions of the base model for several coordinate systems rigidly attached to an input point cloud. It means that all the used coordinate systems rotate together with the given input. This, in turn, leads to equivariance. Mathematically, one can consider that in Eq. (3) every rotation applied to the structure would automatically be applied to all the $\\hat{R}_{jj'}$, which directly translates into proof of equivariance. \n\n> In addition, there so no clear and simple mathematical formulation of how the aligning rotations for the symmetrization are selected.\n\nEach coordinate system used for symmetrization is built based on a triplet of atoms. One of the atoms is always the central one, and two others are neighbors within an adaptive cutoff radius. This is stated in Section 5, “$\\hat{R}_{jj'}[\\cdot]$ indicates the operator that rotates its argument in the coordinate system defined by the neighbors $j$ and $j'$ within the $i$-centered environment $A_i$”. \nThe gory details of how we reduce the number of coordinate systems in a smooth manner is discussed in detail in the Appendix F. \n\n> Furthermore, all the input and output dimensions are missing from all the functions stated in the paper.\n\nWe apply the common convention that vectorial quantities are typeset in boldface. All other quantities and functions in the main text are scalars. The only exception is the dimensionality of the output of the base model in the description of our symmetrization scheme in section 5. However, since our symmetrization method is very general, this dimensionality can be arbitrary. In the Appendices, we indicate the dimensionalities of all multidimensional objects. \n\n> It is hard to follow the experiment section. I would expect to see some ablation regarding the symmetrization technique. Furthermore, it seems that symmetrization has a minor effect on the model's generalization capacity, making the benefit from it, given its computational budget increase, questionable.\n\nWe perform this ablation study, as we report both “raw” and symmetrized accuracy for each benchmark. The purpose of our symmetrization technique is not to improve accuracy. Instead, it is necessary for the very applicability of ML models to atomistic applications. A minor improvement in the model's accuracy is just a fortunate side effect. We provide more details in our global response. \n\n > Lastly, the contribution of the proposed transformer architecture seems to be orthogonal to the challenge the paper set to tackle.\n\nAs stated at the beginning of Section 6, we introduce PET to “demonstrate that lifting the design constraint of rotational equivariance makes it possible to construct better models.” The requirement to have exact equivariance has led to almost exclusive focus on models that are intrinsically equivariant, which greatly restricts the design space for new architectures. Being able to achieve exact equivariance *a posteriori* lifts these design constraints, and we need PET to prove that there is a practical advantage in doing so. "}, {"review_id": "SKWQyvc7RS", "reviewer": "Reviewer_mhUe", "summary": "This paper presents a novel protocol for general symmetrization that adds rotational equivariance to any given 3D point-cloud processing model while preserving symmetries with respect to translations and permutations.\nApplying the protocol to a model a-posteriori (i.e., at inference) makes the model rigorously rotation-equivariant.\nTo demonstrate the potential use of the protocol, the authors present Point Edge Transformer (PET). \nAlbeit not intrinsically equivariant, it achieves state-of-the-art performance on several benchmarks of molecules and solids.\n\n---------------------------------------\nPOST REBUTTAL\n\nGiven the clarification the authors provided by the rebuttal, as well as the additional demonstration of the robustness of their proposed method wrt. noise, I have raised my score to 5.", "questions": "Q1. In the abstract, it says (line 11), \"We propose a general symmetrization protocol that adds rotational equivariance to any given model\". However, the symmetrization seems to assume a base model based on the atom-centric local coordinate system. How common is this? Does this assumption cover the related work, or are there methods that do not use this?\n\nQ2. In the analysis (F.8), it is said that one can improve performance further by introducing extra rotations at inference (not on input but on extracted coordinate systems, so the model is still technically invariant/equivariant). I would like to know why this improves performance.\nIf I understood correctly from the training description, it is done in a non-equivariant way with rotational augmentation. Is this the reason for the improvement?\n\nQ3. Given that the most common metric is MAE, have the authors tried using MAE as training loss?\n\nQ4. Figure 2: $\\bar{x}$ seems to refer to edges between atoms in $A_i$, but as far as I can see, $x$ has not been introduced before. The same applies to $E$ and $M$ and $H$. These terms should be clearly introduced in the text before the figure, not in the Appendix.\n\nQ5. It appears that **r**$_i$ is the 3D position per atom, and $A_i$ is a group of atoms centered around **r**$_i$. However, it was unclear how features per atom are extracted from the beginning. Are different atoms treated as one-hot embeddings?", "limitations": "The authors adequately address the limitations and broader impact.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "S1. The proposed idea of training a model that is intrinsically non-SO(3)-equivariant and making it equivariant during inference is novel, and its mathematical presentation appears correct.\n\nS2. The proposed experimental evaluation is extensive and substantial.", "weaknesses": "W1. Contributions: The main contribution the authors claim is the equivariance scheme. \nUnfortunately, there is not much in the article to support this. Instead, the bulk of the improvement appears to come from the PET model, which is already sufficiently equivariant across multiple tasks. \nIt is unclear what benefits smoothness gave in practice (I might have missed something in the experiment part).\n\nW2. Presentation: The presentation is a bit convoluted and seems too verbose; the provided Appendix contains many important details, e.g., about the PET architecture, that should instead be in the main paper.  Furthermore, I had a hard time finding a precise definition of the problem in the paper. The paper addresses an important question of rotational equi- and invariance, but its vague presentation makes it difficult to appreciate the contributions.\n\nW3. Related work: Somewhat lacking related work on geometric deep learning. There is a detailed overview of related work on GNNs, but almost no discussion about group equivariant CNNs (G-CNNs). It is unclear how different the proposed model is and if the proposed protocol could be applied to such models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1** The reviewer is correct in that the bulk of the accuracy improvement comes from the PET model. However, accuracy is not the primary reason to enforce symmetry or smoothness in our primary domain of application. We discuss this point in our global response, but in short the physical constraints in atomistic modeling make exact equivariance and smooth differentiability a strong requirement. \n\nPerformance in benchmark examples is important, and an objective metric to compare models, but most practitioners in atomistic modeling would choose a rigorously equivariant model over one that is only approximately so, even if the latter was more accurate on benchmarks. \n\n**W2** The space constraints forced us to move to the Appendices the detailed description of our approaches, but we had tried to keep the core ideas in the main text. This might give the impression of a somewhat superficial treatment, and we will attempt to highlight more the key ideas. Similarly, it appears that we have not been able to convey the importance of exact equivariance in our (and related) research fields. Besides further clarifying this problem in the abstract and introduction, it might be useful to reiterate this motivation in the discussion and conclusions, to help readers from other communities who may be more focused on benchmark performance as the primary metric of success. In addition, we agree with the reviewer that the main text should be self-consistent in terms of notation, and we'll make sure to do that in the next revision\n\n**W3** Section 2 of our manuscript contains a review of group equivariant architectures relevant to the scope of our paper. Namely, it discusses such models as [41], [42], [43], [17], and [40].\nReference R1 is less relevant, as it focuses on 2D symmetries whereas our focus is on the continuous 3D rotations. Still, we could certainly add a reference to this foundational work.\n\nThe PET model does not belong to this class of models, and therefore requires application of our symmetrization protocol to become exactly equivariant. There would be no need to apply symmetrization to an architecture that is intrinsically equivariant, and if one did so there would be absolutely no change in the predictions. However, our experiments on PET show that, contrary to the common wisdom in this application domain, models with built-in equivariance are not necessarily more accurate than models that are not. Further performance improvements might be possible in this extended design space: thanks to our exact symmetrization protocol, these improved models would be readily applicable to atomistic simulations. \n\n**Q1** Our symmetrization scheme needs some reference point in order to add rotational equivariance on top of any base model. In principle, this reference point might be arbitrary, for instance, the center of mass of the whole system. The use of all atomic positions as multiple reference points is not a restriction, but a choice that is associated with the nature of target properties in atomistic machine learning. We focus on this case because very often, and in particular for all the models we describe in section 2, predictions are given by an atom-centered decomposition of the target, but it is not a hard requirement. \n\n**Q2** Yes, the training is done in a non-equivariant way with rotational augmentations. We would say that the reason for the improvement is that the base model is not rotationally equivariant. For a non-equivariant model, one can apply a rotational augmentation strategy also during testing. This can be done by generating N random rotations, running the model for each of them, and next averaging the predictions. The larger the N, the better the accuracy. The expectation of the error monotonically decreases with N to some limit but never reaches it for any finite N. \nThe functional form of our symmetrization protocol is a weighted average of the model's predictions for different coordinate systems. Thus, it can also be seen as some form of rotational augmentations strategy during testing. Crucially, it uses a finite number of reference systems, that are chosen so that it achieves exact equivariance with a finite (and small) number of coordinate systems. However, it might not reach the asymptotic accuracy discussed above: in this case, if one artificially increases the number of coordinate systems, as described in F.8, the error should (slighlty) decrease. \n\n**Q3** No, we have always used a weighted sum of MSE losses in energies and forces. However, we used a checkpoint with the best MAE on a validation dataset for a final prediction, if MAE metrics was used. \n\n**Q4** Yes, $\\tilde{x}_{ij}^{k}$ is a hidden representation, also a message, associated with an edge between atoms i and j after the layer k. We thank the reviewer for the possibility to improve our manuscript by pointing us to this lack of definition in the main text. In revising our manuscript, this definition will be moved to the main text from the Appendix. \n\n**Q5** Features associated with each neighbor are given by encoding:\n1) A 3D displacement vector from the central atom to the neighbor\n2) Embedding of the atomic species of the neighbor atom\n3) For each layer but the first, an input message\n\nMore details can be found in Appendix A. We use optimizable embedding, where an abstract vector of a dimensionality $d_{PET}$ is assigned to each atomic specie. $d_{PET}$, which controls the dimensionality of virtually all hidden representations of our model, was set to 128 for most of our computational experiments and to 256 for a couple of them. The number of parameters in a single embedding layer is $d_{PET}$ multiplied by the number of unique atomic species in a dataset. Such an embedding is exactly equivalent to a linear layer applied on top of one-hot embedding.\n\n**References:**   \n\n[R1] Cohen, T., & Welling, M. \"Group equivariant convolutional networks.\"\n\nAll the other references coincide with the ones in the main text."}, {"review_id": "fZ1mj72yHf", "reviewer": "Reviewer_QskA", "summary": "This paper presents a novel posteriori rotational equivariant protocol that can be applied to any existing 3D models, irrespective of their initial equivariance to rotations. This protocol ensures not only rotational equivariance but also the preservation of other important model properties, such as permutation equivariance and expressiveness. Additionally, the paper introduces the Point Edge Transformer architecture, which achieves state-of-the-art performance on several benchmarks. The general equivariant design results in a modest performance improvement, highlighting the fact that exact equivariance is not always essential in certain learning tasks.", "questions": "I'm not familiar with the topic of atomic-scale applications. Therefore I will ask more questions about the equivariance part. \n\n1. As stated in the limitation, I think the equivariant design might be sensitive the noise in the sample, please correct me if I'm wrong. Could the authors provide some experiments to validate its robustness?\n\n2. Could the model be applied in large point cloud application such as shape segmentation or scene prediction in computer vision? I'm curious the scalabilty of the model and more applications of the model.\n\n3. Since the authors show that in some tasks the equivariance is not the essential part. Could the authors compare their equivariant model with the other equivariant models in some tasks where the equivariance ideed is important, for example, missing atom prediction in tensor field networks?\n", "limitations": "The authors adequately addressed the limitations. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper alleviates the need to incorporate the rotational equivariance into the model but designs a general method on the top of the trained $3D$ model.  The method is proven to be strictly equivariant to 3D rotations, allowing for the utilization of different non-equivariant 3D models and enabling broad applications across various 3D tasks.\n\n2. Through experimental analysis, this paper demonstrates that exact equivariance of the model is not always necessary for certain atomic applications. They show that by employing a powerful model and rotation augmentation techniques, non-intrinsically equivariant models outperform their counterparts. This finding provides valuable insights into the essential factors to consider when designing models for atomic applications.\n\n3. This paper is well-written, and the experimental results effectively support the authors' argument.", "weaknesses": "1. The proposed equivariant protocol achieves equivariance but comes with increased computational requirements, particularly in scenarios with large sample sizes. Moreover, exact rotation equivariance may be affected by the presence of noise in the data, thus introducing potential sensitivity issues.\n\n2. The proposed model may not applied to the 3D models like pointnet and pointnet++, which doesn't include the interaction of information between the points.\n\n3. As mentioned in the paper, for GNN network, their method either requires more computation to go over the whole neighboring points, or  involves the redesign of the message-passing mechanism.\n\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1: Computational cost**   \n It is true that our method achieves equivariance at the cost of increased computational requirements. However, we made a significant effort to reduce this overhead. Because of the adaptive inner cutoff radius (described in appendix F.4.), the number of coordinate systems per each central atom doesn't depend on sample size, nor on the number of neighbors included in each local environment, and is typically between 1 and 10. A non-equivariant model is simpler, and can be usually implemented very efficiently, so the overall cost for an equivariant prediction can still be low.\n\nOur symmetrization scheme always achieves exact equivariance in a rigorous mathematical sense. Simultaneously it preserves the smoothness of the prediction. Thus, it is not affected by the noise in the data. See our response to related Q1, where we describe the requested numerical experiment to demonstrate this. \n\n**W2: Applicability to pointnet**   \nOur symmetrization scheme can be applied to pointnet. As for pointnet++, we, indeed, are not aware of how to make this architecture as a whole applicable to atomistic modeling. However, the only obstacle pertains to downsampling layers, not to the ones that do interaction of information between the points. Moreover, due to the mostly local nature of atomistic interactions, we expect that one can remove downsampling layers from the pointnet++ architecture without a noticeable drop in accuracy. After such a modification, our symmetrization protocol can be applied to pointnet++. More details can be found in Appendix G.2.\n\n**W3: Message passing architectures**   \nIt is correct that the overhead when applying our symmetrization protocol to message passing architectures is steeper. As we discuss, it is possible to avoid this by re-designing the message-passing mechanism, and we are convinced that this can be done without substantially affecting the performance of the model. We provide a discussion of several approaches to do this in the manuscript in Appendix F.10. We think however that it would be impossible to perform a thorough investigation of these ideas in a manuscript that already discusses the symmetrization algorithm and benchmarks of a novel architecture. \n\n**Q1:** Our symmetrization scheme produces predictions that are smooth and exactly equivariant in a rigorous mathematical sense. Thus, it is robust to noise. We have conducted two numerical experiments to verify each of the statements. \n\nThe first one shows that the symmetrized model is indeed exactly invariant. We have empirically measured the discrepancy between predictions for different rotations for single (torch.float32) and double (torch.float64) precisions. We find that for single precision evaluation, the rotational discrepancy is 1.1e-7, while for the torch.float64, the rotational discrepancy is 1.86e-15. The magnitude of the predictions themselves was designed to be of the order of 1e0. Such a dramatic difference, along with an overall minimal rotational discrepancy for torch.float64, strongly indicates that machine precision is the only source of difference in predictions for different rotations. \n\nThe second experiment verifies the smoothness of our construction. We add random noise to the positions of all the points in an input point cloud and measure how much predictions change. Figure 1 in the figure response (the pdf file attached to the global response) illustrates that when the L2 norm of the noise converges to zero, so does the difference in predictions. Thus, this experiment verifies the smoothness of the symmetrized model.\n\n**Q2:** Yes, the model can be applied for point cloud applications beyond atomistic systems. The computational cost of the model scales linearly with the input size. Thus, it is, in principle, applicable to larger point clouds. We are also interested in the performance of the model for tasks in computer vision. Therefore this likely might be a direction of future research. \n\n**Q3.** Our experiments show that a non-equivariant model can achieve very good accuracy if trained with rotational augmentations, i.e. that *approximate* equivariance is enough to improve accuracy. For the suggested task of missing atom prediction, the target, i.e., the location of a missing atom, is not invariant. In other words, it is moving together with an input object. Thus, the model also must produce a non-invariant prediction that transforms in a certain manner with shifts and rotations of an input point cloud. \n\nOur benchmarks section already contains an experiment with similar settings, i.e. dipole moment prediction on the QM9 dataset. A dipole moment is a vector rigidly attached to a molecule. This vector rotates together with an input molecule. Thus, the model also must predict different outputs for different orientations. We achieve state-of-the-art on this benchmark, outperforming all the other models, as it is illustrated in Fig. 3 f. This is achieved by modifying a rotational augmentation strategy during the training. In addition to the random rotation of an input point cloud, the same rotation should be applied to a target property. We also note that our symmetrization scheme allows for ensuring exact vectorial equivariance during inference. This means that after applying our symmetrization protocol, the prediction of the overall model transforms exactly as a vector with rotations of an input point cloud. We believe that this numerical experiment is sufficient to demonstrate the performance of our model for a task where equivariance is crucial.\n"}, {"review_id": "EcPXOlLVA2", "reviewer": "Reviewer_oCKF", "summary": "The authors propose a method for uniting rotational symmetries with translational and permutation symmetries in neural networks for point cloud data. Hereby, the authors particularly focus on molecule data. This is achieved by adding rotational augmentations at training time and evaluating the resulting NN on an ensemble of coordinate systems at test time.", "questions": "see weaknesses", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The incorporation of rotational symmetries into NNs applied to point data seems very important to the application of DL to atomic data as it may reduce the data hunger of NNs and improve generalization.\n- The  approach using an ensemble of coordinate systems seems innovative and makes sense.\n", "weaknesses": "- What is the typical number of coordinate systems in eq 3? Since each coordinate system is defined by three atoms, I guess it becomes very large very quickly.\n- I wonder whether using y_s is even necessary. It seems to greatly increase the runtime at inference while not really adding anything in terms of performance. I.e. Simple training with rotational augmentation seems to yield all of the performance gains. Could the authors elaborate on this? I think it would be necessary to clearly demonstrate the usefulness of this addition\n- Are competing methods also using rotational augmentations?\n- -Do the authors have any idea why the SOTA in tab 1 is so much worse on MnO & HEA?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1: Number of coordinate systems**  \nThe reviewer is correct that a naïve implementation of our core idea would involve summing over a very large number of coordinate systems. Our symmetrization scheme, however, doesn't suffer from this problem. The reason is that not all triplets of atoms are used to construct coordinate systems. Instead, we select only a few closest neighbors to the central atom. This is achieved by defining in a smooth way an adaptive cutoff radius that is as small as possible but sufficient to define at least one \"good\" (where the corresponding triplet of atoms is sufficiently non-collinear) coordinate system. \n\nThe overall pipeline of our symmetrization protocol is quite complicated. Thus, for better readability, we decided to first explain (in Sec. 5, \"Ensemble of coordinate systems\") the simplest possible idea, which, as the reviewer notes, suffers from high computational cost and the presence of numerically unstable corner cases. Next, we provide several optimizations and solutions to all corner cases. The main ideas are sketched out in the paragraph \"Adaptive cutoff and limiting cases\" in sec. 5 of the main text. Details of the implementation are given in the Appendix F.4, and potential trade-offs between computational cost and smoothness are discussed in Appendix F.6. Empirically, for the parameters we used in our benchmarks our algorithm selects between 1 and 10 coordinate systems per environment. \n\n**W2: Need of symmetrization**   \nWe clarify the purpose of our symmetrization scheme in our global response. In short, the primary goal of our symmetrization protocol is not to further increase the accuracy. Instead, it is required for the very applicability of a model to practical simulations in atomistic modeling, a domain which considers exact compliance with physical symmetries as a basic requirement. The notion that models should be *intrinsically* equivariant is a constraint that has directed, and perhaps slowed down, the progress of ML applications to this domain. To sum up, given this domain requirement, without our symmetrization scheme, it is unlikely that PET-like models would be used in practical applications, despite their excellent accuracy. With our symmetrization scheme, PET becomes exactly invariant with respect to rotations and, thus, is applicable for atomistic modeling. \n\nIn our manuscript, we described this in the introduction, but we acknowledge that our description might not be sufficiently clear, especially for readers from a different domain. We will do our best to provide a more compelling explanation in a revision of our manuscript.\n\n**W3: Rotational augmentation in competing methods**  \nCompeting methods are exactly, and intrinsically, rotationally invariant. Thus, applying rotational augmentations to them would change nothing at all. For instance, one of the methods [39] relies on distances between the atoms, angles, and dihedral angles. If an atomic configuration is rotated, the distances between the atoms and internal angles stay precisely the same, and so it is irrelevant whether rotational augmentation is used or not. To sum up, competing methods cannot be improved by using rotational augmentations (and in fact they are usually advertised for not requiring augmentation); thus, the comparison is fair. \n\n**W4: SOTA on MnO and HEA**   \nThe case of MnO corresponds to magnetic potentials, which is a more complicated case compared to the usual task in atomistic machine learning; see more details in [92]. Only few models currently support this setup. For instance, Dimenet, Gemnet, and MACE have not yet been reported to be extended to magnetic case and benchmarked on a magnetic dataset such as MnO. We acknowledge that if extended to a magnetic setup, these, or maybe some other, models might perform noticeably better than the current SOTA[92] on the MnO dataset. But this hasn't happened at the time of the submission, and moreover, nor until now, to the best of our knowledge. This aspect highlights the strength of PET, that it can be extended to nonstandard setups with very little effort. In addition, the current SOTA[92] on MnO is also a recent model (from 2021) that is used successfully in practical application, and outperforming it almost 3 times highlights the overall excellent performance of PET. To sum up, we regard the superior performance of PET on the MnO dataset as a strength, not a weakness. \n\nThe HEA dataset has been very recently introduced, and is characterized by extreme levels of chemical diversity, with up to 25 chemical elements simultaneously considered in a single structure. The architecture proposed in the original paper relies heavily on physical priors to improve robustness and transferability. The Reviewer will notice how, in the discussion of the results, we emphasise how the dramatic increase in validation set performance is accompanied with a deterioration of accuracy in strongly extrapolative tasks. We believe this to be an interesting observation to guide future improvements of the PET architecture. \nBesides these examples, that serve as much as case studies for specific dataset features as much as empirical benchmarks, we also report a comparison with the COLL and HME21 datasets, to which Dimenet, Gemnet, and MACE have been also applied. The difference in performance between these models and PET is not as large (also reported in Table 1), but also in this case our model yields remarkable performance, particularly for the COLL dataset that was originally proposed by the authors of Dimenet and Gemnet.\n\n**References**    \nReference numbers match the ones in our manuscript. \n"}]

Author rebuttal:
We thank all reviewers for providing their comments, which highlight opportunities to improve the presentation of our work. We would like to start by addressing a misunderstanding that appears to recur in several places, concerning the relation between the two main contributions in our manuscript - the *a posteriori* symmetrization and the PET model. 

**Exact equivariance is de facto a requirement for atomistic modeling**   
As we should probably have made clearer, and as we will certainly clarify in later revisions, the main reason why several application domains, such as atomic-scale simulations, require exact equivariance is to preserve fundamental physical constraints, lack of which can lead to misleading end results (e.g. spontaneous ordering of molecules in an isotropic fluid). For this reason, equivariance that is exact in a rigorous mathematical sense is highly sought after. An approximate one, achievable by supplementing any model with rotational augmentations, would be considered insufficient by most practitioners. This is the primary reason why, at variance with other fields, the vast majority of ML models for atomistic modeling are built around intrinsically equivariant architectures. 

**Non-equivariant models can outperform intrinsically equivariant ones**  
One could argue that, given that such equivariant architectures exist, our symmetrization protocol solves a non-existent problem. This is why we devote a substantial part of our manuscript to describe and benchmark a non-equivariant model on atomistic-simulation datasets. The excellent performance of PET demonstrates that it can be beneficial to relax the requirement of intrinsic equivariance when designing ML models for atomistic simulation tasks, an observation that we suspect may come as surprising to many domain experts. 

**There is a wealth of models and ideas that can now be applied to new domains**   
There are dozens of algorithms for generic point clouds (some of which we summarize in Sec. 3, e.g. PointNet, PointNet++, O-cnn, Octnet, kd-network, Kpconv, PointMLP, PAConv, Point Transformer, PVT) that have never been applied to atomistic simulations because of lack of equivariance: this untapped potential for models, or architectural features, that can now be seamlessly applied to a new domain, is in our opinion a major outcome of our study.  
The small performance improvement that comes from symmetrizing *a posteriori* the PET model is a fortunate side effect of a procedure whose primary goal is to make PET, and other models, applicable to domains of science that require exact compliance with fundamental physical requirements. 

**Choice of benchmarks, and insights on architectural features**   
Our model, PET, achieves state of the art performance on several benchmarks. For instance, it outperforms such models as DimeNet++ and GemNet on the COLL dataset, developed by the authors of the mentioned models. Another general consideration, however, and we would like to better highlight it in the manuscript, is that we picked most of our benchmarks not only as a way to demonstrate a favourable comparison with the state of the art, but also as a way to highlight some of the strengths and/or shortcomings of our model. For instance, MnO and the dipole models allow us to highlight the simplicity by which non-standard setups (namely atomic attributes associated with magnetism, and the vectorial nature of the target) can be incorporated in the model. The water dataset is well-suited for ablation studies on the receptive field of the model. The HEA dataset allows us to highlight the need to improve the extrapolative capabilities of PET. We believe that this choice adds an interesting dimension to our analysis, while some reviewers seem to have focused their comments primarily on the bare benchmark numbers. 

**Summary of contributions**  
To summarize, the impact of our work is best understood in the context of applications to chemical and materials simulations, for which exact equivariance and smoothness are considered a hard requirement, that have guided and somewhat restricted the development of geometric machine-learning schemes to frameworks that are intrinsically equivariant. In this context our manuscript contributes:
* A practical algorithm to make a generic point-cloud model exactly, rigorously equivariant: this makes multiple existing algorithms applicable to a new domain.
* A demonstration that relaxing the architectural constraint of equivariance can be beneficial in terms of model accuracy. This will come as a surprise to many practitioners, as “built-in” equivariance is usually expected to be also beneficial in terms of accuracy
* A novel architecture that is comparable or better than several state-of-the-art approaches, and that will likely find its own practical applications


Meta-review: After the rebuttal discussion, all reviewers agreed to accept this work (one did so in a comment but did not update their review).  All recognize the importance of the problem tackled in the paper, although the amount of space in the paper for equivariance relative to PETS rightfully seems imbalanced given the focus of the work.

For the camera-ready version of the paper, please add the assorted changes and experiments introduced during the rebuttal phase.

(Also, the AC wonders why the "Vector Neurons" architecture and follow-on work isn't referenced here?)
Decision: Accept (poster)