Title: Smooth, exact rotational symmetrization for deep learning on point clouds
Venue: NeurIPS 2023
Accepted: yes
Keywords: geometric deep learning

point clouds

equivariance

machine learning potentials

GNN

transformer

atomic-scale modeling
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

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.

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, as well as in chemistry and materials modeling, and are the input of a variety of dedicated deep-learning techniques[1]. 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[2-4]. In the case of atomistic modeling, 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[5, 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 high-throughput calculations[8, 9]) to qualitatively incorrect and even absurd[10, 11] simulation outcomes.

These concerns have led to the development of dedicated models for atomistic simulations that rigorously incorporate all these constraints [12-15] by using only symmetry preserving operations. These restrictions limit the design space of these models, for example, leading to the lack of universalapproximation 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++[16], 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 is 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 generic 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, it 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-art performance on several benchmarks ranging from high-energy CH 4 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 a 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 mechanisms. The first involves using only invariant internal coordinates such as interatomic distances, angles, or even dihedral angles[17]. The second involves using equivariant hidden representations that transform in a 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[18]. 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 i , 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 = ∑ i y ( A i ) . Here, y ( A i ) indicates a learnable function of the environment of the i -th atom, defined as all the neighbors within the cutoff radius R c. This local decomposition is 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 y ( A i ) is to first compute smooth, invariant descriptors for each environment[14, 22-27] and then feed them to a 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.[29].

Distance-based message-passing. Another popular method involves constructing a molecular graph and feeding it to a Graph Neural Network (GNN). Early, and very popular, models rely on invariant two-body messages based only on interatomic distances[30, 31]. In this case, a 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

could decorate the edges of a molecular graph with x, y, and z components of the corresponding displacement vectors. This would immediately ensure the universal approximality of the model. However, the necessity to enforce rotational equivariance 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.

Higher-order message-passing. Subsequently, more expressive models were developed. For example, some of them use angles[33], or even dihedrals[17], in addition to distances between the atoms. Others, such as Tensor Field Network[34], Nequip[35], MACE[36], SE(3) transformers[37], 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 reduce to the simple presence of a 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 as 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 a structure can be encoded in a permutation-invariant manner by projecting the point cloud onto a regular 3D voxel grid, and further manipulated by applying three-dimensional convolutions. The computational cost of a 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[45, 46], one can evaluate an expression such as

