SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[100] 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.
[101] So Takamoto, Satoshi Izumi, and Ju Li. Teanet: Universal neural network interatomic potential inspired by iterative electronic relaxations. Computational Materials Science, 207:111280, 2022.
[102] Nataliya Lopanitsyna, Guillaume Fraux, Maximilian A. Springer, Sandip De, and Michele Ceriotti. Dataset: Modeling high-entropy transition-metal alloys with alchemical compression: Dataset HEA25, April 2023.
[103] Bingqing Cheng, Edgar Engel, Jörg Behler, Christoph Dellago, and Michele Ceriotti. Dataset: Ab initio thermodynamics of liquid and solid water, 2018.
[104] Sergey Pozdnyakov, Michael Willatt, and Michele Ceriotti. Dataset: Randomly-displaced methane configurations, 2020.
[105] Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
[106] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages 9323-9332. PMLR, 2021.
[107] Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Velicković, James Kirkpatrick, and Peter Battaglia. Simple gnn regularisation for 3d molecular property prediction & beyond. arXiv preprint arXiv:2106.07971, 2021.
[108] 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.
[109] Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations, 2022.
[110] Andrea Grisafi and Michele Ceriotti. Incorporating long-range physics in atomic-scale machine learning. J. Chem. Phys., 151(20):204105, November 2019.
[111] Tsz Wai Ko, Jonas A Finkler, Stefan Goedecker, and Jörg Behler. A fourth-generation high-dimensional neural network potential with accurate electrostatics including non-local charge transfer. Nature communications, 12(1):398, 2021.
[112] 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.
[113] Stephen R Xie, Matthias Rupp, and Richard G Hennig. Ultra-fast interpretable machine-learning potentials. arXiv preprint arXiv:2110.00624, 2021.
[114] Cardinal b-splines. https://pages.cs.wisc.edu/~deboor/887/lec1new.pdf.
[115] 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.
[116] Emil Prodan and Walter Kohn. Nearsightedness of electronic matter. Proceedings of the National Academy of Sciences, 102(33):11635-11638, 2005.
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 post-processing units. Two embedding layers,  $E_{\mathrm{c}}^{(k)}$  and  $E_{\mathrm{n}}^{(k)}$ , are used to encode the chemical species of the central atom and all the neighbors, respectively. A  $3 \rightarrow d_{\mathrm{PET}}$  linear layer  $L^{(k)}$  is used to map the 3-dimensional Cartesian coordinates of displacement vectors  $\mathbf{r}_{ij}$  to feature vectors of dimensionality  $d_{\mathrm{PET}}$ , which are further transformed by an activation function  $\sigma$  (we use SiLU[96] for all our experiments) to form a position embedding layer  $E_{T}^{(k)} = \sigma(L^{(k)}(\mathbf{r}_{ij})$ . An MLP with one hidden layer, denoted as  $M^{(k)}$ , with dimensions  $3d_{\mathrm{PET}} \rightarrow d_{\mathrm{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:
$$
t _ {j} ^ {(k)} = M ^ {(k)} \left[ \text {c o n c a t e n a t e} \left(\tilde {x} _ {i j} ^ {(k - 1)}, E _ {\mathrm {n}} ^ {(k)} \left(s _ {j}\right), E _ {r} ^ {(k)} \left(\mathbf {r} _ {i j}\right)\right) \right], \tag {5}
$$
where  $\tilde{x}_{ij}^{(k-1)}$  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  $2d_{\mathrm{PET}} \rightarrow d_{\mathrm{PET}}$  transformation. The token associated with the central atom is given simply as  $E_{\mathrm{c}}^{(k)}(s_i)$ .
In addition, each message-passing block has two post-processing units: two heads  $H_{\mathrm{c}}^{(k)}$  and  $H_{\mathrm{n}}^{(k)}$  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_{\mathrm{c}}^{(k)}(x_{i}^{(k)})$ , where  $x_{i}^{(k)}$  is the output token associated with central atom after application of transformer. Bond (or edge) contributions are given as  $H_{\mathrm{n}}^{(k)}(x_{ji}^{(k)})f_{c}(r_{ij}|R_{c},\Delta_{R_{c}})$ , where  $x_{ji}^{(k)}$  indicates the output token associated with the  $j$ -th neighbor. We modulate them with  $f_{c}(r_{ij}|R_{c},\Delta_{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  $\mathbf{r}_{ij}$ . Thus, the linear layer  $L^{(k)}$  defines a  $4 \rightarrow d_{\mathrm{PET}}$  transformation, in contrast to the  $3 \rightarrow d_{\mathrm{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:
$$
y _ {i} ^ {\text {b o n d} (k)} = \frac {\sum_ {j \in A _ {i}} H _ {\mathfrak {n}} ^ {(k)} \left(x _ {j i} ^ {(k)}\right) f _ {c} \left(r _ {i j} \mid R _ {c} , \Delta_ {R _ {c}}\right)}{\sum_ {j \in A _ {i}} f _ {c} \left(r _ {i j} \mid R _ {c} , \Delta_ {R _ {c}}\right)} \tag {6}
$$
rather than
$$
y _ {i} ^ {\text {b o n d} (k)} = \sum_ {j \in A _ {i}} H _ {\mathrm {n}} ^ {(k)} \left(x _ {j i} ^ {(k)}\right) f _ {c} \left(r _ {i j} \mid R _ {c}, \Delta_ {R _ {c}}\right). \tag {7}
$$
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  $2d_{\mathrm{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  $\mathrm{H}_2\mathrm{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:
$$
L = w _ {E} \frac {(\tilde {E} - E) ^ {2}}{\mathrm {M S E} _ {E}} + \frac {\frac {1}{3 N} \sum_ {i \alpha} \left(- \frac {\partial \tilde {E}}{\partial \mathbf {r} _ {i \alpha}} - F _ {i \alpha}\right) ^ {2}}{\mathrm {M S E} _ {F}}, \tag {8}
$$
where  $E$  and  $F$  are the ground-truth energies and forces, respectively,  $\tilde{E}$  is the energy prediction given by the model,  $N$  is the number of atoms in the sample,  $\alpha$  runs over  $x,y,z$ , and  $\mathrm{MSE}_E$  and  $\mathrm{MSE}_F$  are the exponential moving averages of the Mean Squared Errors in energies and forces, respectively, on the validation dataset. We update  $\mathrm{MSE}_E$  and  $\mathrm{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} \approx 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_{\mathrm{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_{\mathrm{PET}} = 128$ , we set the initial learning rate at  $10^{-4}$ . However, some models with  $d_{\mathrm{PET}} = 256$  were unstable at this learning rate. Consequently, we adjusted the initial learning rate to  $5 \cdot 10^{-5}$  for two cases: 1)  $\mathrm{CH}_4 \mathrm{E} + \mathrm{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  $\mathbf{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.
Table 2: Comparison of the accuracy of different versions of PET on the COLL dataset, including also the models reported in Ref. [17]. PET-256 is the main model with  $d_{PET} = 256$ . PET-128-2-body is the auxiliary internally invariant model with  $d_{PET} = 128$ . PET-128 is a normal PET model  $d_{PET} = 128$  and all the other settings match the ones of PET-128-2-body. PET-ECSE is an overall rotationally invariant model constructed with PET-256 as the main model and PET-128-2-body as the auxiliary one. The accuracy of SchNet is reported in Ref. [17].
model	MAE f, meV/Å	MAE E, meV/molecule
SchNet[31]	172	198
DimeNet++[86]	40	47
GemNet[17]	26.4	53
PET-128-2-body	56.7	27.4
PET-128	29.8	15.3
PET-256	23.1	12.0
PET-ECSE	23.1	11.9
C.2 HME21
For the HME21 dataset 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.
Figure 4: Comparison of a standard StepLR learning rate schedule and the strategy we use for the HME21 and water datasets.
In addition, we slightly modify the loss function for this dataset. The error metric for energies used in prior work is the Mean Absolute Error (MAE) per atom, computed first by determining energies per atom for both the ground truth energies and predictions, then by calculating the MAE for these normalized values. Therefore, it is logical to adjust the part of the loss function related to energies to reflect this metric. We re-define the loss as the mean squared error between the ground truth energies and predictions per atom. This is equivalent to the weighted loss function where weights are given by the inverse of the number of atoms in the structure. We observe that the performance of the model fitted with such 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 slightly improve the test accuracy.
Table 3: A comparison between the accuracy of PET and that of SchNet, TeaNet, PaiNN and NequIP (values reproduced from Ref. [100]) and MACE [89]. We compute 5 different PET models, and report both the mean errors (including also the standard deviation between the models) as well as the error on the ensemble models obtained by averaging the predictions of the 5 models. The regularizing effect of ensemble averaging leads to a noticeable improvement in accuracy. We also compare the results from a the non-symmetrized model (PET-1), those obtained by averaging 10 random orientations (PET-10) and those with ECSE applied upon inference.
model	MAE E, meV/at.	MAE |f|, meV/Å	MAE f, meV/Å
MACE[89]	15.7	138
TeaNet[101]	19.6	174	153
SchNet[31]	33.6	283	247
PaiNN[18]	22.9	237	208
NequIP[35]	47.8	199	175
PET-1	17.8 ± 0.1	140.5 ± 2.0	124.0 ± 1.6
PET-1-ens	16.8	128.1	113.6
PET-10	17.7 ± 0.1	139.9 ± 1.9	123.5 ± 1.6
PET-10-ens	16.8	127.9	113.5
PET-ECSE	17.8 ± 0.1	141.6 ± 1.9	125.1 ± 1.6
PET-ECSE-ens	16.8	128.5	114
C.3 MnO
We randomly split the dataset into a training subset of 2608 structures, a validation subset of 200 structures, and a testing subset of 293 structures. Thus, the total size of our training and validation subsets matches the size of the training subset in Ref. [88]. One should note that we use all forces in each epoch, while Ref. [88] only uses 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  $(-\infty, -0.25] \to -1$ ,  $(-0.25, 0.25) \to 0$ ,  $[0.25, \infty) \to 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
Table 4: Comparison between the accuracy of the high-dimensional neural network model of Ref. 88, and that of several (ECSE-symmetrized) PET models. Spins-discretized models approximate the atomic spins to integer values, PET-ECSE-spins uses the real-valued spins as atomic labels, while PET-ECSE-no-spins treats all Mn atoms as if they had no magnetization.
model	MAE E, meV	MAE f, meV/Å
mHDNNP-spins-discretized[88]	1.11	125
PET-ECSE-spins-discretized	0.30	22.7
PET-ECSE-spins	0.14	8.5
PET-ECSE-no-spins	11.6	94.8
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 = 300K$  and  $T = 5000K$ . For the trajectory at  $T = 300K$ , 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 = 5000K$  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 yields 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_{\mathrm{GNN}} = 6$ ,  $n_{\mathrm{TL}} = 2$  and  $R_{c} = 4.25\AA$ . 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 CH4
This dataset contains more than 7 million configurations of a  $\mathrm{CH}_4$  molecule, whose atomic positions are randomized within a  $3.5\AA$  sphere, discarding configurations where two atoms would be closer than  $0.5\AA[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_{\mathrm{train}}$  training configurations, we use the ones with indices from 0 to  $n_{\mathrm{train}}$ . These settings match those of Ref. [80]. We use a full validation dataset for fitting the  $d_{\mathrm{PET}} = 256$  models and a 5k subset for all the  $d_{\mathrm{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} \ll 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\mathrm{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].
Table 5: A comparison of the accuracy of several recent models for predicting the atomization energy of molecules in the QM9 dataset.
Model	Test MAE (meV)
Cormorant[105]	22
Schnet[31]	14
EGNN[106]	11
NoisyNodes[107]	7.3
SphereNet[108]	6.3
DimeNet++[86]	6.3
ET[109]	6.2
PaiNN[18]	5.9
Allegro[72]	4.7 ± 0.2
WK[81]	4.3 ± 0.1
PET	6.7
C.8 QM9 dipoles
The dataset from Ref. [82] contains vectorial dipole moments computed for 21000 molecules from the QM9 dataset. These are not to be confused with the 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:
https://doi.org/10.5281/zenodo.7967079