SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Inductive bias. Although theoretically there is no need to distinguish tuples with different equality patterns explicitly, we still do so to incorporate inductive bias into the model for better learning. Specifically, we modify the initialization function and the output function as follows: 1. At the initialization step, we learn embeddings for different equality patterns and incorporate them into the results of Equation (36) through a Hadamard product. 2. At the output step, we separate tuples where all sub-nodes are the same and the others into different multisets. These modifications are both beneficial for training and generalization.
D Detailed Experiments
D.1 Experimental Setup
Table 4: Training settings.
MD17	revised MD17	QM9
Train set size	1000	950	110000
Val. set size	1000	50	10000
batch size	2	2	16
warm-up epochs	25	25	5
initial learning rate	0.001	0.001	0.0005
decay on plateau patience (epochs)	15	15	10
decay on plateau cooldown (epochs)	15	15	10
decay on plateau threshold	0.001	0.001	0.001
decay on plateau factor	0.7	0.7	0.5
Training Setting. For QM9, we use the mean squared error (MSE) loss for training. For MD17, we use the weighted loss function
$$
\mathcal {L} (\boldsymbol {X}, \boldsymbol {z}) = (1 - \rho) | f _ {\theta} (\boldsymbol {X}, \boldsymbol {z}) - \hat {t} (\boldsymbol {X}, \boldsymbol {z}) | + \frac {\rho}{N} \sum_ {i = 1} ^ {N} \sqrt {\sum_ {\alpha = 1} ^ {3} \left(- \frac {\partial f _ {\theta} (\boldsymbol {X} , \boldsymbol {z})}{\partial \boldsymbol {x} _ {i \alpha}} - \hat {F} _ {i \alpha} (\boldsymbol {X} , \boldsymbol {z})\right) ^ {2}}, \tag {37}
$$
where the force ratio  $\rho$  is fixed as 0.999 (For revised MD17, we set  $\rho$  to 0.99 for several targets since the data quality is better).
We follow the same dataset split as GemNet [Gasteiger et al., 2021]. We optimize all models using Adam [Kingma and Ba, 2014] with exponential decay and plateau decay learning rate schedulers, and also a linear learning rate warm-up. To prevent overfitting, we use early stopping on validation loss and an exponential moving average (EMA) with decay rate 0.99 for model parameters during validation and test. We follow DimeNet's [Klicpera et al., 2020a] setting to calculate  $\Delta \epsilon$  by taking  $\epsilon_{\mathrm{LUMO}} - \epsilon_{\mathrm{HOMO}}$  and use the atomization energy for  $U_0, U, H$  and  $G$  on QM9. Experiments are conducted on Nvidia RTX 3090 and Nvidia RTX 4090. Results are average of three runs with different random seeds. Detailed training setting can be referred to Table 4.
Model hyperparameters. The key model hyperparameters we coarsely tune are the rbf dimension, hidden dimension, and number of message passing blocks. For rbf dimension, we use 16 for MD17 and 32 for QM9. We choose the number of message passing blocks from  $\{4,5,6\}$ . For hidden dimension, we use 512 for 2-DisGNNs and 320 for 3-DisGNNs. The detailed hyperparameters can be found in our codes.
D.2 Supplementary Experimental Information
Discussion of results in MD17. In our experiments on MD17, most of the state-of-the-art performance we achieve is on force targets, and the loss on energy targets is relatively higher, see Table 1. On force prediction tasks,  $k$ -DisGNNs rank top 2 on force prediction tasks on average, and outperform the best results by a significant margin on several molecules, such as aspirin and malonaldehyde. However, the results on the energy prediction tasks are relatively lower, with 2-F-DisGNN ranking  $2^{\text{nd}}$  and 3-E-DisGNN ranking  $4^{\text{th}}$ .
This is due to the fact that we have assigned a quite high weight (0.999) to the force loss during training, similar to what GemNet[Gasteiger et al., 2021] does. In comparison, TorchMD [Thölke and De Fabritiis, 2021] assigned a weight of 0.8 to the force loss, resulting in better results on energy targets, but not as good results on force targets.
Actually, in molecular simulations, force prediction is a more challenging task. It determines the accuracy of molecular simulations and reflects the performance of a model better [Gasteiger et al., 2021] (One can see that the energy performance of different models does not vary much). Therefore, it makes more sense to focus on force prediction. Furthermore, previous researches [Batzner et al., 2022, Christensen and Von Lilienfeld, 2020] have found that models can achieve significantly lower
energy loss on the revised MD17 dataset than on the original MD17 dataset, while holding similar force accuracy on the two datasets. As analyzed in Batzner et al. [2022], this suggests that the noise floor on the original MD17 dataset is higher on the energies, indicating that better force prediction results are more meaningful than energy prediction results on the original MD17 datasets.
Table 5: MAE loss on revised MD17. Energy (E) in kcal/mol, force (F) in kcal/mol/Å.
Target	MACE	Allegro	2F-Dis.
aspirin	E	0.0507	0.0530	0.0465
F	0.1522	0.1683	0.1515
azobenz.	E	0.0277	0.0277	0.0315
F	0.0692	0.0600	0.1121
benzene	E	0.0092	0.0069	0.0013
F	0.0069	0.0046	0.0085
ethanol	E	0.0092	0.0092	0.0065
F	0.0484	0.0484	0.0379
malonal.	E	0.0184	0.0138	0.0129
F	0.0945	0.0830	0.0782
napthal.	E	0.0115	0.0046	0.0103
F	0.0369	0.0208	0.0478
paracet.	E	0.0300	0.0346	0.0310
F	0.1107	0.1130	0.1178
salicyl.	E	0.0208	0.0208	0.0174
F	0.0715	0.0669	0.0860
toluene	E	0.0115	0.0092	0.0051
F	0.0346	0.0415	0.0284
uracil	E	0.0115	0.0138	0.0131
F	0.0484	0.0415	0.0828
Revised MD17. For a comprehensive comparison, we also conducted experiments on the revised MD17 dataset, which has higher data quality than original MD17, and compared with two state-of-the-art models, MACE [Batatia et al., 2022] and Allegro [Musaelian et al., 2023]. The results are shown in Table 5. 2-F-DisGNN achieved 10 best and 4 the second best results out of 20 targets. Note that MACE and Allegro are both models that leverage complex equivariant representations. This further demonstrates the high potential of distance-based models in geometry learning and molecular dynamic simulation.
Supplementary Results on QM9. We present the full results on QM9 in Table 6. We compare our model with 7 other models, including those that use invariant geometric features: SchNet [Schütt et al., 2018], DimeNet [Klicpera et al., 2020a], DimeNet++[Klicpera et al., 2020b], a model that uses group irreducible representations: Cormorant[Anderson et al., 2019], those that use first-order equivariant representations: TorchMD [Thölke and De Fabritiis, 2021], PaiNN [Schütt et al., 2021] and a model that uses local frame methods: GNN-LF [Wang and Zhang, 2022]. We calculate the average improvements of all the models relative to DimeNet and list them in the table.
Table 6: MAE loss on QM9.
Target	Unit	SchNet	Cormor.	DimeNet++	Torchmd	PaiNN	GNN-LF	DimeNet	2F-Dis.
μ	D	0.033	0.038	0.0297	0.002	0.012	0.013	0.0286	0.0100
α	a03	0.235	0.085	0.0435	0.01	0.045	0.0353	0.0469	0.0431
εHOMO	meV	41	34	24.6	21.2	27.6	23.5	27.8	21.81
εLUMO	meV	34	38	19.5	17.8	20.4	17	19.7	21.22
Δε	meV	63	61	32.6	38	45.7	37.1	34.8	31.3
(R2)	a02	0.073	0.961	0.331	0.015	0.066	0.037	0.331	0.0299
ZPVE	meV	1.7	2.027	1.21	2.12	1.28	1.19	1.29	1.26
U0	meV	14	22	6.32	6.24	5.85	5.3	8.02	7.33
U	meV	19	21	6.28	6.3	5.83	5.24	7.89	7.37
H	meV	14	21	6.53	6.48	5.98	5.48	8.11	7.36
G	meV	14	20	7.56	7.64	7.35	6.84	8.98	8.56
cv	cal/mol/K	0.033	0.026	0.023	0.026	0.024	0.022	0.0249	0.0233
Avg improv.	-78.99%	-98.22%	9.42%	25.00%	17.50%	27.82%	0.00%	18.83%
Rank	7	8	5	2	4	1	6	3
Although our models exhibit significant improvements on the MD17 datasets, their performance on the QM9 datasets is relatively less remarkable. This discrepancy may be attributed to the fact that the QM9 datasets pose a greater challenge in terms of a model's generalization performance: Unlike MD17, QM9 comprises various kinds of molecules and regression targets. However, since  $k$ -DisGNNs rely purely on the fundamental geometric feature of distance, they may inherently require a larger amount of data to learn the significant geometric patterns within geometric structures and enhance their generalization performance, compared to models that explicitly incorporate equivariant representations or pre-calculated high-order geometric features. Nevertheless, we believe that such learning paradigm, which employs highly expressive models to extract information from the most basic components in a data-driven manner, has immense performance potential and will exhibit even greater performance gains when provided with more data points, increased model scales, better data quality or advanced pre-training methods, as already validated in fields such as natural language processing [Brown et al., 2020, Devlin et al., 2018], computer vision [Tolstikhin et al., 2021, Liu et al., 2021, Dosovitskiy et al., 2020], and molecule-related pre-training methods [Xia et al., 2023, Lu et al., 2023].
D.3 Time and Memory Consumption
We performed experiments on the MD17 dataset to compare the training and inference time, as well as the GPU memory consumption of our models with their counterparts DimeNet and GemNet. The experiments are conducted on Nvidia RTX 3090, and the results are shown in Table 7.
Table 7: Training time, inference time and GPU memory consumption on MD17. The batch size is 32 for 2-F-DisGNN, DimeNet and TorchMD, and 12 for 3-E-DisGNN and GemNet. Training time in ms, inference time in ms, inference GPU memory consumption in MB. Evaluated on Nvidia A100.
2FDis.	DimeNet	TorchMD	3EDis.	GemNet
Asp.	234/66/2985	270/74/5691	101/39/2072	873/273/18585	556/182/7169
Ben.	92/24/1058	162/44/1815	110/41/905	224/61/3644	380/130/1789
Eth.	61/18/637	163/43/791	103/43/532	112/30/1561	462/158/804
Mal.	61/18/640	161/43/791	103/42/532	113/30/1561	356/118/804
Nap.	184/50/2193	238/68/4578	111/45/1692	652/185/11804	444/141/5253
Sal.	137/37/1758	206/58/3468	119/42/1401	458/127/8281	405/136/3778
Tol.	131/36/1560	200/54/3126	107/25/1326	399/109/6915	473/153/3328
Ura.	93/24/1058	169/45/1782	113/42/906	225/61/3644	369/136/1776
Avg.	124/34/1486	196/53/2755	108/39/1170	382/109/6999	430/144/3087
2-F-DisGNN demonstrates significantly better time and memory efficiency compared to DimeNet and GemNet. Even when compared to the quite efficient model, TorchMD, 2-F-DisGNN shows competitive efficiency. This is partly because it utilizes high-order tensors and employs dense calculation, resulting in significant acceleration on GPUs. However, due to the high theoretical complexity and high hidden dimension, 3-E-DisGNN shows relatively worse performance. Regarding this limitation, we have listed several possible solutions in Section 8 and recent work on simplifying and accelerating  $k$ -WL [Zhao et al., 2022, Morris et al., 2022] when  $k$  is large may also be applied, which is left for future work. Nonetheless, given that 2-F-DisGNN already satisfies the requirements for both theoretical (Section 6.3) and experimental (Section 7) performance, our primary focus lies on this model, and the analysis further emphasizes the practical potential of 2-F-DisGNN.
E Supplementary Related Work
There is a wealth of research in the field of interatomic potentials, which introduced various methods for modeling the atomic environment as representations [Bartók et al., 2010, 2013, Shapeev, 2016, Behler and Parrinello, 2007]. In particular, ACE (Atomic Cluster Expansion) [Drautz, 2019] has presented a framework that serves as a unifying approach for these methods and can calculate high-order complete polynomial basis features with small cost regardless of the body order. Building on Drautz [2019], Batatia et al. [2022] proposed a GNN variant called MACE that can effectively exploit the rich geometric information in atoms' local environment. Additionally, Joshi et al. [2023] proposed
geometric variants of the WL test to characterize the expressiveness of invariant and equivariant geometric GNNs for the general graph setting. All of Drautz [2019], Batatia et al. [2022], Joshi et al. [2023] leverage the many-body expansion and can include  $k$ -tuples during embedding atom's neighbors, allowing atoms to obtain  $k$ -order geometric information from their local environments like  $k$ -DisGNNs.
However, we note that there are key differences between these works (based on many-body expansion) and  $k$ -DisGNNs in terms of their message passing units and use of representations. The former assumes that the total energy can be approximated by the sum of energies of atomic environments of individual atoms and therefore use atoms as their message passing units, while  $k$ -DisGNNs pass messages among  $k$ -tuples and pool tuple representations to obtain global representation. In addition, while interatomic potential methods leverage equivariant representations to enhance their expressiveness,  $k$ -DisGNNs completely decouple E(3)-symmetry and only use the most basic geometric invariant feature, distance, to achieve universality and good experimental results.