where, f i n is the n -th feature of the point i , and g is a collection of N in N out learnable threedimensional functions, where N in and N out are the numbers of input and output features respectively. The output features ( F ∗ g )( r ) m can be evaluated at an arbitrary point r , allowing for complete freedom in the construction of the grid. PointNet[47] is a paradigmatic example of an architecture that eliminates the distinction between point features and positions. The core idea of these approaches is that an expression such as f ( r 1 , r 2 , ..., r n ) = γ (max i ( { h ( r i ) } )) , where γ and h are learnable functions, can approximate to arbitrary precision any continuous function of the point set { r i } . PointNet++[16] and many other models[46, 48, 49] apply similar functional form in a hierarchical manner to extract features of local neighborhoods. Many of these methods have graph-convolution or message-passing forms, similar to those discussed in Section 2, even though they do not enforce invariance or equivariance of the representation. Many transformer models for point clouds have also been proposed[50-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 to a 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

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. (1), that introduces discontinuities related to (dis)appearance of new points at the cutoff sphere, even for smooth functions g mn . These discontinuities can be removed by modifying the convolutional operator as:

where the cutoff function f c, illustrated in Fig. 5c, can be chosen as an infinitely differentiable switching function that smoothly zeroes out contributions from the points approaching the cutoff sphere, and the parameter ∆ R c > 0 controls how fast it converges to 1 for r < R c. 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 a 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. 3 differentiable. Finally, most - but not all[59] - of the generic models for point clouds are invariant to permutations.

In summary, there is 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 a fully-symmetric behavior, such as those that are common in atomistic machine learning, is invariance or 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 a 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 it is one of the nodes in the point cloud. More specifically, we formulate our method in the context of a local decomposition of the (possibly tensorial) target y ( A ) = ∑ i y ( A i ) where the reference point is given by the position of the central atom r i , and the model estimates an atomic contribution y ( A i ) given the local neighborhood. We name our approach Equivariant Coordinate System Ensemble (ECSE, pron. e S e).

Local coordinate systems. A simple idea to obtain a rotationally equivariant model would be to define a coordinate system rigidly attached to an atomic environment. 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 a 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 a close relative of Behler-Parrinello symmetry functions[22], which guarantee smoothness and invariance.

Ensemble of coordinate systems. The basic formulation of our symmetrization protocol, illustrated in Fig. 5a, 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:

a

yo(R[Ai])

Rjyo(F

atomic

environment

Ai

e1

architecture

Backbone

R

jj'

Backbone

e2

architecture

R

e3

e1

architecture

Backbone

R

jj'

ys(Ai)

Rout

Rin

where y S indicates the symmetrized model, y 0 the initial (non-equivariant) backbone architecture, and ˆ R jj ′ [ · ] indicates the rotation operator for the coordinate system defined by the neighbors j and j ′ within the i -centered environment A i . In other words, ˆ R -1 jj ′ [ A i ] 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 c .

Given a pair of neighbors with displacement vectors from central atom r ij and r ij ′ the corresponding coordinate system consists of the vectors ˆ r ij , ˆ r ij × ˆ r ij ′ , and ˆ r ij × [ˆ r ij × ˆ r ij ′ ] . If the model predicts a vectorial (or tensorial) output, it is rotated back into the original coordinate system by an outer application of operator ˆ R jj ′ . Thus, our symmetrization scheme allows for getting not only invariant predictions, but also covariant ones, such as 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 polynomially-computable invariant metric to measure the similarity between crystal structres[64-66], and to construct provably complete invariant density-correlation descriptors [67].

Fig. 5b depicts the necessity to use the weighted average with weights w jj ′ instead of a plain one. The use of plain average in Eq. 3 would lead to the lack of smoothness. Indeed, if a new atom enters the cutoff sphere, it immediately yields new terms into the summation in eq. 3, 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 is assigned a different weight, depending on the positions ( r ij , r ij ′ ) of the neighbors that define the reference frame:

where the cutoff functons f c and q c are illustrated in Fig. 5c. Thanks to the presence of f c, the terms in Eq. (3) that are associated with atoms entering or exiting the cutoff sphere have zero weights. Similarly, q c 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 are

(nearly) collinear, which would make Eq. (3) ill-defined, falling to 0 0 ambiguity. We discuss in the Appendix F.3 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 R in ( A i ) , which is determined separately for each atomic environment A i instead of being a global, user-specified constant. Eq. (3) requires that at least one pair of neighbors is inside the cutoff sphere. Simultaneously, it is desirable to make it as small as possible for computational efficiency. Thus, it makes sense to define R in ( A i ) along the lines of the distance from the central atom to the second closest neighbor. It should be larger, but not much larger, than the second-nearest-neighbor distance.

Our definition of R in ( A i ) , 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, R in ( A i ) 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 k nearest neighbors are collinear, R in ( A i ) is 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 ECSE-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 a 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 backbone architecture, playing the role of test augmentation.

Message-passing. Message-passing schemes can be regarded as local models with a cutoff radius given by the receptive field of the GNN, and therefore ECSE can be applied transparently to them. In a naïve implementation, however, the same message would have to be computed for the coordinate systems of all atoms within the receptive field. This implies a 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 a deeper integration within the model, that also changes the nature of the message-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 equivariance 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, a completely different approach is introduced by the vector neurons framework[70] that converts 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 pointcloud architectures discussed in Section 3, making them directly-applicable to atomistic modeling and to any other domain with strict symmetry requirements. Here we want instead to focus on a less-

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 a new architecture that we name Point Edge Transformer (PET), which is 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[72] that loosely resembles a single message-passing block of PET) as well as with early attempts by Behler et al. applying MLPs to invariant edge features[73].

The PET architecture is illustrated in Fig. 2. More details are given in Appendix A. The core component of each message-passing block is a permutationally-equivariant transformer that takes a set of tokens of size d PET associated with the central atom and all its neighbors, and generates new tokens that are used both to make predictions and as output messages. The transformer we use in PET is a straightforward implementation of the classical one[74], with a modification to the attention mechanism that ensures smoothness with respect to (dis)appearance of the neighbors at the

cutoff radius. The attention coefficients α ij determining the contribution from token j to token i are modified as α ij ← α ij f c ( r j | R c , ∆ R c ) , 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-free 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-increase 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[72].

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-passing mechanism. For the current moment, our proof-of-principle implementation follows a simpler approach (see more details in Appendix F.11) and favors ease of implementation instead of computational efficiency (e.g., all weights and most of the intermediate values used in the ECSE scheme are stored in separate zero-dimensional PyTorch[77] tensors). This leads to a 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 equivariant PET inference is much cheaper than the reference first-principles calculations.

7 Benchmarks

We benchmark PET and the ECSE scheme over six different datasets, which have been previously used in the literature and which allow us to showcase the performance of our framework, and the ease with which it can be adapted to different use cases. The main results, are compared with the state of the art in Figure 3 and Table 1, while in-depth analyses can be found in the Appendix C.

As a first benchmark, we conduct several experiments with the liquid-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 molecular 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 a single layer performs considerably worse than one with two or more layers. Given that stacking multiple GNN blocks increases both the receptive field and the flexibility in describing local interactions, it is interesting to look at the trend for a fixed total number of transformer layers (Fig. 3b), that shows that a 6 × 2 model outperforms a large 12 × 1 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 error over the state-of-the-art equivariant model NEQUIP[35] 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 R c beyond 3.7Å (a value we chose to ensure approximately 20 neighbors on average), reaching a force MAE of 14.4 meV/Å when using a cutoff of 4.25Å.

We then move to the realm of small molecules with the COLL dataset[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 error on forces (the main optimization target in previous studies) while reducing by a 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 a very wide energy range, we consider the database of CH 4

1

2

3

4

n

GNN

2×10 1

3×10 1

4×10 1

MAE

f

/ meV/Å

water

a

n

TL=1

n

TL=2

n

NEQUIP

1

2

3

4

6

12

n

GNN

water

b

n

GNN ×

n

TL=12

NEQUIP

3.0

3.5

4.0

4.5

R

c

/ Å

water

c

n

GNN,

n

TL =6,2

NEQUIP

10 2

10 4

10 6

n

train

10 0

10 1

10 2

10 3

RMSE

E

/ meV

CH4

E

d

WK[81]

SOAP[23]

NICE[78]

NN[80]

LE-ACE[79]

PET

dPET

=128

10 2

10 4

10 6

n

train

CH4

E

,

f

e

REANN[84]

PET

dPET

=128

PET

dPET

=256

10 1

10 2

10 3

10 4

n

train

10

1

10 0

MAE |

| / D

QM9 dipoles

f

FCHL*[83]

GPR[82]

WK[81]

PET

dPET

=128

configurations first introduced in Ref. 80. This dataset contain random arrangements of one C and 4 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 in Figure 3c and d demonstrate the steep, monotonic improvement of performance for increasing train set size, surpassing after a few thousand training points the best existing models (a linear many-body potential for the energy-only training[81] and REANN, an angle-based GNN[84] for combined energy/gradient training).

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 [90]. 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 of compression. However, the more flexible form of the model compared to HEA25-4-NN allows for a a 3(5)-fold 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 meV/at.

MAE at 5000 K) even though room-temperature MD trajectories are better described by PET. The case of the HME21 dataset, which contains high-temperature molecular-dynamics 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[36], inheriting much of the robustness of shallow models. PET, however, comes very close: the simple regularizing effect of a 5-models PET ensemble is sufficient to tip the balance, bringing the force error to 128.5 meV/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[91]. PET achieves a MAE of 6.7 meV/molecule, in line with the accuracy of DimeNet++[87], but not as good as the 4.3 meV/molecule MAE achieved by Wigner kernels[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 MnO 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 in the description of the atomic environments. The improvement with respect to the original model is quite dramatic (Tab. 1). 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 a recently-developed many-body kernel regression scheme[81] 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, we also develop the PET model, a deep-learning architecture that is not intrinsically rotationally invariant but achieves state-of-the-art 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 a 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-ofthe-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] while relaxing the exact equivariance of the model.

It 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 an example, one can examine the models proposed for, e.g., the ModelNet40 dataset[95], 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-inspired inductive biases (such as smoothness, range, and body order of interactions) that have this far been conflated only with a specific class of equivariant architectures.

References

Yulan Guo, Hanyun Wang, Qingyong Hu, Hao Liu, Li Liu, and Mohammed Bennamoun. Deep learning for 3d point clouds: A survey. IEEE transactions on pattern analysis and machine intelligence , 43(12):4338-4364, 2020.

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.

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.

George Em Karniadakis, Ioannis G. Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nat Rev Phys , 3(6):422-440, May 2021.

Jirasak Wong-ekkabut and Mikko Karttunen. The good, the bad and the user in soft matter simulations. Biochimica et Biophysica Acta (BBA) - Biomembranes , 1858(10):2529-2538, October 2016.

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 , 119(41):23752-23766, October 2015.

Michael F Herbst and Antoine Levitt. Black-box inhomogeneous preconditioning for selfconsistent field iterations in density functional theory. J. Phys.: Condens. Matter , 33(8):085503, February 2021.

Wahyu Setyawan and Stefano Curtarolo. High-throughput electronic band structure calculations: Challenges and tools. Computational Materials Science , 49(2):299-312, August 2010.

Valerio Vitale, Giovanni Pizzi, Antimo Marrazzo, Jonathan R. Yates, Nicola Marzari, and Arash A. Mostofi. Automated high-throughput Wannierisation. npj Comput Mater , 6(1):66, December 2020.

Xiaojing Gong, Jingyuan Li, Hangjun Lu, Rongzheng Wan, Jichen Li, Jun Hu, and Haiping Fang. A charge-driven molecular water pump. Nature Nanotech , 2(11):709-712, November 2007.

Jirasak Wong-ekkabut, Markus S. Miettinen, Cristiano Dias, and Mikko Karttunen. Static charges cannot drive a continuous flow of water molecules through a carbon nanotube. Nature Nanotech , 5(8):555-557, August 2010.

Jörg Behler and Michele Parrinello. Generalized Neural-Network Representation of HighDimensional Potential-Energy Surfaces. Phys. Rev. Lett. , 98(14):146401, April 2007.

Albert P. Bartók, Mike C. Payne, Risi Kondor, and Gábor Csányi. Gaussian Approximation Potentials: The Accuracy of Quantum Mechanics, without the Electrons. Phys. Rev. Lett. , 104(13):136403, April 2010.

O. Anatole von Lilienfeld, Raghunathan Ramakrishnan, Matthias Rupp, and Aaron Knoll. Fourier series of atomic radial distribution functions: A molecular fingerprint for machine learning models of quantum chemical properties. Int. J. Quantum Chem. , 115(16):1084-1093, August 2015.

Stefan Klus, Patrick Gelß, Feliks Nüske, and Frank Noé. Symmetric and antisymmetric kernels for machine learning problems in quantum physics and chemistry. Mach. Learn.: Sci. Technol. , 2(4):045016, December 2021.

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.

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.

Kristof Schütt, 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.

Jérôme Cuny, Yu Xie, Chris J. Pickard, and Ali A. Hassanali. Ab Initio Quality NMR Parameters in Solid-State Materials Using a High-Dimensional Neural-Network Representation. J. Chem. Theory Comput. , 12(2):765-773, February 2016.

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):4501, December 2018.

