Title: SaVeNet: A Scalable Vector Network for Enhanced Molecular Representation Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: geometric deep learning

molecule property prediction

geometric representation learning
Abstract: Geometric representation learning of molecules is challenging yet essential for applications in multiple domains. Despite the impressive breakthroughs made by geometric deep learning in various molecular representation learning tasks, effectively capturing complicated geometric features across spatial dimensions is still underexplored due to the significant difficulties in modeling efficient geometric representations and learning the inherent correlation in 3D structural modeling. These include computational inefficiency, underutilization of vectorial embeddings, and limited generalizability to integrate various geometric properties. To address the raised concerns, we introduce an efficient and effective framework, Scalable Vector Network (SaVeNet), designed to accommodate a range of geometric requirements without depending on costly embeddings. In addition, the proposed framework scales effectively with introduced direction noise. Theoretically, we analyze the desired properties (i.e., invariance and equivariant) and framework efficiency of the SaVeNet. Empirically, we conduct a comprehensive series of experiments to evaluate the efficiency and expressiveness of the proposed model. Our efficiency-focused experiments underscore the model's empirical superiority over existing methods. Experimental results on synthetic and real-world datasets demonstrate the expressiveness of our model, which achieves state-of-the-art performance across various tasks within molecular representation learning.

SaVeNet: A Scalable Vector Network for Enhanced Molecular Representation Learning

Abstract

Geometric representation learning of molecules is challenging yet essential for applications in multiple domains. Despite the impressive breakthroughs made by geometric deep learning in various molecular representation learning tasks, effectively capturing complicated geometric features across spatial dimensions is still underexplored due to the significant difficulties in modeling efficient geometric representations and learning the inherent correlation in 3D structural modeling. These include computational inefficiency, underutilization of vectorial embeddings, and limited generalizability to integrate various geometric properties. To address the raised concerns, we introduce an efficient and effective framework, S c a lable Ve ctor Network (SAVENET), designed to accommodate a range of geometric requirements without depending on costly embeddings. In addition, the proposed framework scales effectively with introduced direction noise. Theoretically, we analyze the desired properties (i.e., invariance and equivariant) and framework efficiency of the SAVENET. Empirically, we conduct a comprehensive series of experiments to evaluate the efficiency and expressiveness of the proposed model, which achieves state-of-the-art performance across various tasks within molecular representation learning.

1 Introduction

The field of geometric deep learning (GDL) has seen a rapid expansion in recent years, thanks to the successful application of Graph Neural Networks (GNNs) for modeling graph structures [4]. The ability to learn complex geometric representations has driven significant breakthroughs across various disciplines and has proven to be particularly beneficial in diverse areas such as social science [24], physics [14], and neuroscience [6]. Within the realm of molecular representation learning, the use of message-passing-based GNNs has demonstrated remarkable outcomes, especially in understanding the 3D structures of molecules.

Although geometric deep learning has recently seen promising results in various applications in molecular representation learning, its full potential in the field is still untapped. To begin with, a notable research gap persists within this field: the trade-off between expressiveness and efficiency in GNNs. Many studies often resort to the use of multi-hop neighbors [9, 17, 21] or complex embeddings [19] in their pursuit to augment the expressiveness of message-passing based GNNs [9, 29, 33]. While these approaches can enhance the expressiveness and predictive power of the model, they invariably compromise scalability and efficiency due to their computationally intensive nature. As these networks grow more expressive, they demand more computational resources, impeding their ability to scale to larger or more complex datasets. Addressing the challenge of efficiently handling molecular structures to sufficiently capture the complete, multi-level structural information during learning without compromising the computational efficiency of the network is of critical importance. Our approach circumvents these issues by avoiding using multi-hop GNNs and

expensive embeddings. Novel approaches for initializing and processing vector-typed embeddings have been designed to work with vector representations with the goal of achieving a balance between maintaining numerical stability, facilitating faster convergence, and enhancing the model's ability to generalize to new datasets [25, 33]. Utilizing one-hop vector embeddings work with our proposed SAVENET significantly enhances the network's efficiency, demonstrating their effective integration in this context. Figure 1 provides a visual illustration of the advantages of our proposed SAVENET , which compares our method with state-of-the-art methods on the QM9 dataset. The values on the x -axis show the latency of the model, and the values on the y -axis show the standard mean absolute error (std. MAE) of the models' performance. Because lower is better for both latency and std. MAE, the best model in both efficiency and expressiveness is the one that is closest to the (0,0) point. In terms of performance, our base model, denoted as SAVENET -B, surpasses the existing state-of-the-art method Equiformer [19], while concurrently enhancing training duration efficiency by a factor of 14.6. This substantial improvement underscores the efficacy of our model in achieving superior results with increased efficiency.

