SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sergey N. Pozdnyakov and Michele Ceriotti
Laboratory of Computational Science and Modelling,
Institute of Materials, Ecole Polytechnique Federale de Lausanne,
Lausanne 1015, Switzerland
sergey.pozdnyakov@epfl.ch, michele.ceriotti@epfl.ch
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.
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  $\mathrm{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 = \sum_{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_{\mathrm{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 explicitly 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
$$
\left(\mathcal {F} * g\right) (\mathbf {r}) _ {m} = \sum_ {i} \sum_ {n} ^ {N _ {i n}} g _ {m n} \left(\mathbf {r} _ {i} - \mathbf {r}\right) f _ {n} ^ {i}, \tag {1}
$$
where,  $f_{n}^{i}$  is the  $n$ -th feature of the point  $i$ , and  $g$  is a collection of  $N_{in}N_{out}$  learnable three-dimensional functions, where  $N_{in}$  and  $N_{out}$  are the numbers of input and output features respectively. The output features  $(\mathcal{F}*g)(\mathbf{r})_m$  can be evaluated at an arbitrary point  $\mathbf{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(\mathbf{r}_1,\mathbf{r}_2,\dots,\mathbf{r}_n) = \gamma (\max_i(\{h(\mathbf{r}_i)\}))$ , where  $\gamma$  and  $h$  are learnable functions, can approximate to arbitrary precision any continuous function of the point set  $\{\mathbf{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:
$$
\left(\mathcal {F} * g\right) (\mathbf {r}) _ {m} = \sum_ {i} \sum_ {n} ^ {N _ {i n}} g _ {m n} \left(\mathbf {r} _ {i} - \mathbf {r}\right) f _ {n} ^ {i} f _ {c} \left(\| \mathbf {r} _ {i} - \mathbf {r} \| \mid R _ {\mathrm {c}}, \Delta_ {R _ {\mathrm {c}}}\right), \tag {2}
$$
where the cutoff function  $f_{\mathrm{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  $\Delta_{R_{\mathrm{c}}} > 0$  controls how fast it converges to 1 for  $r < R_{\mathrm{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  $\mathbf{y}(A) = \sum_{i} \mathbf{y}(A_i)$  where the reference point is given by the position of the central atom  $\mathbf{r}_i$ , and the model estimates an atomic contribution  $\mathbf{y}(A_i)$  given the local neighborhood. We name our approach Equivariant Coordinate System Ensemble (ECSE, pron. efe).
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:
$$
\mathbf {y} _ {\mathrm {S}} \left(A _ {i}\right) = \sum_ {j j ^ {\prime} \in A _ {i}} w _ {j j ^ {\prime}} \hat {R} _ {j j ^ {\prime}} \left[ \mathbf {y} _ {0} \left(\hat {R} _ {j j ^ {\prime}} ^ {- 1} [ A _ {i} ]\right) \right] / \sum_ {j j ^ {\prime} \in A _ {i}} w _ {j j ^ {\prime}}, \tag {3}
$$
Figure 1: (a) Equivariant coordinate-system ensemble: Each ordered pair of neighbors defines a local coordinate system. Next, an atomic environment is projected on all of them (which is equivalent to rotation) and used as input for a backbone architecture. If outputs are covariant, such as vectors, they are rotated back to the initial coordinate system. Finally, predictions are averaged over. (b) Discontinuities related to plain average. The weighted average with weights  $w_{jj'}$  resolves these problems. (c) Cutoff functions  $f_c$  and  $q_c$  used to define weights  $w_{jj'}$  (d) To reduce the computational cost, an adaptive cutoff  $R_{\mathrm{in}}$  is used, which adjusts to a given geometry instead of being a global user-specified constant.
where  $\mathbf{y}_S$  indicates the symmetrized model,  $\mathbf{y}_0$  the initial (non-equivariant) backbone architecture, and  $\hat{R}_{jj'}[\cdot]$  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,  $\hat{R}_{jj'}^{-1}[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  $\mathbf{r}_{ij}$  and  $\mathbf{r}_{ij'}$  the corresponding coordinate system consists of the vectors  $\hat{r}_{ij}, \hat{r}_{ij} \times \hat{r}_{ij'}$ , and  $\hat{r}_{ij} \times [\hat{r}_{ij} \times \hat{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  $\hat{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 structures[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  $(\mathbf{r}_{ij},\mathbf{r}_{ij'})$  of the neighbors that define the reference frame:
$$
w _ {j j ^ {\prime}} = w \left(\mathbf {r} _ {i j}, \mathbf {r} _ {i j ^ {\prime}}\right) = f _ {\mathrm {c}} \left(r _ {i j} \mid R _ {\mathrm {c}}, \Delta_ {R _ {\mathrm {c}}}\right) f _ {\mathrm {c}} \left(r _ {i j ^ {\prime}} \mid R _ {\mathrm {c}}, \Delta_ {R _ {\mathrm {c}}}\right) q _ {\mathrm {c}} \left(\left| \hat {r} _ {i j} \times \hat {r} _ {i j ^ {\prime}} \right| ^ {2} \mid \omega , \Delta_ {\omega}\right), \tag {4}
$$
where the cutoff functors  $f_{\mathrm{c}}$  and  $q_{\mathrm{c}}$  are illustrated in Fig. 5c. Thanks to the presence of  $f_{\mathrm{c}}$ , the terms in Eq. (3) that are associated with atoms entering or exiting the cutoff sphere have zero weights. Similarly,  $q_{\mathrm{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  $\frac{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_{\mathrm{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_{\mathrm{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_{\mathrm{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_{\mathrm{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_{\mathrm{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 reduces the cost of evaluating the ECSE-symmetrized model.