E Prodan and W Kohn. Nearsightedness of electronic matter. Proc. Natl. Acad. Sci. , 102(33):11635-11638, August 2005.

Jörg Behler. Atom-centered symmetry functions for constructing high-dimensional neural network potentials. The Journal of Chemical Physics , 134(7):074106, February 2011.

Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Phys. Rev. B , 87(18):184115, May 2013.

Mitchell A. Wood and Aidan P. Thompson. Extending the accuracy of the SNAP interatomic potential form. The Journal of Chemical Physics , 148(24):241721, June 2018.

Haoyan Huo and Matthias Rupp. Unified representation for machine learning of molecules and crystals. ArXiv Prepr. ArXiv170406439 , 13754, 2017.

Alexander V. Shapeev. Moment Tensor Potentials: A Class of Systematically Improvable Interatomic Potentials. Multiscale Model. Simul. , 14(3):1153-1173, January 2016.

Ralf Drautz. Atomic cluster expansion for accurate and transferable interatomic potentials. Phys. Rev. B , 99(1):014104, January 2019.

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):10073-10141, August 2021.

Jörg Behler. Four Generations of High-Dimensional Neural Network Potentials. Chem. Rev. , 121(16):10037-10072, August 2021.

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.

Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. In NIPS , 2017.

Sergey N Pozdnyakov and Michele Ceriotti. Incompleteness of graph neural networks for points clouds in three dimensions. Mach. Learn.: Sci. Technol. , 3(4):045020, December 2022.

Yaolong Zhang, Junfan Xia, and Bin Jiang. Reann: A pytorch-based end-to-end multifunctional deep neural network package for molecular, reactive, and periodic systems. The Journal of Chemical Physics , 156(11), 2022.

Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds. arxiv:1802.08219 , 2018.

Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. 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.

Ilyes 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)-equivariant atom-centered interatomic potentials. arxiv:2205.06643 , 2022.

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.

Hang Su, Subhransu Maji, Evangelos Kalogerakis, and Erik Learned-Miller. Multi-view convolutional neural networks for 3d shape recognition. In Proceedings of the IEEE international conference on computer vision , pages 945-953, 2015.

Seyed Saber Mohammadi, Yiming Wang, and Alessio Del Bue. Pointview-gcn: 3d shape classification with multi-view point clouds. In 2021 IEEE International Conference on Image Processing (ICIP) , pages 3103-3107. IEEE, 2021.

Xiaozhi Chen, Huimin Ma, Ji Wan, Bo Li, and Tian Xia. Multi-view 3d object detection network for autonomous driving. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition , pages 1907-1915, 2017.

Asako Kanezaki, Yasuyuki Matsushita, and Yoshifumi Nishida. Rotationnet: Joint object categorization and pose estimation using multiviews from unsupervised viewpoints. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 5010-5019, 2018.

Peng-Shuai Wang, Yang Liu, Yu-Xiao Guo, Chun-Yu Sun, and Xin Tong. O-cnn: Octree-based convolutional neural networks for 3d shape analysis. ACM Transactions On Graphics (TOG) , 36(4):1-11, 2017.

Gernot Riegler, Ali Osman Ulusoy, and Andreas Geiger. Octnet: Learning deep 3d representations at high resolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 3577-3586, 2017.

Roman Klokov and Victor Lempitsky. Escape from cells: Deep kd-networks for the recognition of 3d point cloud models. In Proceedings of the IEEE international conference on computer vision , pages 863-872, 2017.

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 computer vision and pattern recognition , pages 2589-2597, 2018.

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 IEEE/CVF international conference on computer vision , pages 6411-6420, 2019.

Charles R Qi, Hao Su, Kaichun Mo, and Leonidas J 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.

Xu Ma, Can Qin, Haoxuan You, Haoxi Ran, and Yun Fu. Rethinking network design and local geometry in point cloud: A simple residual mlp framework. arXiv preprint arXiv:2202.07123 , 2022.

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.

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 1625916268, 2021.

Cheng Zhang, Haocheng Wan, Xinyi Shen, and Zizhao Wu. Pvt: Point-voxel transformer for point cloud learning. International Journal of Intelligent Systems , 37(12):11985-12008, 2022.

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 IEEE/CVF International Conference on Computer Vision , pages 3164-3173, 2021.