Furthermore, although several prior works [2, 3, 12, 18, 19, 25, 27, 29, 30] proposed equivariant neural networks that handle scalar and vector representations together for molecular representation learning, one primary problem with these existing works is that vector-valued representations show limited performance or scalability and cannot equally contribute towards performance improvements in the network [12, 18, 29]. It remains challenging to scale the models to improve the expressiveness of the model further. We propose an efficient framework to tackle this challenge. Our model's scalability is exemplified by its capacity to stack N layers, effectively serving as an encoder for learning the latent geometric representations, thereby augmenting the model's

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.6

0.7

0.8

0.9

1.0

Latency (ms)

std. MAE

ComENet

SphereNet

DimeNet++

SaVeNet-B

SaVeNet-L

EqGAT

ET

PaiNN

expressiveness. These superior scaling capabilities enhance the model's expressiveness and effectively distinguish our work from existing research, which often struggles to scale efficiently [9, 19]. The scalability of our approach is not merely theoretical but is empirically demonstrated in Section 5.1. Figure 1 illustrates one example of the scaled larger version of our model: SAVENET -L, which exhibits enhanced capability in contrast to our base model and existing baselines. This advancement substantiates the prospective capacity of our model to scale effectively, thereby indicating potential for superior performance in future implementations.

Our paper's main contributions are: (i) Proposed an efficient and effective message-passing framework coupled with several decoder mechanisms for leveraging the geometric representations within the 3D molecular graphs. (ii) Novel approaches for initializing and processing vector-typed embeddings have been designed to work with vector representations to ensure numerical stability, speed up model convergence, and improve the model's ability to adapt to new datasets. (iii) A thorough evaluation of the expressiveness, efficiency, and scalability of our models compared with existing benchmarks, as well as detailed ablation studies, underscore the advantages of our proposed SAVENET .

2 Related work

Equivariant Graph Neural Networks The properties of graph representations can change with transformations such as translation, rotation, and permutation. By design, GNNs inherently permutation equivalence [15]. While it is possible to introduce diverse transformations to the network through data augmentation, this approach can lead to computational inefficiencies.

The concepts of invariance and equivariance have been long acknowledged as critical to the success of various tasks [2, 13, 35, 36]. Recent research in molecular representation learning has predominantly addressed this challenge with equivariant network architectures. Within the realm of molecules,

numerous equivariant models have been proposed for tasks such as chemical property prediction [16], protein structure prediction [13], and energy prediction [17].

Molecular Property Prediction Over recent years, there have been considerable advancements in molecular property prediction [9, 16, 19, 28]. These developments have been primarily driven by the adoption of graph models for molecular representation, coupled with the application of kernel methods for structural learning. Such kernels operate based on geometric features, such as interatomic distances, which are encoded using radial basis or Bessel basis functions. Several other studies employ spherical embeddings to model message-passing [3, 9, 17, 19, 21, 33]. Dimenet++ [9], GemNet [17] and SphereNet [21] depend on a multi-hop message-passing to model angular properties, which introduces computational complexity that curtails scalability with respect to the model size of these approaches. These advanced representations have resulted in more precise property predictions, including but not limited to dipole moment, HOMO-LUMO gap, energy, and force. However, the precision of these predictions is accompanied by increased computational complexity due to their dependence on quadratic or even cubic operations. Recent works have strived to address this issue, proposing innovative techniques that improve computational efficiency with minimal or no detriment to prediction accuracy [21, 33]. These novel techniques continue to operate on higher-order features during feature generation, underscoring the ongoing challenge of balancing computational efficiency with predictive precision. Both ET [30] and EQGAT [18] incorporate attention mechanisms into vector representations. However, the complex interactions in these works often limit the scalability and efficiency of the networks. More recently, networks such as SEGNN [3] and Equiformer [19] have utilized Clebsch-Gordan tensor products to achieve high performance. However, these tensor products are empirically found to be significantly slower than the previously discussed methods, posing a challenge when applying such methods to large, real-world datasets.

Our model utilizes 1-hop geometric representations, offering a balance of expressivity and computational efficiency. This design choice circumvents the need for complex and potentially lossy conversions often associated with higher-order representations. By harnessing the rich information in 1-hop representations, we maintain data integrity and enhance model performance without substantial computational overhead.

