SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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-
a
point cloud decorated with chemical species
b
Figure 2: Architecture of the Point-Edge Transformer (PET). White and colored boxes represent layers; gray boxes and lines represent data. (a) PET is a message-passing architecture. At each of the  $n_{\mathrm{GNN}}$  message-passing (MP) interactions, messages are communicated between all the pairs of atoms closer than a certain cutoff distance  $R_{c}$ . At each stage, the corresponding MP block computes output messages and predictions of the target property given the input messages and geometry of the point cloud. (b) For each atom in the system, we define atom-centered environment  $A_{i}$  as a collection of all the neighbors within the cutoff distance  $R_{c}$ . The MP block is applied to each such atomic environment. Given 1) the geometry of the atomic environment, 2) the chemical species of the atoms, and 3) input messages from all the neighbors to the central atom it produces output messages from the central atom to all the neighbors and contribution to the prediction of the target property. The first step is to encode all the information associated with each neighbor to an abstract token of dimensionality  $d_{\mathrm{PET}}$ . Next, the collection of such tokens (with the one associated with the central atom) is fed into the transformer with  $n_{\mathrm{TL}}$  self-attention layers. The transformer does permutationally covariant transformation. Thus, the association between the tokens and neighbors is preserved. Therefore, we can simply treat output tokens as output messages to the corresponding neighbors. (c) The Encoder layer first maps all the sources of information into dimensionality  $d_{\mathrm{PET}}$ . Next, all 3 tokens are concatenated and compressed into a single one of the desired size.
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_{\mathrm{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  $\alpha_{ij}$  determining the contribution from token  $j$  to token  $i$  are modified as  $\alpha_{ij} \gets \alpha_{ij} f_c(r_j | R_c, \Delta_{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 \times 2$  model outperforms a large  $12 \times 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  $\sim 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_{\mathrm{c}}$  beyond  $3.7\AA$  (a value we chose to ensure approximately 20 neighbors on average), reaching a force MAE of  $14.4\mathrm{meV}/\AA$  when using a cutoff of  $4.25\AA$ .
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 concise rotaional invariance and universal approximation[17, 87]. PET improves by  $\sim 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  $\mathrm{CH}_4$
Figure 3: (a-c) Accuracy of PET potentials  $(y_0)$  of liquid water, compared with NEQUIP[35]. (a) Accuracy for different numbers of message-passing blocks  $n_{\mathrm{GNN}}$  and transformer layers  $n_{\mathrm{TL}}$ ; (b) Accuracy as a function of  $n_{\mathrm{GNN}}$ , for constant  $n_{\mathrm{GNN}} \times n_{\mathrm{TL}} = 12$ ; (c) Accuracy as a function of cutoff. (d-f) Learning curves for different molecular data sets, comparing symmetrized PET models  $(y_{\mathrm{S}})$  with several previous works[23, 78-84], including the current state of the art. (d) Random  $\mathrm{CH}_4$  dataset, training only on energies; (e) Random  $\mathrm{CH}_4$  dataset, training on energies and forces; (f) Vectorial dipole moments in the QM9 dataset[82].
configurations first introduced in Ref. 80. This dataset contains 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).
Table 1: Comparison of the accuracy of PET and current state-of-the-art models for the COLL, MnO, HM21 and HEA data sets. A more comprehensive comparison with leading models is provided in the Appendix C. Energy errors are given in meV/atom, force errors in meV/Å. In all cases the PET model is nearly equivariant, and the difference between the accuracy of  $y_{0}$  and  $y_{S}$  is minuscule. For HME21, we report error bars over 5 random seeds, and results for an ensemble of symmetrized models.
dataset metric	COLL	MnO	HME21	HEA
MAE f	MAE E	RMSE f	RMSE E/at.	MAE |f|	MAE E/at.	MAE f	MAE E/at.
SOTA model	26.4[17] GemNet	47[86] DimeNet++	125[88] mHDNNP	1.11[88]	138[89] MACE	15.7[89]	190[90] HEA25-4-NN	10[90]
PET (y0)	23.1	12.0	22.7	0.312	140.5 ± 2.0	17.8 ± 0.1	60.2	1.87
PET (yS)	23.1	11.9	22.7	0.304	141.6 ± 1.9	17.8 ± 0.1	60.1	1.87
PET (ens.)	128.5	16.8
The PET performs well even for condensed-phase datasets that contain dozens of different atomic types. The HEA dataset contains distorted crystalline structures with up to 25 transition metals simultaneously[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 3(5)-fold reduction of force的能量) 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\mathrm{meV / at}$
MAE at  $5000\mathrm{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\mathrm{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\mathrm{meV / molecule}$ , in line with the accuracy of DimeNet++[87], but not as good as the  $4.3\mathrm{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.
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-of-the-art performance of PET on several benchmarks reinforces early indications that non-equivariant models trained with rotational augmentation can outperform rigorously equivariant ones. Similar observations were also made independently by other groups. Spherical Channel Network[92] and ForceNet[93] achieve excellent performance on the Open Catalyst dataset[94] 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.
We thank Marco Eckhoff and Jörg Behler for sharing the MnO dataset with the collinear spins and Artem Malakhov for useful discussions.
This work was supported by the Swiss Platform for Advanced Scientific Computing (PASC) and the NCCR MARVEL, funded by the Swiss National Science Foundation (SNSF, grant number 182892).
[1] 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.
[2] Michael M. Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geometric Deep Learning: Going beyond Euclidean data. IEEE Signal Process. Mag., 34(4):18-42, July 2017.
[3] Giuseppe Carleo, Ignacio Cirac, Kyle Cranmer, Laurent Daudet, Maria Schuld, Naftali Tishby, Leslie Vogt-Maranto, and Lenka Zdeborova. Machine learning and the physical sciences. Rev. Mod. Phys., 91(4):045002, December 2019.
[4] 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.
[5] 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.
[6] Danny E. P. Vanpoucke, Kurt Lejaeghere, Veronique Van Speybroeck, Michel Waroquier, and An Ghysels. Mechanical Properties from Periodic Plane Wave Quantum Mechanical Codes: The Challenge of the Flexible Nanoporous MIL-47(V) Framework. J. Phys. Chem. C, 119(41):23752-23766, October 2015.
[7] Michael F Herbst and Antoine Levitt. Black-box inhomogeneous preconditioning for self-consistent field iterations in density functional theory. J. Phys.: Condens. Matter, 33(8):085503, February 2021.
[8] Wahyu Setyawan and Stefano Curtarolo. High-throughput electronic band structure calculations: Challenges and tools. Computational Materials Science, 49(2):299-312, August 2010.
[9] Valerio Vitale, Giovanni Pizzi, Antimo Marrazzo, Jonathan R. Yates, Nicola Marzari, and Arash A. Mostofi. Automated high-throughput Wannierisation. npj Comput Mater, 6(1):66, December 2020.
[10] Xiaojing Gong, Jingyuan Li, Hangjun Lu, Rongzheng Wan, Jichen Li, Jun Hu, and Haiping Fang. A charge-driven molecular water pump. Nature Nanotech, 2(11):709-712, November 2007.
[11] Jirasak Wong-ekkabut, Markus S. Miettinen, Cristiano Dias, and Mikko Karttunen. Static charges cannot drive a continuous flow of water molecules through a carbon nanotube. Nature Nanotech, 5(8):555–557, August 2010.
[12] Jörg Behler and Michele Parrinello. Generalized Neural-Network Representation of High-Dimensional Potential-Energy Surfaces. Phys. Rev. Lett., 98(14):146401, April 2007.
[13] Albert P. Bartók, Mike C. Payne, Risi Kondor, and Gábor Csányi. Gaussian Approximation Potentials: The Accuracy of Quantum Mechanics, without the Electrons. Phys. Rev. Lett., 104(13):136403, April 2010.