Meng-Hao Guo, Jun-Xiong Cai, Zheng-Ning Liu, Tai-Jiang Mu, Ralph R Martin, and Shi-Min Hu. Pct: Point cloud transformer. Computational Visual Media , 7:187-199, 2021.

Abien Fred Agarap. Deep learning using rectified linear units (relu). arXiv preprint arXiv:1803.08375 , 2018.

Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:1511.07289 , 2015.

Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415 , 2016.

Diganta Misra. Mish: A self regularized non-monotonic activation function. arXiv preprint arXiv:1908.08681 , 2019.

Jonathan T Barron. Continuously differentiable exponential linear units. arXiv preprint arXiv:1704.07483 , 2017.

Yangyan Li, Rui Bu, Mingchao Sun, Wei Wu, Xinhan Di, and Baoquan Chen. Pointcnn: Convolution on x-transformed points. arXiv preprint arXiv:1801.07791 , 2018.

Tristan Bereau, Denis Andrienko, and O. Anatole Von Lilienfeld. Transferable Atomic Multipole Machine Learning Models for Small Organic Molecules. J. Chem. Theory Comput. , 11(7):3225-3233, July 2015.

Chungwen 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 , 96(4):041407, July 2017.

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):143001, April 2018.

Han Wang, Linfeng Zhang, Jiequn Han, and Weinan E. DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics. Computer Physics Communications , 228:178-184, July 2018.

Vitaliy Kurlin. Exactly computable and continuous metrics on isometry classes of finite and 1-periodic sequences. arXiv preprint arXiv:2205.04388 , 2022.

Vitaliy Kurlin. Computable complete invariants for finite clouds of unlabeled points under euclidean isometry. arXiv preprint arXiv:2207.08502 , 2022.

Daniel Widdowson and Vitaliy Kurlin. Resolving the data ambiguity for periodic crystals. Advances in Neural Information Processing Systems (Proceedings of NeurIPS 2022) , 35, 2022.

Jigyasa Nigam, Sergey N. Pozdnyakov, Kevin K. Huguenin-Dumittan, and Michele Ceriotti. Completeness of atomic structure representations. arxiv:2302.14770 , 2023.

Omri Puny, Matan Atzmon, Heli Ben-Hamu, Edward J Smith, Ishan Misra, Aditya Grover, and Yaron Lipman. Frame averaging for invariant and equivariant network design. arXiv preprint arXiv:2110.03336 , 2021.

Alexandre Agm Duval, Victor Schmidt, Alex Hernández-Garcıa, Santiago Miret, Fragkiskos D 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.

Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas J Guibas. Vector neurons: A general framework for so (3)-equivariant networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 12200-12209, 2021.

Leon Bergen, Timothy O'Donnell, and Dzmitry Bahdanau. Systematic generalization with edge transformers. Advances in Neural Information Processing Systems , 34:1390-1402, 2021.

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.

K. V.Jovan Jose, Nongnuch Artrith, and Jörg Behler. Construction of high-dimensional neural network potentials using environment-dependent atom pairs. J. Chem. Phys. , 136(19):194111, 2012.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Chaitanya Joshi. Transformers are graph neural networks. The Gradient , 12, 2020.

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077 , 2019.

Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.

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.

Filippo Bigi, Kevin K Huguenin-Dumittan, Michele Ceriotti, and David E Manolopoulos. A smooth basis for atomistic machine learning. The Journal of Chemical Physics , 157(23):234101, 2022.

Sergey N Pozdnyakov, Michael J Willatt, Albert P Bartók, Christoph Ortner, Gábor Csányi, and Michele Ceriotti. Incompleteness of Atomic Structure Representations. Phys. Rev. Lett. , 125:166001, 2020.

Filippo Bigi, Sergey N. Pozdnyakov, and Michele Ceriotti. Wigner kernels: body-ordered equivariant machine learning without a basis. arxiv:2303.04124 , 2023.

Max Veit, David M. Wilkins, Yang Yang, Robert A. DiStasio, and Michele Ceriotti. Predicting molecular dipole moments by combining atomic partial charges and atomic dipoles. J. Chem. Phys. , 153(2):024113, July 2020.

Felix A. Faber, Anders S. Christensen, Bing Huang, and O. Anatole Von Lilienfeld. Alchemical and structural distribution based representation for universal quantum machine learning. J. Chem. Phys. , 148(24):241717, June 2018.

Yaolong Zhang, Junfan Xia, and Bin Jiang. REANN: A PyTorch-based end-to-end multifunctional deep neural network package for molecular, reactive, and periodic systems. J. Chem. Phys. , 156(11):114801, 2022.

Bingqing Cheng, Edgar A. Engel, Jörg Behler, Christoph Dellago, and Michele Ceriotti. Ab initio thermodynamics of liquid and solid water. Proc. Natl. Acad. Sci. U. S. A. , 116(4):11101115, January 2019.

Johannes Gasteiger, Shankari Giri, Johannes T Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14115 , 2020.

Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional Message Passing for Molecular Graphs. 2020.

Marco Eckhoff and Jörg Behler. High-dimensional neural network potentials for magnetic systems using spin-dependent atom-centered symmetry functions. npj Computational Materials , 7(1):170, 2021.

James P Darby, Dávid P Kovács, Ilyes Batatia, Miguel A Caro, Gus LW Hart, Christoph Ortner, and Gábor Csányi. Tensor-reduced atomic density representations. arXiv preprint arXiv:2210.01705 , 2022.

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):045802, April 2023.

Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Sci. Data , 1:1-7, August 2014.

Larry Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood. Spherical channels for modeling atomic interactions. Advances in Neural Information Processing Systems , 35:8054-8067, 2022.

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.

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) dataset and community challenges. Acs Catalysis , 11(10):6059-6072, 2021.

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.

Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks , 107:3-11, 2018.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

Scipy generator of random rotations. https://docs.scipy.org/doc/scipy/reference/ generated/scipy.spatial.transform.Rotation.random.html .

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

So Takamoto, Chikashi Shinagawa, Daisuke Motoki, Kosuke Nakago, Wenwen Li, Iori Kurata, Taku Watanabe, Yoshihiro Yayama, Hiroki Iriguchi, Yusuke Asano, et al. Towards universal neural network potential for material discovery applicable to arbitrary combination of 45 elements. Nature Communications , 13(1):2991, 2022.

So Takamoto, Satoshi Izumi, and Ju Li. Teanet: Universal neural network interatomic potential inspired by iterative electronic relaxations. Computational Materials Science , 207:111280, 2022.

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.