3 Preliminaries

Notations. We formulate the atomic structure graph representation, where each input molecule is represented as a graph G = ( V , E ) . V is the set of N nodes in the graph, each node representing an atom in the molecule, and E ⊆ V × V is the set of edges that connects the nodes. The average degree of a node, denoted by k , represents the average number of edges per node. The geometric information of a node for each v i ∈ V is represented by the coordinates c i = ( x i , y i , z i ) in the Cartesian coordinate system. The scalar representations s generally include properties like relative distance r , while the vector representations V include direction vectors d ⃗ . We use ⊙ to denote the element-wise (Hadamard) product and || · || 2 to denote the row-wise L 2 norm. The cross-product between two vectors is denoted as × .

Invariance and Equivariance. A function f : X → Y is said to be equivariant to a transformation if applying the transformation to the input of the function is the same as applying the transformation to the output of the function. Formally, for a transformation T , a function f is equivariant if for all x ∈ X : T ( f ( x )) = f ( T ( x )) . This means that the function preserves the structure of the transformation. On the other hand, a function is invariant to a transformation if its output does not change when the transformation is applied to its input. Formally, for a transformation T , a function f is invariant if for all x ∈ X : f ( T ( x )) = f ( x ) . This implies that the function is invariant to the transformation; the function's output remains unchanged regardless of the applied transformation.

Euclidian Transformations. In this paper, we frequently use the special orthogonal group SO (3) , comprising all 3x3 rotation matrices R that satisfy R T R = I and det ( R ) = 1 , and the special Euclidean group SE (3) , consisting of all 4 × 4 special Euclidean transformation matrices T that can be written as T = [ R t 0 1 ] , where R ∈ SO (3) , t ∈ R 3 is a translation vector, and 0 is a row vector

of zeros. These matrices satisfy T -1 T = I and det ( T ) = 1 . SE (3) forms a manifold, the product of SO (3) and R 3 , allowing operations combining rotations and translations in 3D space.

4 Scalable Vector Network

Our goal is to design a message-passing neural network that is aware of the geometric characteristics of the molecular graph, with consideration of expressiveness and efficiency, and preserve these characteristics beneficial for downstream tasks. We first introduce the representations used in SAVENET and how we handle these geometric representations. Then we describe the overall message-passing framework to model the geometric properties in the graph and show how to integrate our proposed framework into the representation learning process on downstream tasks.

4.1 Efficient Representations for 3D graphs

The initial representations fed into the network are instrumental in achieving complete geometric representations [21, 33]. It is crucial to efficiently represent geometric graphs without the loss of their properties. Our framework operates by employing embeddings from one-hop neighbors exclusively. Invariant properties are encoded as scalar values, which enable the network to maintain invariance without resorting to computationally expensive operations. For equivariant properties, we extract the inherent directional information from the input graphs and encode these features into vector representations. The directional information between nodes i and j is represented with ⃗ β ij and defined as follows:

Here, c represents the average of c , and × denotes the cross-product operation. With the geometric representation defined, the efficacy of our scalar and vector representations is further demonstrated in the following lemma:

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 ij ) properties.

The proof for Lemma 1 is provided in the Appendix A. Lemma 1 can be further generalized, requiring no positional information by translating the p 0 reference point to the origin using a t o translation vector. In light of this, we propose the following theorem, which demonstrates the SAVENET's ability to reconstruct the rotation-equivariant geometric structure of an input graph:

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.

The proof for Theorem 1 is provided in Appendix B. The vector representations are an effective way to model equivariant networks and are widely used networks [11, 12, 25, 29] on different challenging tasks. With the proof of Lemma 1 and Theorem 1, we present a provably efficient and lossless geometric encoding for the 3D graphs without reliance on the multi-hop neighborhoods. Despite the vector representation offering effective encodings, several unresolved challenges hinder its successful application towards enhancing the model's expressiveness and scalability. In response to this, we propose the incorporation of direction noise initialization and a vector activation class within our SAVENET framework, specifically designed to address this issue.

Direction Noise. Vector representations are typically initialized

X

Y

Z

using dataset-specific features [12, 25], such as direction vectors between sequential neighbors in a protein's amino acid sequence. However, in the absence of such data, existing works [18, 29] often default to initializing these vectors as null, allowing the network to learn these vector representations intrinsically. Such initialization can unnecessarily consume computational resources, particularly in the network's initial layers. In this work, as illustrated in Figure 2, we propose a novel strategy of

