SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 three-dimensional functions for each atomic species, captures the long-range information. This global potential is then used to define local descriptors, that are combined with other short-range features to model a local energy decomposition. Thus, the ECSE protocol can 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_{AB}^{1}$ , to atom  $B$ . At the second, atom  $B$  sends messages  $m_{BC}^{2}$  and  $m_{BD}^{2}$  to atoms  $C$  and  $D$ , respectively. The second layer additionally 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_{AB}^{1}$ , 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_{AB}^{1}$  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  $\mathrm{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.
Figure 6: Numerical verification of preserving smoothness. Each point represents one random perturbation of one  $\mathrm{CH}_4$  molecule. The horizontal coordinate is an amplitude of Gaussian noise applied to all the positions of all the atoms. Vertical coordinate represents the difference in the predictions. The experiment was conducted for both symmetrized and non-symmetrized PET models.
Figure 7: (a) B-spline of 2nd order. (b) Equidistant collection of B-splines. The figure is adapted from Ref. [112]
G Architectures for generic point clouds made smooth
In the main text, 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:
$$
c _ {i} = \sum_ {k} B _ {i} ^ {p} \left(x _ {k}\right), \tag {26}
$$
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.:
$$
B _ {i _ {1} i _ {2} i _ {3}} ^ {p} (x, y, z) = B _ {i _ {1}} ^ {p} (x) B _ {i _ {2}} ^ {p} (y) B _ {i _ {3}} ^ {p} (z), \tag {27}
$$
and the corresponding projection of a point cloud onto a 3D voxel grid is:
$$
c _ {i _ {1} i _ {2} i _ {3}} = \sum_ {k} B _ {i _ {1} i _ {2} i _ {3}} ^ {p} (\mathbf {r} _ {k}). \tag {28}
$$
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:
$$
\rho (\mathbf {r}) = \sum_ {k} B \left(\mathbf {r} - \mathbf {r} _ {k}\right), \tag {29}
$$
Figure 8: Integral projection. The integral of the red area determines the features of the third left voxel.
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:
$$
\int_ {\substack {\text {voxel} \\ i _ {1} i _ {2} i _ {3}}} \rho (\mathbf {r}) d \mathbf {r} = \sum_ {k} B _ {i _ {1} i _ {2} i _ {3}} ^ {\text {int}} (\mathbf {r} _ {k}), \tag{30}
$$
$$
B _ {i _ {1} i _ {2} i _ {3}} ^ {\text {i n t}} \left(\mathbf {r} _ {k}\right) = \int_ {\substack {\text {v o x e l} \\ i _ {1} i _ {2} i _ {3}}} B \left(\mathbf {r} - \mathbf {r} _ {k}\right) d \mathbf {r}. \tag{31}
$$
If  $B$  is given by a B-spline of the  $p$ -th order, then  $B^{\mathrm{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:
$$
c _ {i _ {1} i _ {2} i _ {3}} = \sum_ {k} f _ {c} \left(r _ {k} \mid R _ {c}, \Delta_ {R _ {c}}\right) B _ {i _ {1} i _ {2} i _ {3}} ^ {p} \left(\mathbf {r} _ {k}\right), \tag {32}
$$
and Eq. (29) as:
$$
\rho (\mathbf {r}) = \sum_ {k} f _ {c} \left(r _ {k} \mid R _ {c}, \Delta_ {R _ {c}}\right) B \left(\mathbf {r} - \mathbf {r} _ {k}\right). \tag {33}
$$
$(\mathbf{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.
G.2 Point models
A common building block of various point methods is the following functional form [48, 115]:
$$
\mathbf {x} _ {i} ^ {\prime} = \gamma (\square (\{h \left(\mathbf {x} _ {i}, \mathbf {x} _ {j}\right) | j \in A _ {i} \}), \tag {34}
$$
where  $\mathbf{x}_i$  is a feature vector associated with a central point  $i$ ,  $\mathbf{x}_j$  are feature vectors associated with the neighbors,  $\gamma$  and  $h$  are learnable functions,  $\square$  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  $\mathbf{x}$  at the current layer of the neural network.
The construction above contains several sources of discontinuities even for smooth functions  $\gamma$  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  $\square$  is given by summation, then, as it was already outlined in the main text for point convolutions, the smooth form can be constructed as:
$$
\mathbf {x} _ {i} ^ {\prime} = \gamma \left(\sum_ {j} f _ {c} \left(d _ {i j} \mid R _ {c}, \Delta_ {R _ {c}}\right) h \left(\mathbf {x} _ {i}, \mathbf {x} _ {j}\right)\right), \tag {35}
$$
where  $d_{ij}$  is the distance between points  $i$  and  $j$ .
If  $\square$  is given by a maximum, then one can ensure smoothness using an expression such as:
$$
\mathbf {x} _ {i} ^ {\prime} = \gamma (\operatorname {S m o o t h M a x} \left(\left\{f _ {c} ^ {\max } \left(d _ {i j} \mid R _ {c}, \Delta_ {R _ {c}}\right) h \left(\mathbf {x} _ {i}, \mathbf {x} _ {j}\right) \mid j \in \text {n e i g h b o r h o o d} \left(\mathbf {x} _ {i}\right) \right\}\right)), \tag {36}
$$
where  $f_{c}^{max}(d_{ij}|R_{c},\Delta_{R_{c}})$  is some smooth function which converges to  $-\infty$  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.