Bingqing Cheng, Edgar Engel, Jörg Behler, Christoph Dellago, and Michele Ceriotti. Dataset: Ab initio thermodynamics of liquid and solid water, 2018.

Sergey Pozdnyakov, Michael Willatt, and Michele Ceriotti. Dataset: Randomly-displaced methane configurations, 2020.

Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems , 32, 2019.

Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning , pages 9323-9332. PMLR, 2021.

Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Veliˇ ckovi´ c, James Kirkpatrick, and Peter Battaglia. Simple gnn regularisation for 3d molecular property prediction & beyond. arXiv preprint arXiv:2106.07971 , 2021.

Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d graph networks. arXiv preprint arXiv:2102.05013 , 2021.

Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations , 2022.

Andrea Grisafi and Michele Ceriotti. Incorporating long-range physics in atomic-scale machine learning. J. Chem. Phys. , 151(20):204105, November 2019.

Tsz Wai Ko, Jonas A Finkler, Stefan Goedecker, and Jörg Behler. A fourth-generation highdimensional neural network potential with accurate electrostatics including non-local charge transfer. Nature communications , 12(1):398, 2021.

Sergey Pozdnyakov, Artem R Oganov, Efim Mazhnik, Arslan Mazitov, and Ivan Kruglov. Fast general two-and three-body interatomic potential. Physical Review B , 107(12):125160, 2023.

Stephen R Xie, Matthias Rupp, and Richard G Hennig. Ultra-fast interpretable machinelearning potentials. arXiv preprint arXiv:2110.00624 , 2021.

Cardinal b-splines. https://pages.cs.wisc.edu/~deboor/887/lec1new.pdf .

Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E Sarma, Michael M Bronstein, and Justin M Solomon. Dynamic graph cnn for learning on point clouds. Acm Transactions On Graphics (tog) , 38(5):1-12, 2019.

Emil Prodan and Walter Kohn. Nearsightedness of electronic matter. Proceedings of the National Academy of Sciences , 102(33):11635-11638, 2005.

Appendices

A Point Edge Transformer

In this section we provide additional details on the specific implementation of the PET model that we 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 is the central component of PET is complemented by several pre- and postprocessing units. Two embedding layers, E ( k ) c and E ( k ) n , are used to encode the chemical species of the central atom and all the neighbors, respectively. A 3 → d PET linear layer L ( k ) is used to map the 3-dimensional Cartesian coordinates of displacement vectors r ij to feature vectors of dimensionality d PET, which are further transformed by an activation function σ (we use SiLU[96] for all our experiments) to form a position embedding layer E ( k ) r = σ ( L ( k ) ( r ij ) . An MLP with one hidden layer, denoted as M ( k ) , with dimensions 3 d PET → d PET, 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 to the transformer associated with neighbor j is computed as follows:

where ˜ x ( k -1) ij denotes input message from atom j to the central atom i , s j is the chemical species of neighbor j . For the very first message-passing block, there are no input messages, and so M (1) performs a 2 d PET → d PET transformation. The token associated with the central atom is given simply as E ( k ) c ( s i ) .

In addition, each message-passing block has two post-processing units: two heads H ( k ) c and H ( k ) n to compute atomic and bond contributions to the total energy or any other target at each message passing block . Atomic contributions are computed as H ( k ) c ( x ( k ) i ) , where x ( k ) i is the output token associated with central atom after application of transformer. Bond (or edge) contributions are given as H ( k ) n ( x ( k ) ji ) f c ( r ij | R c , ∆ R c ) , where x ( k ) ji indicates the output token associated with the j -th neighbor. We modulate them with f c ( r ij | R c , ∆ R c ) 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.

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 r ij . Thus, the linear layer L ( k ) defines a 4 → d PET transformation, in contrast to the 3 → d PET of a standard model. Additionally, we 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. (5).

There is great flexibility in the pre-and post-processing steps before and after the application of the transformer. Our implementation supports several modifications of the procedure discussed above, which are controlled by a set of microarchitectural hyperparameters. For example, one can average bond contributions to the target property instead of summing them. In this case, the total contribution from all the bonds attached to central atom i is given as:

rather than

Another possibility is to not explicitly concatenate embeddings of neighbor species in Eq. (5), and instead define the input messages for the very first message-passing block as these neighbor species embeddings. This change ensures that all M ( k ) have an input dimensionality of 2 d PET across all the message-passing blocks. A brief description of these hyper-parameters accompanies the reference PET implementation that we discuss in Section D.

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-contributions from the targets. These atomic terms often constitute a 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 3 atomic species - H, C, and O. The 'bag of atoms' features for an H 2 O molecule are [2 , 0 , 1] . During inference, self-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 simultaneously on energies and forces, we use the following loss function:

where E and F are the ground-truth energies and forces, respectively, ˜ E is the energy prediction given by the model, N is the number of atoms in the sample, α runs over x, y, z , and MSE E and MSE F are the exponential moving averages of the Mean Squared Errors in energies and forces, respectively, on the validation dataset. We update MSE E and MSE F once per epoch. We found that for such a loss definition the choice of the dimensionless parameter w E is crelatively robust, with the best results achieved approximately for w E ≈ 0 . 03 -0 . 1 .

General details. In most cases, we use the StepLR 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 d PET = 128 , n GNN = n TL = 3 , multi-head attention with 4 heads, SiLU activation, and the dimension of the feedforward network model in the transformer is set to 512. We use a rotational augmentation strategy during fitting, which involves randomly rotating all samples at 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 d PET = 128 , we set the initial learning rate at 10 -4 . However, some models with d PET = 256 were unstable at this learning rate. Consequently, we adjusted the initial learning rate to 5 · 10 -5 for two cases: 1) CH 4 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.

C Detailed benchmark results

C.1 COLL

The COLL dataset contains both 'energy' and 'atomization_energy' targets, that can be used together with the force data as gradients, and that are entirely equivalent after subtracting self-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 a fallback a 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 r ij , as opposed to the entire displacement vector r ij , as in the standard PET. As a result, the model becomes intrinsically rotationally invariant but loses some of its expressive power: the rotationally-invariant 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 of the auxiliary model has a negligible impact on the accuracy of the overall model. We summarize the performance of several models in Table 2.

C.2 HME21

For the HME21 dataset we employ a learning rate scheduler that differs from the standard StepLR. 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 to a noticeable improvement in validation accuracy.

110

120

130

140

150

val mae

f

/ meV/Å

HME21

standard StepLR

fast lr decay

15

20

25

30

val MAE

f

/ meV/Å

Water 6_2_4.25

standard StepLR

fast lr decay

0

50

100

150

200

250

300

350

0.0

0.2