introducing directional noise utilizing the spherical coordinate system as an initialization measure for the vector representations generated based on the node type. This approach is designed to mitigate numerical instability issues and empower the network to explore the latent space we propose through direction noises. We follow the following convention to generate a direction vector for a given node type F :

where θ i and ϕ i are learnable parameters for node type i and γ is a control parameter governing the magnitude of the direction embeddings. This approach effectively generates direction vectors using the spherical coordinate system. The value of γ is associated with the progress of the training, similar to learning rate schedulers. The value of γ is adjusted throughout training, following a decay schedule that ensures it reaches zero before training concludes. This methodology offers an effective means for creating direction vectors using a spherical coordinate system while allowing for the adaptive control of the direction embeddings' magnitude throughout training.

Vector Activations. Applying similar activation functions used in scalars on vector representations, where the function operates on each element of the vector independently, can break the model's equivariance. This is because these functions are not designed to handle the geometric transformations that the vectors might undergo. Consequently, SAVENET introduces a vector activation function in which the activation is computed based on the vector's L 2 norm. For vector representation V , the vector activation function is written as follows:

where σ v is a scalar activation function. To further regulate the activation, the vector activation function incorporates a bias term b , where b is the same shape as the representation dimension C . Our experimentation found that vector activation functions significantly enhance the numerical stability during training, particularly with larger-scale models.

4.2 SAVENET for 3D Graphs

This section describes the overall architecture of SAVENET in detail. The overall SAVENET framework is a SO(3)-equivariant message-passing neural network incorporating the vector initialization and activation for scalable molecular representation learning. It is based on an encoder-decoder structure comprising multiple stacked encoder and decoder layers. The encoder layers learn the latent representation of the geometric features, and the decoder layers decode the learned latent representations for downstream tasks on either invariant or equivariant targets.

Encoder. The encoder in SAVENET follows the message-passing paradigm and incorporates the above-defined geometric representations. The main two components of the encoder are inter-atomic interactions followed by atom-wise blocks.

Vector Updates. Vector updates are employed to compute channel-wise interactions in vector representations. By conducting identical operations across spatial dimensions, vector updates ensure that the equivariance property is maintained. More specifically, given a vector representation V ∈ R C × 3 , the vector update applies a learnable weight matrix W ∈ R C × C ′ , where C and C ′ denote the dimensions of the input and output representations, respectively. Therefore, the output vector V ′ is computed as V ′ = W V .

Inter-atomic interactions. Our model encapsulates inter-atomic interactions by integrating both equivariant direction vectors and invariant distance filters, which are encoded using radial basis functions. We have designed the scalar interaction path to interface with the vector interactions in a manner that preserves the equivariance property, achieved by appropriately scaling the magnitudes of the vectors. The overall formulation of the inter-atomic interactions can be expressed as follows:

In this equation, N i represents the set of neighbors of node i , and e is the interaction function, ⃗ β denotes the equivariant vectors. To accommodate a variety of equivariant properties, we fuse these vectors with vector updates VA. The interaction function is formalized as follows:

where ϕ denotes sequentially stacked perceptron layers and η denotes distance encoding. Specifically, ϕ b serves as a vector update layer as it operates on vector inputs. The distance encoding is calculated using the formula η ( r ij ) = ϕ r ( χ ( r ij ) ∗ ω ( r ij )) , where χ ( · ) denotes a basis function and ω ( · ) stands for a cutoff function.

Atom-wise blocks. Following the execution of the interaction block, the workflow proceeds to the atom-wise blocks, which are responsible for computing interactions between invariant and equivariant representations, as well as performing channel-wise updates. The formal expression of this process is defined as:

where ϕ { m,v } denotes multi-layer perceptron (MLP) and ϕ vu denotes vector updates. Here, the vector representations interact with scalar representations via the L 2 norm of updated vector representations. Moreover, the scalar representations scale the vector representations after applying channel-wise interactions ϕ v . This process presents bi-directional communication between different representations, thus enhancing the model's ability to learn and extract meaningful information from data.

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.

Equivariance 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 ∈ 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.

5 Experiments

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

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 × the quantity of graphs as the QM9 dataset and an approximate 1 . 6 × and 1 . 9 × 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-

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.

5.1 Results

Accuracy results on QM9. Table 2 listed the experimental results of the QM9 dataset. SAVENETB 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, ϵ HOMO, ϵ LUMO, C ν , 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 ϵ HOMO, ϵ 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.

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 × 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 × and 5.6 × 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 × 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.

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 µ task and 92% on the ⃗ µ task compared to the next best model, SphereNet. Our SAVENET-L model, further improved the performance of SAVENET-B. For the µ task, SAVENET-L showed an improvement of approximately 26% over SAVENET-B; for the ⃗ µ 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.

