SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The invariant input vectors of SAVENET, which represent atomic numbers and interatomic distances, are clearly invariant under any transformation in  $SE(3)$  because these properties do not depend on the geometrical information of the system. Hence, atomic numbers remain constant under global rotations and translations. Similarly, the distances between atoms, being local properties, remain unchanged under global transformations.
Network Propagation: Having established the symmetries at the input level, we extend our analysis to the entire network of SAVENET. Two critical components are Inter-atomic Interactions and Atom-wise Blocks, which we shall now scrutinize.
Inter-atomic interactions. The interaction function  $e$  can be dissected to analyze its operations on the equivariant and invariant branches. Here,  $\phi_b$ ,  $\phi_d$ , and  $\phi_v$  denote the vector update layer, the linear transformation, and the multiplication of  $v$  representation with the linear transformation of  $s_j$ , respectively. In the equivariant branch, operations can be expressed as:
$$
e _ {v} \left(s _ {j}, V _ {j}, r _ {i j}, \vec {\beta} _ {i j}\right) = \phi_ {b} \left(\vec {\beta} _ {i j}\right) \odot \phi_ {d} \left(s _ {j}\right) \odot \eta_ {d} \left(r _ {i j}\right) + V _ {j} \odot \phi_ {v} \left(s _ {j}\right) \odot \eta_ {r} \left(r _ {i j}\right) \tag {7}
$$
Each component in the above expression preserves symmetry properties. More specifically, the first term  $\phi_b(\vec{\beta}_{ij})$  preserves equivariance because the same transformation is applied across spatial dimensions. The second term,  $\phi_d(s_j)$ , is a linear transformation applied to an invariant, and therefore, remains invariant under any transformation in  $T$ . The third term,  $\eta_d(r_{ij})$ , which multiplies a basis function and a cutoff function for the distance  $r_{ij}$ , does not change under transformations in  $T$ . Thus, it is also invariant under  $T$ . The fourth term,  $\phi_v(s_j)$ , involves a multiplication operation of  $v$  representation with a linear transformation of  $s_j$ , hence, it produces an equivariant vector since the linear transformation does not affect the spatial dimensions. Finally, the term  $\eta_r(r_{ij})$  represents an invariant scalar, remaining unchanged under transformation.
The first part  $\phi_{b}(\vec{\beta}_{ij})\odot \phi_{d}(s_{j})\odot \eta_{d}(r_{ij})$  involves the Hadamard product of an equivariant vector  $(\phi_b(\vec{\beta}_{ij})\in \mathbb{R}^{C\times 3})$  with invariant scalar values  $(\phi_d(s_j),\eta_d(r_{ij})\in \mathbb{R}^C)$ . Hadamard product operates element-wise and does not alter the spatial dimensions, hence the resulting vector preserves the equivariant properties. The second part  $V_{j}\odot \phi_{v}(s_{j})\odot \eta_{r}(r_{ij})$  involves the Hadamard product of an equivariant vector  $(V_{j}\in \mathbb{R}^{C\times 3})$  with an invariant scalar representations. Again, the resulting vector is equivariant since the same values are multiplied across spatial dimensions. Finally, the addition operation in  $e_v$  combines two equivariant vectors. As the addition of equivariant vectors preserves equivariance, the result will be equivariant under any rotation matrix  $R\in SO(3)$ . This concludes our examination of the equivariance of the interaction function  $e_v$  under any rotation  $R\in SO(3)$ . Similarly, the operations in the invariant update branch can be captured as:
$$
e _ {s} \left(s _ {j}, v _ {j}, r _ {i j}, \vec {\beta} _ {i j}\right) = \phi_ {s} \left(s _ {j}\right) \odot \eta_ {s} \left(r _ {i j}\right) \tag {8}
$$
where  $\phi_s$  denotes a linear transformation, and  $\eta_s$  represents a transformation applied using basis and cutoff functions. Both these functions are invariant under any transformation in  $T$ .
The updated  $s$  and  $V$  representations, determined for each adjacent node, are subsequently aggregated through summation. Given that summation inherently upholds symmetry characteristics, the procedures employed in the Inter-atomic interactions layer maintain the equivariant nature of  $V$  for all rotations  $R$  in  $SO(3)$ , as well as the invariant nature of  $s$  under any transformation  $T$  in  $SE(3)$ .
Atom-wise Blocks. Assuming that the inputs to this block are invariant for scalar representations and equivariant for vector representations, we can extend our proof to the operations within this block. The operations in the invariant update branch can be encapsulated as follows:
$$
\phi_ {m} \left(\left| \left| \phi_ {v u} (V) \right| \right| \cup s\right) \tag {9}
$$
In this equation,  $\phi_{vu}$  denotes the vector update function which is, as previously established, equivariant under any rotation matrix  $R$ . The norm operation  $\|\cdot\|$  applied to  $V$  is invariant, as the magnitude of the vector does not change with rotation. Consequently, the composition  $\|\phi_{vu}(V)\|$  is invariant to global rotations. Furthermore, applying the linear transformation  $\phi_m$  to the concatenation of two invariants, in this case  $\|\phi_{vu}(V)\|$  and  $s$ , remains invariant under any transformation in  $SE(3)$ . The operations in the equivariant branch of the Atom-wise blocks can be described as:
$$
\mathrm {V A} \left(\phi_ {v u} (V) \odot \phi_ {v} \left(\left| \left| \phi_ {v u} (V) \right| \right| \cup s _ {i}\right)\right) \tag {10}
$$
This equation is similar to the one describing the invariant update branch, with a difference lying in the linear transformation  $\phi_v$  that uses different learnable parameters. Hence, the second term is invariant. The first term,  $\phi_{vu}(V)$ , is equivariant as it is a vector update function applied to  $V$ . The Hadamard product of this equivariant representation with the invariant term preserves its equivariance under any rotation  $R \in SO(3)$ . In conclusion, the Atom-wise blocks preserve the symmetrical properties for both  $(s,V)$  representations. The proof for Theorem 2 is thus complete.
The three datasets used in this study evaluate our proposed SAVENET on three tasks in geometric representation learning over both synthetic and real-world datasets. In addition to their diversity in their real-world use cases, these tasks demonstrate a rich diversity not just in their practical implications, but also in their requisite output specifications. In other words, they demand the prediction of both invariant and equivariant targets, each necessitating different computational considerations. This multi-faceted evaluation strategy, which includes varying input data types and output requirements, provides a comprehensive view of the model's capabilities, and furthers the understanding of its strengths and potential areas for refinement. By analysing the model's performance across these distinct scenarios, we aim to exhibit SAVENET's adaptability and potential as a valuable tool in advancing the field of geometric representation learning.
QM9. Widely adopted small molecule property prediction dataset. It consists of 131,831 stable molecules with nine heavy atoms. We utilize the established split from [16], a standard in subsequent works [9, 21, 33]. QM9 contains twelve tasks that involve various aspects of molecular properties. The multi-task nature of the dataset demonstrates the effectiveness of the evaluated models.
Newtonian many-body system. Newtonian many-body system is a synthetic dataset that is generated based on particle simulations. The task is to predict the future positions of the particles based on the initial location and velocity. The performance was measured with the mean square error (MSE) metric between predicted and ground truth. Recently the existing work by [27] extended by increasing the particle size and introducing various force fields. By introducing the force fields, the existing task became more difficult. The initial work evaluated the models' performance based on the future position on 5-body (ES(5)). More challenging tasks to predict future locations of the 20-body (ES(20)) task is introduced in [5]. Furthermore, complex conditions integrated with gravity field (G+ES(20)) and Lorentz-like force field  $(\mathrm{L} + \mathrm{ES}(20))$ . The performances of the models were evaluated under challenging conditions. The hyper-parameters of our work SAVENET were chosen based on the best-performing hyper-parameters on ES(5) based on the best validation score.
Molecule3D. This dataset comprises 3,899,647 molecules, making it approximately 29.5 times larger than the QM9 dataset, providing a significant platform to assess the scalability of our model. Predefined splits are utilized, based on random sampling as detailed in [37]. Though Molecule3D has fewer targets than QM9, with only four targets including dipole moment vector  $(\vec{\mu})$ , energies of the highest occupied and lowest unoccupied molecular orbitals  $(\epsilon_{\mathrm{HOMO}}, \epsilon_{\mathrm{LUMO}})$ , and orbital energy gap  $(\Delta \epsilon)$ , it offers an essential testbed for model evaluation. The dipole magnitude  $(\mu)$  can be easily derived from  $\vec{\mu}$ , highlighting the significance of equivariance in models due to the influence of global rotations on  $\vec{\mu}$ . For a fair comparison, we have reported standardized MAE for invariant properties.
E Theoretical complexity analysis
In the following section, we present a detailed theoretical complexity analysis of our proposed model and baselines, which is summarized in a comprehensive table for understanding the model's computational efficiency. In the table, we provide a breakdown of the computational com
plexities associated with the main components of the model's architecture. This includes complexities during the graph embedding phase, the encoder, and the decoder phase. By presenting this theoretical complexity analysis, we aim to provide an in-depth understanding of the operational demands of our model and highlight its computational advantages.
Table 8: Complexity analysis.
Complexity	Dimenet++	SphereNet	ComENet	SAVENET
Initialization	n3	n4	n2	n2
Encoder	le n3d2	le n3d2	le n2d2	le n2d2
Decoder	ldnd2	ldnd2	ldnd2	ldnd2
Overall	n3d2	n3d2	n2d2	n2d2
F Implementation details
The reference implementation can be found in https://github.com/EfficientGraphs/SaVeNet.
The hyperparameter search space of SAVENET-B on QM9 and Molecule3D is shown in Table 9. The final hyperparameters are chosen based on the hyperparameters with the best performance on the validation set of each dataset. The hyperparameters of SAVENET-L is determined based on the best hyperparameters found for SAVENET-B on the validation set. Then the hidden dimensions and the number of layers are scaled on SAVENET-L. On the QM9 dataset, SAVENET-L has two times more hidden dimensions and decoder hidden dimensions compared to SAVENET-B. SAVENET-L utilized 256 hidden dimensions compared to 128 dimensions, and SAVENET-L has 1.5 more layers increasing from eight layers to twelve layers. On the other hand for the Molecule3D dataset, hidden dimensions are increased 1.5 times from 128 to 192 and the number of layers increased to eight layers from six layers. The hyperparameter search space of SAVENET on N-Body is shown in Table 10.
Table 9: Hyperparameter search space of SAVENET-B on QM9 and Molecule3D.
Hyperparameters	QM9	Molecule3D
Epochs	1000	300
Batch size	32, 64, 128	64, 128, 256
Number of layers	6, 8	4, 6
Hidden dimensions	64, 128	128, 192
Decoder dimensions	128, 256	256, 384
Number of RBF	12, 16, 24	12, 16, 24
Base Learning Rate	1e-3, 5e-4, 2e-4	1e-3, 5e-4, 2e-4
Learning rate decay	0.6, 0.8, 0.9	0.6, 0.8, 0.9
Learning rate patience	10, 15, 20	10, 15, 20
Vector dropout	0.0, 0.1, 0.25	0.0, 0.1, 0.25
Direction noise decay	150, 250	150, 250
Table 10: Hyperparameter search space of SAVENET-B on N-Body.
Hyperparameters	N-Body
Epochs	500
Number of layers	4
Hidden dimensions	128
Decoder dimensions	128, 256
Number of RBF	12, 16, 24
Base Learning Rate	1e-3, 5e-4, 2e-4
Learning rate decay	0.6, 0.8, 0.9
Learning rate patience	10, 15, 20
Vector dropout	0.0, 0.1, 0.25
Direction noise decay	150, 250
G Efficiency evaluation
In this section, we conduct a rigorous evaluation of our proposed model, employing a set of carefully devised efficiency metrics. These metrics are instrumental in providing a comprehensive understanding of the model's operational efficacy and practical viability.
The key metrics used include:
Training Speed: This metric is crucial in assessing the efficiency of the model from a computational perspective. It gauges the rate at which the model processes data and adjusts parameters during training, which is particularly pertinent in large-scale applications. For the purpose of deriving accurate
results, each model underwent a warm-up phase of 5,000 steps before we timed the forward/backward passes across the entire training dataset. The training dataset was deliberately left unshuffled to maintain consistent batches, thereby facilitating efficient and comparable measurements. The final result was then determined by calculating the average time spent on these passes.
Model Complexity: A model's complexity often directly correlates with its performance. However, a balance between complexity and computational efficiency is desirable. Hence, we consider this metric to quantify the intricacy of our model's architecture and algorithms, which also informs us about its interpretability and potential for scaling. Mirroring the approach adopted during the training speed experiment, we maintained uniform batches throughout the entire training dataset, and the resultant metrics were averaged upon the conclusion of an epoch. The metrics were computed with a consistent batch size of 64 to accommodate memory-demanding models.
Memory Consumption: The amount of memory consumed during the model's operation is a significant factor, especially for deployment in resource-constrained environments. This metric provides insights into the model's memory requirements, both during the training phase. Since evaluating the complete dataset evaluation was impractical, we employed an effective measurement strategy. To accommodate the impact of various sample sizes on memory usage, we constructed a worst-case scenario, selecting batches containing the largest samples in a given dataset. An efficient strategy for identifying these substantial samples involved considering the number of nodes in each sample. Consequently, we extracted the  $n$  largest samples from the dataset, where  $n$  is equivalent to the batch size. The memory consumption was then repeatedly measured for this batch, and the results were averaged to yield the final value.