0.4

0.6

0.8

1.0

learning rate / 1e-4

standard StepLR

fast lr decay

0

500

1000

1500

2000

2500

0.0

0.2

0.4

0.6

0.8

1.0

learning rate / 1e-4

standard StepLR

fast lr decay

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 a modified loss is nearly identical to that obtained when targeting total energies.

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

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 a 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 a 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, and therefore it is impractical to use them for inference. Similar to what is done in Ref. 88, we discretize the local spin values, rounding them to the according to the rule ( -∞ , -0 . 25] → -1 , ( -0 . 25 , 0 . 25) → 0 , [0 . 25 , ∞ ) → 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 in energies of a model that doesn't use spins at all, which agrees within two significant digits with the analogous results reported in Ref. [88]. This suggests that both models were able to extract all possible energetic information from an incomplete description that disregards magnetism.

C.4 High-entropy alloys

The main dataset of High-Entropy-Alloy configurations, which we used for training, was generated by randomly assigning atom types chosen between 25 transition metals to lattice sites of fcc or bcc 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. 90. We used 24630 samples for training, 500 for validation, and 500 for testing. While Ref. 90 uses all the energies available for the training dataset, we acknowledge that it uses only part of the forces. This, however, can be attributed primarily to the considerable Random

Access Memory (RAM) demands imposed by the model introduced in Ref.90. The efficiency of PET allowed us to use all the available forces on a standard node of our GPU cluster.

In addition, Ref. 90 proposes additional 'out-of-sample' test scenarios. In particular, it provides molecular-dynamics trajectories of a configuration containing all 25 atom types, performed at temperatures of T = 300 K and T = 5000 K . For the trajectory at T = 300 K , the errors of the PET model (ECSE-symmetrized) are 9.5 meV/atom and 0.13 eV/Å for energies and forces, respectively, which compares favorably with the values of 14 meV/atom and 0.23 eV/Å obtained with the HEA25-4-NN model. For the T = 5000 K trajectory, which undergoes melting and therefore differs substantially from the crystalline structures the model is trained on, PET yields errors of 152 meV/atom and 0.28 eV/Å. HEA25-4-NN yelds lower energy error and comparable force error, 48 meV/atom and 0.29 eV/Å respectively.

C.5 Water

We randomly split the dataset[103] into a training subset of size 1303, a validation subset of size 100, and a testing subset of 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 n GNN = 6 , n TL = 2 and R c = 4 . 25 Å. We use this trick for the ablation studies illustrated in panels b and c of Fig. 3 in the main text and not for those depicted in panel a.

C.6 CH 4

This dataset contains more than 7 million configurations of a CH 4 molecule, whose atomic positions are randomized within a 3.5Å sphere, discarding configurations where two atoms would be closer than 0.5Å[104]. Our testing subset consists of configurations with indices from 3050000 to 3130000 in the original XYZ file. For the validation set, we use indices from 3000000 to 3005000. When we train a model with n train training configurations, we use the ones with indices from 0 to n train. These settings match those of Ref. [80]. We use a full validation dataset for fitting the d PET = 256 models and a 5k subset for all the d PET = 128 ones in order to reduce the computational cost.

For methane, the only metric of interest is the Root Mean Square Error (RMSE) in energies, even when the fitting involves both energies and forces. The loss function we use is defined in equation (8). This definition includes the dimensionless hyperparameter w E , which determines the relative importance of energies. We have observed that, surprisingly, the best accuracy in energies is achieved when w E ≪ 1 . For all our training runs, we used w E = 0 . 125 .

C.7 QM9 energies

We randomly split the dataset into a training subset of size 110000, a validation subset of size 10000, and a testing subset of size 10831. The reported MAE of 6.7 meV/molecule corresponds to the base model, without application of the ECSE protocol. Table 5 compares the accuracy of PET with several other models, including the current state-of-the-art model, Wigner Kernels[81].

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 scalar values of the norms of dipole moments contained in the original QM9 dataset. We use 1000 configurations as a testing subset, similarly to WK[81]. Our validation subset contains 500 molecules, and we use up to 19500 samples for training.

D Reproducibility Statements

In order to facilitate the reproducibility of the results we present, we have released the following assets: 1) the source code for the PET model, 2) the source code for our proof-of-principle implementation of the ECSE protocol, 3) a complete set of hyperparameters for each training procedure, organized as a collection of 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 are available at:

In addition, the most important details of the different models we trained are discussed in Appendix B.

E Computational resources

Training times substantially depend on the dataset used for fitting. For instance, fitting the PET model with d PET = 128 on the CH 4 dataset using 1000 energy-only training samples (the very first point of the learning curve in Fig. 3d) takes about 40 minutes on a V100 GPU. In contrast, fitting the PET model with d PET = 256 on energies and forces using 300,000 samples is estimated to 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 d PET = 256 model on the COLL dataset. It is worth mentioning, however, that the model already improves upon the previous state-of-the-art model in forces after approximately one-third of the total fitting time on the validation dataset. We used d PET = 256 models only for 3 cases: 1) CH 4 E+F, 100k samples, 2) CH 4 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 a single d PET = 128 model.

F ECSE

The overall idea behind the ECSE symmetrization protocol is relatively simple, but an efficient implementation requires a rather complicated construction, that we discuss here in detail.

0

1

2

3

r

0.0

0.5

1.0

f

c

(

r

|

Rc

=4,

Rc

=2)

0

1

2

3

z

q

1

c

(

z

|

=1,

=2)

0

1

2

3

z

q

2

c

(

z

|

=1,

=2)

F.1 Preliminary definitions

We begin by defining several mathematical functions that we use in what follows. In order to obtain smooth predictions, PET uses several cutoff functions. While there is a degree of flexibility in defining them, we found the following functional forms, that are also illustrated in Fig. 5, to be effective:

We also need smooth functions to prune the ensemble of coordinate systems. To do so, we use several variations on a well-known smooth approximation for the max function. Given a finite set of numbers { x i } i and a smoothness parameter β > 0 one can define:

This function satisfies the following properties:

For a set of just two numbers { x 1 , x 2 } one can show that SmoothMax is bounded from below:

for T ( β ) = W (exp( -1)) /β , where W is the Lambert W function.

For the case when the numbers x i are associated with weights p i ≥ 0 it is possible to extend the definition of SmoothMax as:

This function satisfies the properties in Eq. (13), where maximum is taken only of the subset of { x i } where p i > 0 . A useful consequence of the definition of SmoothMaxWeighted is that its value is not changed when including a point with zero weight, ( x, 0) :

Corresponding functions that provide a smooth approximation to the min operation, SmoothMin and SmoothMinWeighted , can be defined the same way as SmoothMax and SmoothMaxWeighted using ( -β ) instead of β .