5.2 Ablation Study

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 (V A), 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

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 V A 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.

References

Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems , 32, 2019.

Sarp Aykent and Tian Xia. GBPNet: Universal Geometric Representation Learning on Protein Structures. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery Data Mining , KDD '22, Washington, DC, USA, 2022. Association for Computing Machinery.

Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling. Geometric and physical quantities improve e(3) equivariant message passing. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum? id=_xwr8gOBeV1 .

Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geometric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine , 34 (4):18-42, 2017.

Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and Tie-Yan Liu. SE(3) equivariant graph neural networks with complete local frames. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 5583-5608. PMLR, 2022. URL https://proceedings.mlr.press/v162/du22e.html .

Timothy W Dunn, Jesse D Marshall, Kyle S Severson, Diego E Aldarondo, David GC Hildebrand, Selmaan N Chettih, William L Wang, Amanda J Gellis, David E Carlson, Dmitriy Aronov, et al. Geometric deep learning enables 3d kinematic profiling across species and environments. Nature methods , 18(5):564-573, 2021.

Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In International Conference on Machine Learning , pages 3165-3176. PMLR, 2020.

Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks. In Advances in Neural Information Processing Systems , volume 33, pages 1970-1981. Curran Associates, Inc. URL https://proceedings.neurips.cc/paper/2020/hash/ 15231a7ce4ba789d13b722cc5c955834-Abstract.html .

Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop, NeurIPS , 2020.

Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning , pages 1263-1272. PMLR, 2017.

Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum? id=1YLJDvSx6J4 .

Bowen Jing, Stephan Eismann, Pratham N Soni, and Ron O Dror. Equivariant graph neural networks for 3d macromolecular structure. arXiv preprint arXiv:2106.03843 , 2021b.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

Yashar Kiarashinejad, Mohammadreza Zandehshahvar, Sajjad Abdollahramezani, Omid Hemmatyar, Reza Pourabolghasem, and Ali Adibi. Knowledge discovery in nanophotonics using geometric deep learning. Advanced Intelligent Systems , 2(2):1900132, 2020.

Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations (ICLR) , 2017.

Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. In International Conference on Learning Representations , 2020.

Johannes Klicpera, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. URL https: //openreview.net/forum?id=HS_sOaxS9K-.

Tuan Le, Frank Noé, and Djork-Arné Clevert. Equivariant graph attention networks for molecular property prediction. arXiv preprint arXiv:2202.09891 , 2022.

Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=KwmPfARgOTD .

Meng Liu, Youzhi Luo, Limei Wang, Yaochen Xie, Hao Yuan, Shurui Gui, Haiyang Yu, Zhao Xu, Jingtun Zhang, Yi Liu, Keqiang Yan, Haoran Liu, Cong Fu, Bora M Oztekin, Xuan Zhang, and Shuiwang Ji. DIG: A turnkey library for diving into graph deep learning research. Journal of Machine Learning Research , 22(240):1-9, 2021. URL http://jmlr.org/papers/v22/ 21-0343.html .

Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=givsRXsOt9r .

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum? id=Bkg6RiCqY7 .

Microsoft. Deepspeed: A deep learning optimization library for pytorch. https://github. com/microsoft/DeepSpeed , 2020.

Federico Monti, Fabrizio Frasca, Davide Eynard, Damon Mannion, and Michael M Bronstein. Fake news detection on social media using geometric deep learning. arXiv preprint arXiv:1902.06673 , 2019.

Alex Morehead and Jianlin Cheng. Geometry-complete perceptron networks for 3d molecular graphs. AAAI Workshop on Deep Learning on Graphs: Methods and Applications , 2023.

Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.

Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) Equivariant Graph Neural Networks. In Proceedings of the 38th International Conference on Machine Learning , pages 9323-9332. PMLR. URL https://proceedings.mlr.press/v139/satorras21a.html .

Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. Advances in neural information processing systems , 30, 2017.

Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning , pages 9377-9388. PMLR, 2021.

Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations , 2022.

Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds. URL http://arxiv.org/abs/1802.08219 .

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.

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 .

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 .

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 .

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.

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.

Appendix

A Proof of Lemma 1

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 ⃗ ij ) 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 ∈ 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 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 d ⃗ ij from c i to c j . This fact contradicts our assumption, and hence concludes the proof.

