SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Decoder. The framework utilizes decoder variations to accommodate various geometric requirements of a given task. By taking the latent representation pair  $(s_e, V_e)$  from the output of the encoder, the decoder can effectively leverage these representations to maximize task performance. For invariant targets where the tasks do not depend on global rotations, transformations, and permutation, the decoder takes encoded  $s_e$  representations and processes them with MLP layers. For equivariant targets where tasks are sensitive to global rotations, the decoder takes the scalar and vector representations and processes them with stacked AW blocks. The final vector representations are scaled with scalar representations to generate the predictions. If the task is graph level, final representations are aggregated with global sum pooling.
Equivalence and Invariance. Addressing the strategic initialization of the model and the design of operations that uphold invariant and equivariant properties is equally important. SAVENET uphold such properties, which can be stated more formally as follows:
Theorem 2. The equivariant representation of SAVENET is equivariant to any given rotation matrix in  $R \in SO(3)$ . The invariant representation of SAVENET is invariant to any given transformation matrix in  $SE(3)$ .
The formal proof of Theorem 2 is provided in Appendix C, demonstrating the validity of these assertions and the robustness of the model under various transformations.
This section evaluates our proposed SAVENET on three tasks in geometric representation learning over both synthetic and real-world datasets. We implement the model with PyTorch [26]. All experiments are conducted on an NVIDIA 3090 GPU with 24 GB memory. All models use the AdamW optimizer [22] for the optimization. Detailed descriptions of the datasets are compared in Table 1, which includes the average number of nodes, the average number of edges, and the total number
Table 1: Dataset details. The number of edges are computed with radius graph where  $\mathrm{d}\left( \overset{ \circ  }{\mathrm{A}}\right)  = {5.0}$
Datasets	QM9	N-Body	Molecule3D
Avg. # of Nodes	18.02	20	29.11
Max # of Nodes	29	20	137
Avg. # of Edges	280.72	380	553.00
# of Graphs	130,831	7,000	3,899,647
# of Tasks	12	4	5
Splits	84:8:8	42:29:29	6:2:2
Invariant	✓	×	✓
Equivariant	×	✓	✓
Table 2: Performance comparisons on QM9. † denotes using different data partitions.
Task	α	Δε	εHOMO	εLUMO	μ	Ccalmol K	G	H	R2a0	U	U0	ZPVE	std. %	log -
Units	a03	meV	meV	meV	D	meV	meV	meV	meV	meV	meV	meV	-
NMP†	.092	69	43	38	.030	.040	19	17	.180	20	20	1.50	1.78	-5.08
SchNet	.235	63	41	34	.0330	.0330	14	14	.073	19	14	1.70	1.76	-5.17
Cormorant†	.085	61	34	38	.038	.026	20	21	.961	21	22	2.03	2.14	-4.75
LieConv†	.084	49	30	25	.032	.038	22	24	.800	19	19	2.28	1.35	-4.99
PhysNet	.061	42.5	32.9	24.7	.0529	.0280	9.4	8.42	.765	8.34	8.15	1.39	1.37	-5.35
DimeNet	.047	34.8	27.8	19.7	.029	.0249	8.98	8.11	.331	7.89	8.02	1.29	1.05	-5.57
DimeNet++	.044	32.6	24.6	19.5	.0297	.0230	7.56	6.53	.331	6.28	6.32	1.21	0.98	-5.67
TFN†	.223	58	40	38	.064	.101	-	-	-	-	-	-	-	-
SE(3)-Tr†	.142	53	35	33	.051	.054	-	-	-	-	-	-	-	-
EGNN†	.071	48	29	25	.029	.031	12	12	.106	12	11	1.55	1.23	-5.43
PaiNN	.045	45.7	27.6	20.4	.0120	.024	7.35	5.98	.066	5.83	5.85	1.28	1.01	-5.85
ET	.059	36.1	20.3	17.5	.011	.026	7.62	6.16	.033	6.38	6.15	1.84	0.84	-5.90
SphereNet	.046	31.1	22.8	18.9	.0245	.0215	7.78	6.33	.268	6.36	6.26	1.12	0.91	-5.73
ComENet	.045	32.4	23.1	19.8	.0245	.0220	7.98	6.86	.259	6.82	6.69	1.20	0.93	-5.69
SEGNN†	.060	42	24	21	.023	.031	15	16	.660	13	15	1.62	1.08	-5.27
EQGAT	.053	32	20	16	.011	.024	23	24	.382	25	25	2.00	0.86	-5.28
Equiformer	.046	30	15.4	14.7	.0117	.0230	7.63	6.63	.251	6.74	6.59	1.26	0.70	-5.82
SAVENET-B	.039	24.8	18.4	16.3	.0093	.0227	6.64	5.43	.058	5.48	5.43	1.18	0.69	-6.04
SAVENET-L	.035	22.7	16.6	15.1	.0085	.0210	6.10	4.83	.049	4.74	4.83	1.10	0.63	-6.14
of graphs. The three datasets include both invariant and equivariant targets for comprehensive evaluation of our proposed model. Further dataset details are provided in Appendix D. The rationale behind our design decisions and hyperparameters, as well as their specific values, are documented in Appendix F for reference.
QM9: QM9 dataset was used to evaluate the performance and efficiency of the models across twelve tasks for invariant target predictions. We present two distinct configurations of the SAVENET model, distinguished by their respective sizes: The SAVENET-B denotes the base variant, while the SAVENET-L is the larger model. We report various baseline models, including NMP [10], Schnet [28], Cormorant [1], LieConv [7], PhysNet [32], Dimenet [16], Dimenet++ [9], TFN [31], SE(3) Transformer [8], EGNN [27], PaiNN [29], ET [30], SphereNet [21], ComENet [33], SEGNN [3], EQGAT [18], and Equiformer [19].
N-Body: To assess the effectiveness of SAVENET with equivariant tasks, we conducted experiments on the extended N-Body dataset with harder targets [5]. The proposed SAVENET is compared with GNN, Tensor Field Networks [31], SE(3) Transformer [8], Radial Fields, PaiNN [29], EGNN [27], ClofNet [5], GCPNet [25]. Most of the baseline results are adopted from [5, 25]. The results for PaiNN are not reported in previous work [25] due to numerical stability. Notably, we incorporated vector activation functions into PaiNN's interaction layers to address this issue, allowing us to include PaiNN in the performance comparison.
Molecule3D: The Molecule3D dataset, possessing over  $29 \times$  the quantity of graphs as the QM9 dataset and an approximate  $1.6 \times$  and  $1.9 \times$  increase in the average number of nodes and edges respectively, is employed to evaluate the scalability of the model with respect to dataset size. For the Molecule3D dataset, we adopted the baselines as per [33]. As only the gap target has been reported in preceding works, we utilized the hyperparameters provided in [33] for evaluating additional targets. The details of the hyperparameters are described in Appendix F.1 The comparative baseline models for the Molecule3D include GIN-Virtual [33], Dimenet++ [9], Spherenet [21], and ComENet [33].
Evaluation Metrics. A set of metrics are used to measure the performance of the models. For the QM9 and Molecule3D datasets, each task is assessed using the Mean Absolute Error (MAE). For the N-Body dataset, the performance is measured by the Mean Squared Error (MSE) of the future position predictions. Given that each dataset encompasses multiple tasks, we employ two additional aggregate measures: standardized (std.) and logarithmic (log) metrics, to assess overall performance. As the error can be disproportionately influenced by a few outliers, we report the logarithmic error to prevent these outliers from dominating the evaluation.
Efficiency Analysis. To critically assess the efficiency of the models, we employ a set of metrics including training speed, model complexity, and memory consumption. All the measurements are conducted utilizing hyperparameters from the respective model authors, and the reference imple
Table 3: Efficiency comparison of the proposed SAVENET with the baselines on QM9.
Model Metric	std. MAE ↓	log MAE ↓	Batch Size↑	Memory GB ↓	Latency ms ↓	Samples/s ↑	FLOPs G ↓	Param. M ↓	MACs G ↓
DimeNet++	0.98	-5.67	357	4.82	38.9	1646	30.15	1.89	15.08
ComENet	0.93	-5.69	1174	0.96	26.4	2422	14.36	3.81	7.18
SphereNet	0.91	-5.73	238	6.65	51.6	1240	31.80	1.89	15.90
EQGAT	0.86	-5.28	980	1.74	53.3	1200	17.00	0.93	8.49
Equiformer	0.70	-5.83	94	14.04	337.3	190	-	3.53	-
SAVENET-B	0.69	-6.04	1660	0.85	23.1	2767	11.60	1.37	5.76
SAVENET-L	0.63	-6.14	580	2.48	50.6	1237	55.36	7.72	27.51
mentations are based on [18-20, 33]. To ensure consistent comparison, the experimental runs are executed on the same configurations utilizing an RTX 3090 GPU with 24 GB of graphics memory. The details of the efficiency analysis procedure are described in Appendix G and the theoretical time complexity is analyzed in Appendix E.
- Training Speed. To gauge the training speed of each model, we recorded the time consumed by a single forward/backward pass, expressed in milliseconds. Alongside this latency, we computed the number of samples that could be processed per second, employing similar measurement techniques.
- Model Complexity. The complexity of the models was evaluated by considering two key parameters: floating point operations (FLOPs) and multiply-accumulate operations (MACs). These metrics were computed utilizing the DeepSpeed library [23], with a consistent batch size.
- Memory Consumption. The memory consumption of each model was evaluated by assessing the amount of memory utilized across multiple batches. After processing each batch, the cache was cleared to ensure a realistic appraisal of memory usage. However, this practice of cache clearing significantly impedes the evaluation process, rendering a full dataset evaluation impractical.
Accuracy results on QM9. Table 2 listed the experimental results of the QM9 dataset. SAVENET-B demonstrated superior performance on seven out of the twelve targets, surpassing all baseline models on these measures. While SAVENET-B did not yield the highest performance for certain parameters—namely,  $\epsilon_{\mathrm{HOMO}}$ ,  $\epsilon_{\mathrm{LUMO}}$ ,  $C_\nu$ ,  $R^2$ , and ZPVE—it is noteworthy that it still delivered competitive results, indicating its robustness across diverse molecular properties. SAVENET-L further improved on SAVENET-B's performance, securing the best results on nine out of the twelve targets, with the exceptions of  $\epsilon_{\mathrm{HOMO}}$ ,  $\epsilon_{\mathrm{LUMO}}$ , and  $R^2$ . Furthermore, in terms of the overall performance measures, SAVENET-L outperformed all other models, exhibiting the lowest standardized MAE and logarithmic MAE, reflecting its superior accuracy across all target properties. These results highlight the effectiveness of our proposed models, particularly SAVENET-L, in accurately predicting a wide range of molecular properties in the QM9 dataset.
Table 4: Performance comparisons on N-Body dataset.
Method	ES(5)	ES(20)	G+ES(20)	L+ES(20)	std.MSE	logMSE
GNN	0.0131	0.0720	0.0721	0.0908	0.0151	-4.3876
TFN	0.0236	0.0794	0.0845	0.1243	0.0192	-4.0978
SE(3)-Tr.	0.0329	0.1349	0.1000	0.1438	0.0252	-3.8037
Radial Field	0.0207	0.0377	0.0399	0.0779	0.0110	-4.6212
PaiNN	0.0158	0.09971	0.10291	0.13561	0.0835	-2.5329
ET	0.1653	0.1788	0.2122	0.2989	0.0535	-2.9587
EGNN	0.0079	0.0128	0.0118	0.0368	0.0044	-5.6241
ClofNet	0.0065	0.0073	0.0072	0.0251	0.0029	-6.0324
GCPNET	0.0070	0.0071	0.0073	0.0173	0.0024	-6.1104
SAVENET	0.0062	0.0063	0.0082	0.0123	0.0020	-6.2268
Efficiency results on QM9. The results of efficiency experiments are presented in Table 3. We select the top five models based on their std. MAE performance for the efficiency comparison. Regarding efficiency, SAVENET-B stands out with its superior results, processing the highest number of samples per second and the lowest memory consumption. This demonstrates an impressive balance between performance and resource utilization. Additionally, SAVENET-B shows the lowest latency and number of FLOPs and MACs. The latency of SAVENET-B is  $14.5 \times$  lower compared to Equiformer, the previously most performant model with the lowest std. and log MAE. Even with a significant reduction in compute time, SAVENET-B surpasses the Equiformer in performance, illustrating efficiency advancements of SAVENET without performance compromise.
SAVENET-L manages to maintain a competitive efficiency despite its larger size. Its memory consumption is  $2.6 \times$  and  $5.6 \times$  times smaller compared to high-performing models such as SphereNet and Equiformer, respectively. Notably, SAVENET-L processes 1237 samples per second, which is comparable to the rates of SphereNet and EqGAT. Compared to Equiformer, the baseline model with the highest accuracy, SAVENET-L exhibits approximately  $6 \times$  lower latency. This suggests that SAVENET-L, even with its larger size, can maintain high performance without sacrificing speed.
Results on N-Body dataset. Table 4 presents the results of the N-Body systems. SAVENET outperformed all baseline methods on the ES(5), ES(20), and L+ES(20) tasks, whilst demonstrating competitive performance on the G+ES(20) task. Remarkably, SAVENET achieved the lowest average MSE and logarithmic MSE across all four tasks, illustrating its superior prediction precision for equivariant targets. These results emphasize the capability of our proposed model to accurately predict the future positions of particles in N-Body systems, even under challenging conditions where particles are under the influence of the gravity and force fields.
Table 5: Performance comparisons on Molecule3D dataset.
Symmetry Task	μ	Invariant	Equivariant μ
εHOMO	εLUMO	Δε	std. MAE	logMAE
GIN-Virtual	0.0882	0.0692	0.0632	0.1036	0.0592	-2.8677	1.5233
SchNet	0.0532	0.0275	0.0265	0.0428	0.0263	-3.6633	1.2082
DimeNet++	0.0293	0.0240	0.0190	0.0306	0.0188	-4.0139	1.2014
SphereNet	0.0288	0.0239	0.0183	0.0301	0.0184	-4.0327	1.1836
ComENet	0.0345	0.0288	0.0252	0.0326	0.0220	-3.8403	1.3521
SAVENET-B	0.0183	0.0190	0.0173	0.0290	0.0156	-4.2257	0.0108
SAVENET-L	0.0136	0.0159	0.0143	0.0239	0.0128	-4.4408	0.0090
Accuracy results on Molecule3D. Our proposed models, SAVENET-B and SAVENET-L, outperformed several baseline models on the Molecule3D dataset. In terms of relative performance, SAVENET-B outperforms all tasks compared to other established models. SAVENET-B showed improved performance by approximately  $37\%$  on the  $\mu$  task and  $92\%$  on the  $\vec{\mu}$  task compared to the next best model, SphereNet. Our SAVENET-L model, further improved the performance of SAVENET-B. For the  $\mu$  task, SAVENET-L showed an improvement of approximately  $26\%$  over SAVENET-B; for the  $\vec{\mu}$  task, it showed an improvement of about  $17\%$ . Regarding aggregate measures, SAVENET-L also outperformed all other models, showing a relative improvement of around  $18\%$  in the std. MAE metric and  $5\%$  in the log MAE metric compared to SAVENET-B. These results demonstrate the superior performance and scalability of our proposed models, especially SAVENET-L, on the Molecule3D dataset.
Efficiency results on Molecule3D. The results of the efficiency evaluations on the Molecule3D dataset are illustrated in Table 6. Impressively, SAVENET-B continues to dominate with respect to efficiency, manifesting the highest sample processing rate and minimal memory consumption. This shows an exemplary trade-off between accuracy and resource management.
<sup>1</sup>The results are reported after adding vector activation functions. These results were initially unavailable due to numerical instability.
Table 6: Efficiency comparison of the proposed SAVENET with the baselines on Molecule3D.
Model
Metric	std. MAE ↓	log MAE ↓	Batch Size↑	Memory GB ↓	Latency ms ↓	Samples/s ↑	FLOPs G ↓	Param. M ↓	MACs G ↓
DimeNetPP	0.0188	-4.0139	62	12.14	84.2	704	70.53	1.89	35.27
SphereNet	0.0184	-4.0327	43	17.32	134.6	448	92.56	1.90	46.28
ComENet	0.0220	-3.8403	277	3.37	40.7	1536	47.25	7.36	23.63
SAVENET-B	0.0156	-4.2257	329	2.84	27.8	2304	15.33	1.05	7.63
SAVENET-L	0.0128	-4.4408	159	5.37	51.9	1216	38.93	3.01	19.37
Table 7 presents the results of our ablation study on the proposed SAVENET model, where we investigate the impacts of Direction Noise (DN), Vector Activations (VA), and varying the number of layers (L) on the QM9 dataset. In Table 7, we report standard Mean Absolute Error (std. MAE), the percentage change in std. MAE, and the logarithmic Mean Absolute
Table 7: Ablation study of SAVENET
Model	L	std.	Δstd.	log
SAVENET - DN - VA	6	0.79	-	-5.85
SAVENET - VA	6	0.77	-2.5%	-5.95
SAVENET - VA	8	0.71	-10.1%	-5.98
SAVENET	8	0.69	-12.7%	-6.04
SAVENET	12	0.63	-20.3%	-6.14
Error (logMAE), respectively. The importance of DN is illustrated in the first two rows, where the std. MAE is reduced by  $2.5\%$ . Moreover, scaling the model with DA from six to eight layers yields a  $7.8\%$  reduction in error. Introducing VA into an eight-layer model further improves the performance by  $2.8\%$ . Lastly, expanding the model to 12 layers offers a substantial performance improvement of  $8.7\%$ .
6 Conclusion, Limitations and Future Works
This paper has focused on learning geometric representations of molecular structures. We introduced SaVeNet, a scalable message-passing neural network designed for learning geometric representations from 3D graphs. Furthermore, we implemented a unique approach to integrate geometric features, capturing complex geometric relationships within 3D structures without resorting to computationally complex operations. The performance of this framework was demonstrated on both synthetic and real-world datasets. For the latter, SaVeNet was evaluated across various core tasks in molecular representation learning. When compared with state-of-the-art methods, the effectiveness and efficiency of our proposed architecture were clearly validated. Our comprehensive examination of network efficiency and our comprehension of the trade-offs between expressiveness and efficiency will provide a foundation for future advancements in molecular representation learning. However, like many studies in this area, our work does have limitations. The primary constraint, shared with the baseline model discussed in this paper, is the dependency on known geometric structures of molecules at their equilibrium state. In practice, obtaining accurate 3D confirmations can be costly. An intriguing direction for future research is to explore an end-to-end framework that employs sequential representations of molecules, such as SMILES strings [34], to generate their 3D geometry, where the predicted 3D geometry could then be used to model molecular properties. Another promising avenue is to investigate the performance of larger variations of SaVeNet. Furthermore, assessing the impacts of efficiency improvements on larger models, potentially achieved through shared parameters in various configurations, is a worthwhile endeavor. We anticipate that our findings will pave the way for future investigations in this domain, potentially leading to even more efficient and effective modeling techniques.
[32] Oliver T Unke and Markus Meuwly. Physnet: A neural network for predicting energies, forces, dipole moments, and partial charges. Journal of chemical theory and computation, 15(6): 3678-3693, 2019.
[33] Limei Wang, Yi Liu, Yuchao Lin, Haoran Liu, and Shuiwang Ji. ComENet: Towards complete and efficient message passing for 3d molecular graphs. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=mCzMqewSFJ.
[34] David Weininger. Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules. Journal of Chemical Information and Computer Sciences, 28(1):31-36, 1988. doi: 10.1021/ci00057a005. URL https://doi.org/10.1021/ci00057a005.
[35] Tian Xia and Wei-Shinn Ku. Geometric graph representation learning on protein structure prediction. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, KDD '21, page 1873-1883, New York, NY, USA, 2021. Association for Computing Machinery. ISBN 9781450383325. doi: 10.1145/3447548.3467323. URL https://doi.org/10.1145/3447548.3467323.
[36] Tian Xia, Sarp Aykent, and Wei-Shinn Ku. Extrinsic-intrinsic representation learning framework for drug discovery. In Proceedings of the 2023 SIAM International Conference on Data Mining (SDM), pages 523-531. SIAM, 2023.
[37] Zhao Xu, Youzhi Luo, Xuan Zhang, Xinyi Xu, Yaochen Xie, Meng Liu, Kaleb Dickerson, Cheng Deng, Maho Nakata, and Shuiwang Ji. Molecule3d: A benchmark for predicting 3d geometries from molecular graphs. arXiv preprint arXiv:2110.01717, 2021.
Lemma 1: Consider a known geometric graph with at least one node, and assume that each node in this graph has at least one connection. Let a new node  $j$  be added to this graph, stipulating that  $j$  is connected to at least one existing node. Then, the position  $p_j$  of node  $j$  can be determined in constant time using the  $(p_0, r_{ij}, d_{\vec{j}})$  properties.
Proof of Lemma 1: In our context, we are dealing with a geometric graph whose node properties are known a priori. When a new node  $j$  is appended to the graph, such that it is connected to at least one node, we can select any node  $i \in \mathcal{N}_j$  from the neighbors of  $j$  to establish the position  $c_j$  of node  $j$ .
Proceeding by contradiction, assume that multiple feasible positions  $p_j$  could satisfy the conditions  $||c_i - c_j|| = r_{ij}$  and  $\vec{d}_{ij}$  for a given position  $p_i$ . In multi-dimensional spaces, including 3D, there exist multiple positions satisfying the distance constraint  $||c_i - c_j|| = r_{ij}$ . Specifically, any point lying on a sphere in 3D space, with  $c_i$  as the center and  $r_{ij}$  as the radius, could satisfy this condition. Yet, there is only one such position that can satisfy the direction vector  $\vec{d}_{ij}$  from  $c_i$  to  $c_j$ . This fact contradicts our assumption, and hence concludes the proof.
Theorem 1: Assuming the input graph is strongly connected, the input space of SAVENET, defined as  $(r,d)$ , is capable of reconstructing the rotation-equivariant geometric structure of the input graph. The reconstructed graph maintains the relative spatial configuration of nodes and edges in the input graph, the translation transformation is required for exact alignment with the original graph.
Proof of Theorem 1: By the definition of a strongly connected graph, for any given node  $i$ , there exists a path  $\Pi_i^j$  to every other node  $j$ . Hence, we can arbitrarily choose any node from the graph as a reference for reconstructing the rotation-equivariant geometric structure of the input graph. Let's select node  $i$  for this purpose. Since  $c_i$  is arbitrary, we can conveniently set it at the origin without loss of generality.
Utilizing Lemma 1, we can identify the relative position of every other node  $j \in V$  by traversing each edge on the path  $\Pi_i^j$  from  $i$  to  $j$ . In essence, the relative position of a node  $j$  with respect to node  $i$  can be computed by the accumulation of relative positions along the path  $\Pi_i^j$ .
Assuming, by way of contradiction, that multiple relative positions of node  $j$  with respect to node  $i$  exist. However, this is counter to Lemma 1, which affirms that for given  $c_{i}, r_{ij}, \vec{d}_{ij}$ , there exists a unique location. This principle can be extrapolated to any number of nodes in the graph. Further, as we can ascertain the position of a node after it's added to the reference graph, we can similarly determine the position of any subsequent node  $k$  relative to  $V \cup \{j\}$ .
Thus, we conclude that the input graph's rotation-equivariant geometric structure can be reconstructed using  $(d,r)$  properties, thereby completing our proof.
Theorem 2: The equivariant representation of SAVENET is equivariant to any given rotation matrix in  $R \in SO(3)$ . The invariant representation of SAVENET is invariant to any given transformation matrix in SE(3).
Proof of Theorem 2: Our proof unfolds in two stages: first, we discuss the symmetrical properties of SAVENET in the input space. Second, we extend these properties to the entire network.
Inputs. The input to SAVENET consists of direction vectors, denoted  $\vec{\beta}$ . It's clear that these vectors, representing direction, maintain their properties under any global rotation. Formally, if  $f_{b}(\mathcal{G})$  represents the function that extracts  $\vec{\beta}$  from an input graph  $\mathcal{G}$ , we have that  $Rf_{b}(\mathcal{G}) = f_{b}(R\mathcal{G})$  for all
$R \in SO(3)$ . Moreover, any noise introduced during initialization decays over time until reaches zero before the end of the training phase. Therefore, having no effect on the inference process.