Finally, given an ordered pair of noncollinear unit vectors ˆ v 1 and ˆ v 2 we define the associated coordinate system by computing u 1 = ˆ v 1 × ˆ v 2 , ˆ u 1 = u 1 / ∥ u 1 ∥ and ˆ u 2 = ˆ v 1 × ˆ u 1 . The reference frame is given by the vectors ˆ v 1 , ˆ u 1 and ˆ u 2 . This procedure generates only right-handed coordinate systems, and thus, the transformation between the absolute coordinate system to the atom-centered one can always be expressed as a proper rotation, never requiring inversion. When we later refer to the coordinate system given by a central atom i and two neighbors j and j ′ , we mean that the vectors ˆ v 1 and ˆ v 2 are the ones pointing from the central atom to neighbors j and j ′ respectively.

F.2 Assumptions

We assume that (1) there is a finite number of neighbors within any finite cutoff sphere and (2) there is a lower boundary d min 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 a relevant energy scale.

F.3 Fully-collinear problem

The simplest version of ECSE protocol was defined in the main text in the following way:

where

The q c function can be chosen as either q 1 c or q 2 c , while ω and ∆ ω are user-specified constant parameters.

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 ′ , | ˆ r j × ˆ r j ′ | 2 is smaller than ω which implies that all the angular cutoff values q c ( | ˆ r j × ˆ r j ′ | 2 | ω, ∆ ω ) are zeros, and so all the weights w jj ′ . As a result, the ECSE weighted average is undefined, falling to 0 0 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:

In this case, y S ( A i ) can fall to y aux ( A i ) if all the weights w jj ′ 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 to make predictions for corner cases. For this purpose we define w aux as follows:

With this expression, w aux takes a non-zero weight only when the smooth maximum of the coordinatesystem weights w jj ′ is below t aux : in other words, as long as at least one pair of neighbors is non-collinear, the auxiliary model will be ignored. Since SmoothMaxWeighted 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. β ω , t aux and ∆ aux are user-specified constant parameters.

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 ∆ ω that enter the definition (18) of the ECSE weights:

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-collinear singularity, the coordinate systems used by ECSE are not stable. However, given the definition of the coordinate systems discussed in Appendix F.1, the x 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 non-equivariant models is always the same, irrespective of the y and z axes, that will be arbitrarily oriented in the orthogonal plane. Thus, all evaluations of the backbone architecture y 0 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 x for all coordinate systems, making the ECSE average smooth.

One small detail is that, within this protocol, the orientation of the x axis can be arbitrary. This problem is easily solved by always including an additional coordinate system in which ˆ v 1 is oriented along -r ij rather than along r ij . Both coordinate systems should be used with the same weight w jj ′ . One should note, however, that this approach has a number of limitations. In particular, it doesn'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 R in 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

For each pair of neighbors, we use SmoothMax( r j , r ′ j | β ) to select the farthest distance in each pair

Applying a SmoothMin function to the collection of neighbor pairs

selects a distance that contains at least a 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. (13).

We ensure to pick at least one non-collinear pair (if there is one within R out) by defining the adaptive inner cutoff as

where p jj ′

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 at least one good pair of neighbors inside the outer cutoff sphere, there's at least one within the inner cutoff as well. If no such good pair exists within the outer cutoff sphere, the definition falls back to R out . Note that the expression includes also some 'tolerance" parameters ∆ R c and ∆ ω . 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 1.

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:

First, one computes the maximum weight present, using a SmoothMaxWeighted function. Then, pruning factors f jj ′ are computed using the cutoff q 1 c function, aiming to zero out the smallest weights with w jj ′ ← w jj ′ f jj ′ . An additional modification is needed to ensure smooth behavior around fullycollinear configurations, where all weights become zero. If this occurs, q 1 c ( w jj ′ | w MAX T f , w MAX ∆ T f ) converges to a step function. Therefore, it makes sense to activate weight pruning only if w MAX is larger than a certain threshold T e , which is implemented using a smooth activation factor e . We found this pruning to be more efficient if the q c function in Eq. (4) is selected to be q 2 c rather than q 1 c . 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 yields 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.