B Proof of Theorem 1

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 Π j i 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 ∈ V by traversing each edge on the path Π j i 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 Π j i .

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 , ⃗ 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 ∪ { j } .

Thus, we conclude that the input graph's rotation-equivariant geometric structure can be reconstructed using ( d, r ) properties, thereby completing our proof.

C Proof of Theorem 2

Theorem 2: The equivariant representation of SAVENET is equivariant to any given rotation matrix in R ∈ 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 ⃗ β . It's clear that these vectors, representing direction, maintain their properties under any global rotation. Formally, if f b ( G ) represents the function that extracts ⃗ β from an input graph G , we have that Rf b ( G ) = f b ( R G ) for all

R ∈ 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.

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, ϕ b , ϕ d , and ϕ 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:

Each component in the above expression preserves symmetry properties. More specifically, the first term ϕ b ( ⃗ β ij ) preserves equivariance because the same transformation is applied across spatial dimensions. The second term, ϕ d ( s j ) , is a linear transformation applied to an invariant, and therefore, remains invariant under any transformation in T . The third term, η 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, ϕ 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 η r ( r ij ) represents an invariant scalar, remaining unchanged under transformation.

The first part ϕ b ( ⃗ β ij ) ⊙ ϕ d ( s j ) ⊙ η d ( r ij ) involves the Hadamard product of an equivariant vector ( ϕ b ( ⃗ β ij ) ∈ R C × 3 ) with invariant scalar values ( ϕ d ( s j ) , η d ( r ij ) ∈ 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 ⊙ ϕ v ( s j ) ⊙ η r ( r ij ) involves the Hadamard product of an equivariant vector ( V j ∈ R C × 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 ∈ SO (3) . This concludes our examination of the equivariance of the interaction function e v under any rotation R ∈ SO (3) . Similarly, the operations in the invariant update branch can be captured as:

where ϕ s denotes a linear transformation, and η 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:

In this equation, ϕ vu denotes the vector update function which is, as previously established, equivariant under any rotation matrix R . The norm operation || · || applied to V is invariant, as the magnitude of the vector does not change with rotation. Consequently, the composition || ϕ vu ( V ) || is invariant to global rotations. Furthermore, applying the linear transformation ϕ m to the concatenation of two invariants, in this case || ϕ 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:

This equation is similar to the one describing the invariant update branch, with a difference lying in the linear transformation ϕ v that uses different learnable parameters. Hence, the second term is invariant. The first term, ϕ 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 ∈ 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.

D Datasets

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 (L+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 ( ⃗ µ ), energies of the highest occupied and lowest unoccupied molecular orbitals ( ϵ HOMO, ϵ LUMO), and orbital energy gap ( ∆ ϵ ), it offers an essential testbed for model evaluation. The dipole magnitude ( µ ) can be easily derived from ⃗ µ , highlighting the significance of equivariance in models due to the influence of global rotations on ⃗ µ . 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.

F Implementation details

The reference implementation can be found in https://github.com/EfficientGraphs/SaVeNet.

F.1 Hyperparameters

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.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel molecule representation network that enhances the learning capacity and scalability through the integration of innovative initialization techniques and activation functions for vector features. The conducted experiments validate the network's proficiency in three distinct molecule representation tasks, demonstrating its capability to effectively learn from data while maintaining scalability.\n\nSTRENGTHS:\n- A new equivariant network capable of learning scalar and vector features.\n- Superior performance and scalability compared to previous methods.\n\nWEAKNESSES:\n- The presentation of the method can be improved. The symbols and equations employed are unclear, such as the presence of a comma in Eq. 4 and the excessive use of unnecessary new symbols in Eq. 5. Furthermore, the substitution of the symbol for multiplication with the letter 'x' in line 187 could be clarified. Additionally, it is suggested that the authors consider using a figure to illustrate the message passing of scalar and vector features, similar to previous methods like DimeNet. \n\n- The core difference between the proposed network and previous networks for vector features, such as PaiNN, remains unclear. While it appears that the primary distinction lies in the different processing of features, it would be beneficial to highlight and discuss this core difference in greater detail.\n\nQUESTIONS:\nRegarding the direction noise, it would be helpful to clarify whether theta and phi are deterministic functions of the node type. If they are, applying a rotation to the graph would not change the node type, thereby maintaining the invariance instead of the equivariance of VEN(F_i). So why can VEN(F_i) be an equivariant vector initialization?  Further explanations and intuitions could be added to address this point and enhance understanding.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an effective and efficient equivariant graph neural network for geometric learning on molecules. The model encodes 3D graphs with node types and coordinates and outputs scalar and vector representations. The message passing process is purely scalar-based, which enjoys more efficiency than baselines using irreducible representations and high-order tensor objects . A vector initialization based on spherical coordinate system is utilized for better numerical stability. Experiments on QM9, N-body and molecule3D demonstrate the superiority of the proposed model over the baselines.\n\nSTRENGTHS:\n1. The proposed scalar-based message passing scheme enjoys theoretically better efficiency and empirically better performance compared to existing baselines.\n2. The experiments are thorough, including three popular benchmarks and a variety of baselines.\n3. The empirical analysis on efficiency is comprehensive, covering  time for forward/backward pass, FLOPs, MACs and memory consumption.\n\nWEAKNESSES:\n1. Several points claimed in the contribution are not well supported:\n   - The authors state that the vector initialization can ensure numerical stability and speed up model convergence. However, no theoretical analysis on why it can ensure numerical stability is provided, and neither experimental evidence nor theoretical proof of speeding up convergence is provided.\n   - The scaled version of the baselines should also be compared with SAVENET-L to show that the proposed model benefits more from the scaling (i.e. Has better scalability). Moreover, it is not shown whether the scalability of the proposed model can hold when keeping increasing the model size.\n2. The presentation of the model structure can be optimized (e.g. add an overview figure).\n\nQUESTIONS:\n1. What is the $h_j$ in Eq.5?\n2. What do $n, l_e, l_d$ represent in the theoretical complexity analysis?\n3. Are the coordinates fixed during message passing?\n4. Why can the directional initialization improve numerical stability? What is the reason for using spherical coordinate system as initialization? What is the difference if we use random vectors as initialization and anneal their norm to zero during training?\n5. I think the computational scalability of GNNs is mainly restricted by the fully connected message passing which has quadratic complexity with respect to the number of nodes. Therefore in large graphs (e.g. proteins), people often use k-nearest neighbors to reduce the complexity to be linear to the number of nodes. However, the expressiveness of scalar-based equivariant GNNs can only be ensured in nearly fully connected message passing [1]. So can you show that the proposed models can maintain its superiority over the baselines when adapted to larger graphs like proteins, where edges might have to be much more sparser than in fully connected graphs. Or, under the current experiment settings, will the superiority be maintained if the edges get sparser (e.g. use k-nearest neighbors for message passing and decrease the K or decreasing the cutoff in the radial graph)?\n6. Will the scalability holds when keeping increasing the hidden size and the number of layers?\n\n[1] Scalars are universal: Equivariant machine learning, structured like classical physics\n\nLIMITATIONS:\nWhen scaling to larger graphs (e.g. proteins), the edges may have to be sparser (e.g. K-nearest neighbors) and the scalar-based message passing potentially will have worse expressiveness.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an SE(3)-equivariant model called SaVeNet, designed to accommodate various geometric requirements. The proposed framework can effectively scale with the introduction of directional noise. Theoretical analysis and empirical results on several datasets are provided to validate the efficiency and expressiveness of SaVeNet.\n\nSTRENGTHS:\n* The proposed method achieves outstanding performance across several synthetic and real-world datasets.\n* SaVeNet exhibits better efficiency than baseline methods.\n\nWEAKNESSES:\n* Many claims about the proposed method and related works lack supportive evidence. For example,\n    * Line 40, \"Novel approaches for initializing ... achieving a balance between maintaining numerical stability, facilitating faster convergence, and enhancing the model’s ability to generalize to new datasets.\" This sentence states that the proposed approach can maintain numerical stability and facilitate faster convergence, but there is a lack of convincing empirical ablation studies or theoretical evidence supporting these claims.\n    * Line 58-63, \"one primary problem...limited performance or limited performance or scalability and cannot equally contribute towards performance improvements in the network.\" These sentences make informative claims about the limitations of existing methods, but the authors do not provide detailed discussions or analyses to support these points, weakening the motivation for SaVeNet.\n    * Line 63-68, the authors propose to tackle the scalability challenge by stacking multiple encoder layers, which is a widely adopted operation in equivariant models. However, it is unclear why other baseline models (especially those vector-based models such as PaiNN) cannot improve scalability using the same operations.\n* The methodology section (Section 3) is poorly organized and written. \n    * Some notations are not well defined. E.g., $IA$, $s_i$ and $e$ in Equation 4 are not defined. $h_j$ in Equation 5 is not defined.\n    * Some equations are difficult to understand. E.g.,  the procedure of updating vector $V'_i$ from Equation 4 is unclear. It seems that some details of the equation are missing, making it difficult to integrate Equation 5 and Equation 4.\n* The overall novelty of SaVeNet is limited. Some techniques in the framework (e.g., the scalar/vector-based equivariant operations) have been studied in related works such as PaiNN, and ET. The newly proposed techniques, direction noise and vector activations, account for the main contribution of this paper. However, the ablation results in Table 6 show that removing these two modules only results in a slight performance drop, which cannot validate their importance. Accordingly, the technical contribution of SaVeNet is limited. Considering that many details in the methodology section are unclear, all comments are made based on the current version.\n* There are some grammatical mistakes. \n    * Line 47, `model's` should be `model`\n    * Line 187, $R^{Cx3}$ should be $R^{C \\times 3}$\n\nIn summary, while the proposed method demonstrates exceptional performance across various downstream tasks and exhibits promising scalability, the limited technical novelty and issues with the presentation render the current version unsuitable for acceptance in NeurIPS.\n\nQUESTIONS:\n* The metrics std. MAE and log MAE lack detailed descriptions. How did the authors calculate these metrics?\n* The proposed direction noise module's SE(3)-equivariance is unclear, as there is no provided theoretical proof.\n\nLIMITATIONS:\nThe authors well clarify the limitations and potential negative societal impact of their work in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

comments: SUMMARY:
This paper proposes a framework called SaVeNet for geometric representation learning of molecules. The paper includes theoretical analysis and empirical experiments to demonstrate the superiority of SaVeNet over existing methods in terms of efficiency and expressiveness.

STRENGTHS:
1. This paper proposes an efficient geometric encoding, a novel directional noising strategy in the spherical coordinate system, and a novel vector activation function.

2. The experimental results show that SaVeNet achieves very good performance over multiple invariant and equivariant molecular tasks. 

3. As an equivariant model, the efficiency of SaVeNet is pretty good.

WEAKNESSES:
It will help readers to better understand if authors can add a figure of the proposed SaVeNet’s architecture, including the illustration of components described in Sec 3.1 and 3.2.

QUESTIONS:
1. I’m wondering if the direction noise and vector activation also applicable to other equivariant models?

2. Can you add more equivariant baselines like PaiNN to Figure 1 for efficiency comparison?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose an efficient and scalable equivariant GNN (SaVeNet) for 3D molecular conformations. The architecture follows an encoder-decoder style framework, where the encoder is composed of \"Inter-atomic Interactions\" and \"Atom-wise blocks\" learning mechanisms. Several other modeling augmentations are proposed such as 1-hop molecule embeddings in the feature space, non-null vector initializations (\"Direction Noise\"), and norm-aware vector activation functions (\"Vector Activations\").\n\nSTRENGTHS:\nThe thrust of this paper is clear: designing GNNs with SO(3)/SE(3) equivariance can introduce complexity into the model. This slows down inference-time latency, and may inhibit scaling the model to larger sizes (i.e. a RNN with parameter capacity equal to a Transformer may have similar performance, but Transformers scale significantly better).  Accordingly, the authors develop an equivariant GNN that uses only one-hop neighborhood vector embeddings (i.e. cheap & scales), but (assuming a connected graph -- which is the case for molecules) this feature-space representation is sufficient to reconstruct the entire molecule. The empirical results align with the theory, and state-of-the-art results are attainted on a variety of 3D benchmarks.\n\nTo summarize, I believe the paper (a) tackles a relevant problem (b) proposes a relevant solution and (c) demonstrates theoretically and empirically that solution works. Therefore, I recommend acceptance.\n\nI'd also like to note the punctilious empirical evaluation---efficiency analyses can be especially difficult, but it's my opinion that the authors perform a fair comparison among methods to constitute their empirical findings.\n\nWEAKNESSES:\nI am not an expert on equivariant GNNs and am not familiar with much of the related work. I am also unable to recognize any method/evaluation-specific \"red flags\" which may appear and defer to the expertise of the other reviewers. \n\nOne facet that stood out; however, is that the writing at times feels long. It could be significantly shortened in many areas. For instance, the abstract is 19 lines and could be made more \"pointed\". I would also suggest creating a \"Figure 1\" that is a visual depiction of the SaVeNet architecture to supplement the material in Section 3.2 on learning mechanisms.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nYes -- the discussion in Section 5 was appreciated and meaningful.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}