β is 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, R in, 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 a relatively low value to β generates a very smooth approximation of this second minimum, but it significantly overestimates the exact value, and therefore R in potentially includes more neighbor pairs than needed. In contrast, for a large value of β , 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 β > 0 , it might develop unphysical spikes in the derivatives, reducing its accuracy. To see more precisely the impact of β consider the limit of a large value (considering also that other adjustment parameters such as ∆ R c 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 β determines the characteristic size of this transition region. Thus, for excessively large values of β , there is a sharp (but still continuously differentiable) transition from one coordinate system to another. On the other hand, if β is set to a very low value, R in will always be significantly larger than the exact distance to the second neighbor. As a result, the ECSE protocol will constantly utilize multiple coordinate systems, leading to a smoother symmetrized model but with a higher computational cost.

In summary, the parameters used by the ECSE dictate the balance between smoothness and computational efficiency. We name the selection of such parameters that leads to 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.

F.7 Accuracy of a-posteriori-ECSE-symmetrized models

The most computationally efficient method for fitting a model is to initially train a non-equivariant backbone architecture with random rotational augmentations and then apply the ECSE protocol a posteriori. There are general considerations one can make on the expected accuracy of the symmetrized model relative to the initial, non-equivariant one.

The functional form of ECSE is a weighted average of predictions generated by a backbone architecture 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 a 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 to ∂ ˆ R jj ′ ∂ r k and ∂w jj ′ ∂ r k are not associated with the forces produced by a backbone architecture ∂ y 0 ([ A i ]) ∂ r k . Consequently, the error in forces of the final, symmetrized model may significantly exceed that of the backbone architecture. This issue is more pronounced for a 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 a more pronounced computational

overhead.

F.8 Additional rotational augmentations

In some cases, such as for the CH 4 configurations, the total number of all possible coordinate systems is very limited. However, using a larger number of rotational augmentations may still be beneficial during inference. For such cases, it's possible to incorporate additional augmentations into the ECSE protocol as follows:

where S aug is a predefined constant set of random rotations ˆ R aug .

F.9 Application of the ECSE protocol to long-range models

Several models capable of capturing long-range interactions have been developed within the atomistic machine-learning community. These models cannot be directly cast into the local energy decomposition discussed in Sec. 2. However, they are constructed as either pre- or post-processing operations for the local models. Thus, these schemes can also benefit from the ECSE protocol by making it possible to use some models developed for generic point clouds as a local engine.

An example of such models is LODE[110]. As a first step, it constructs a representation that is inspired by the electrostatic potential. This three-dimensional function, or a collection of threedimensional 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 be applied to LODE, or similar schemes, without any modifications.

Another example of a long-range model is given in Ref. [111]. This scheme applies local models to compute partial charges associated with each atom and later uses them to calculate a non-local electrostatic energy contribution. The ECSE protocol can ensure that these partial charges remain invariant with respect to rotations.

F.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-passing layers. At the first layer, atom A sends some message, denoted as m 1 AB , to atom B . At the second, atom B sends messages m 2 BC and m 2 BD 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 a prediction associated with atom D. Both calculations rely on the message m 1 AB , computed at the first layer of the GNN. The problem is that, in contrast to normal GNNs, this message cannot be reused to 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 1 AB 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 or the size of a unit cell is smaller than the receptive field of a 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 to 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 to a 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, symmetrizing 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, one can expect that a 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 to 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 to 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 a 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 is turned on immediately and symmetrized model is trained from the very beginning.

F.11 Current implementation

Our current implementation uses a 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-dimensional PyTorch tensors, each containing only one scalar. Thus, it has the potential to be substantially accelerated. We want to reiterate that we have selected loose parameters for the ECSE protocol corresponding to the best accuracy and substantial computational cost (see more details in Appendices F.6 and F.7). Therefore, even with the current proof-of-principle implementation, ECSE can be substantially accelerated at the cost of a minuscule drop in accuracy.

F.12 Numerical verification of smoothness

To verify that our implementation of the ECSE protocol indeed preserves the smoothness of the backbone architecture, we numerically analyzed alterations of the predictions with respect to random perturbations of the input. For this purpose, we selected the CH 4 dataset (discussed in the Appendix C) since it represents random configurations, thus covering most of the configurational space. For 100 molecules from this dataset, we applied 50 random Gaussian perturbations to 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 a backbone architecture, we used a PET model with hyperparameters similar to the one benchmarked in Appendix C. The parameters of ECSE were selected to be loose in a sense discussed in Appendix F.6.

Results are shown in Fig. 6. One can see that a small amplitude of the applied noise entails a minor alteration of the predictions. No instances were observed where a minor noise led to a significant change in the predictions of the ECSE symmetrized model.

ECSE symmetrized model

10-1

Non-symmetrized model

10-3

predictions

10-5

.

10-7

Difference i

10-9

10-11

10-13

10-9

10-7

10-5

10-3

10-1

L2 norm of the noise

0

1

2

3

x

0.0

0.3

0.6

0.9

x

2

2

2

x

2 +6

x

3

2

x

2

6

x

+9

2

B

2

(

x

)

a

x

B

2

0

B

2

1

B

2

2

b

G Architectures for generic point clouds made smooth

In the main text, we 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 a naive voxelization of a point cloud, voxel features are estimated based on the points inside the corresponding voxel. This means that this approach leads to discontinuities whenever a point moves from one voxel to another. The construction used in GTTP[112] and UF[113] is an example of a smooth projection of a point cloud onto the voxel grid. We first describe the smooth projection onto a one-dimensional pixel/voxel grid for simplicity. It relies on smooth, piecewise polynomial, bell-like functions known as B-splines.

Fig. 7a demonstrates the shape of a second-order B-spline. It spreads over three one-dimensional voxels, on each of them it is defined by a second-order polynomial, and is continuously differentiable. For any p , it is possible to construct an analogous function B p ( x ) which spreads over p +1 voxels, is a piecewise polynomial of order p , and is p -1 times continuously differentiable. Next, an equidistant collection of such B-splines is constructed, as shown in Fig. 7 (b). The projection onto the one-dimensional voxel grid is defined as follows:

where x k is the position of a k-th point. The coefficients c i can be treated as features associated with the i-th voxel.

Three-dimensional B-splines can be built as a tensor product of these functions, i.e.:

and the corresponding projection of a point cloud onto a 3D voxel grid is:

An alternative method to define a smooth projection is to define a smooth point density, and to compute its integrals for each voxel, as demonstrated in Fig. 8 for the one-dimensional case. The point density is defined as follows:

where

x

point density

Point positions

B splines centered at point positions

Point density

where B is any bell-like function. In 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:

If B is given by a B-spline of the p -th order, then B int is given by a B-spline of p +1 -th order[114].

Since the B-spline functions have compact support, the discussed projections preserve the sparsity of the voxel features.

Within a paradigm of local energy decomposition, one can center the voxel grid at 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)appearance of new atoms at the cutoff sphere. This problem can be avoided by modifying Eq. (28) as:

( r k here denotes the displacement vector from central atom to 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 a 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 voxel-based 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 be made smooth and applied using the 2D analogue of the methods discussed above.

and Eq. (29) as:

G.2 Point models

A common building block of various point methods is the following functional form [48, 115]:

where x i is a feature vector associated with a central point i , x j are feature vectors associated with the neighbors, γ and h are learnable functions, □ is a permutationally invariant aggregation function, and the neighborhood A i is defined either as the k nearest neighbors or as all the points within certain cutoff radius R c . The distance used to define local neighborhoods can be computed either in the initial space or between the feature vectors x at the current layer of the neural network.

The construction above contains several sources of discontinuities even for smooth functions γ 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 a fixed cutoff radius. In principle, one can make it adaptive, meaning it could be a smooth function of a neighborhood, embracing k neighbors on average.

If □ is given by summation, then, as it was already outlined in the main text for point convolutions, the smooth form can be constructed as:

where d ij is the distance between points i and j .

If □ is given by a maximum, then one can ensure smoothness using an expression such as:

where f max c ( d ij | R c , ∆ R c ) is some smooth function which converges to -∞ at R c . Minimum and average aggregations can be adapted analogously.

We are unaware of any simple modification to ensure smoothness in the downsampling operations that are applied in PointNet++-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 to 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, or neural network layers, can be made smooth and used to construct 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 a 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.

Peer reviews:
comments: 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.

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.
- The  approach using an ensemble of coordinate systems seems innovative and makes sense.

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.
- 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
- Are competing methods also using rotational augmentations?
- -Do the authors have any idea why the SOTA in tab 1 is so much worse on MnO & HEA?

QUESTIONS:
see weaknesses

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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\nQUESTIONS:\nI'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\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nS1. 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.\n\nWEAKNESSES:\nW1. 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.\n\nQUESTIONS:\nQ1. 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?\n\nLIMITATIONS:\nThe authors adequately address the limitations and broader impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe idea to enforce smoothness with other symmetries such as rotational symmetries seems to be valuable, and indeed missing from existing solutions.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nI would appreciate any response from the authors regarding the weakness stated above.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}