SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sarp Aykent
sarp@auburn.edu
Tian Xia
tian@auburn.edu
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, which achieves state-of-the-art performance across various tasks within molecular representation learning.
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
Figure 1: Comparison of our proposed SAVENET and the state-of-the-art methods: latency vs. std. MAE on QM9 dataset.
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.
**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.
Notations. We formulate the atomic structure graph representation, where each input molecule is represented as a graph  $\mathcal{G} = (\mathcal{V},\mathcal{E})$ .  $\mathcal{V}$  is the set of  $N$  nodes in the graph, each node representing an atom in the molecule, and  $\mathcal{E}\subseteq \mathcal{V}\times \mathcal{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}\in \mathcal{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  $\vec{d}$ . We use  $\odot$  to denote the element-wise (Hadamard) product and  $\| \cdot \| _2$  to denote the row-wise  $L^2$  norm. The cross-product between two vectors is denoted as  $\times$ .
Invariance and Equivariance. A function  $f: X \to 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 \in 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 \in 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.
Euclidean Transformations. In this paper, we frequently use the special orthogonal group  $SO(3)$ , comprising all  $3 \times 3$  rotation matrices  $R$  that satisfy  $R^T R = I$  and  $\operatorname{det}(R) = 1$ , and the special Euclidean group  $SE(3)$ , consisting of all  $4 \times 4$  special Euclidean transformation matrices  $T$  that can be written as  $T = \begin{bmatrix} R & t \\ 0 & 1 \end{bmatrix}$ , where  $R \in SO(3)$ ,  $t \in \mathbb{R}^3$  is a translation vector, and  $0$  is a row vector of zeros. These matrices satisfy  $T^{-1} T = I$  and  $\operatorname{det}(T) = 1$ .  $SE(3)$  forms a manifold, the product of  $SO(3)$  and  $\mathbb{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  $\vec{\beta}_{ij}$  and defined as follows:
$$
\vec {\beta} _ {i j} = \left\{\vec {d} _ {i j} = \frac {\left(c _ {i} - c _ {j}\right)}{r _ {i j}}, \vec {t} _ {i j} = \left(c _ {i} - \bar {c}\right) \times \left(c _ {j} - \bar {c}\right), \vec {o} _ {i j} = \vec {d} _ {i j} \times \vec {t} _ {i j} \right\} \in \mathbb {R} ^ {3 \times 3} \tag {1}
$$
Here,  $\overline{c}$  represents the average of  $c$ , and  $\times$  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}, \vec{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.
Figure 2: Direction noise.
Direction Noise. Vector representations are typically initialized
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$ :
$$
\operatorname {V E N} \left(F _ {i}\right) = \left\{\gamma \left(\sin \theta_ {i} \cos \phi_ {i}\right), \gamma \left(\sin \left(\theta_ {i}\right) \sin \left(\phi_ {i}\right)\right), \gamma \left(\cos \left(\theta_ {i}\right)\right) \mid \theta_ {i} \in \mathbb {R}, \phi_ {i} \in \mathbb {R} \right\} \tag {2}
$$
where  $\theta_{i}$  and  $\phi_{i}$  are learnable parameters for node type  $i$  and  $\gamma$  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  $\gamma$  is associated with the progress of the training, similar to learning rate schedulers. The value of  $\gamma$  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:
$$
\operatorname {V A} (V) = \left(V \odot \sigma_ {v} \left(\left| \left| V \right| \right| + b\right) \mid V \in \mathbb {R} ^ {\mathrm {C} \times 3}, b \in \mathbb {R} ^ {C}\right) \tag {3}
$$
where  $\sigma_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 \in \mathbb{R}^{C \times 3}$ , the vector update applies a learnable weight matrix  $\mathbf{W} \in \mathbb{R}^{C \times C'}$ , where  $C$  and  $C'$  denote the dimensions of the input and output representations, respectively. Therefore, the output vector  $V'$  is computed as  $V' = \mathbf{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:
$$
\begin{array}{l} I A (s, V, r, \vec {\beta}) = \left(s _ {i}, V _ {i}\right) + \sum_ {j \in \mathcal {N} _ {i}} e \left(s _ {j}, V _ {j}, r _ {i j}, \vec {\beta} _ {i j}\right) = \left(s _ {i} ^ {\prime}, V _ {i} ^ {\prime}\right) \\ s _ {i} ^ {\prime} = s _ {i} + \sum_ {j \in \mathcal {N} _ {i}} e _ {s} \left(s _ {j}, V _ {j}, r _ {i j}, \vec {\beta} _ {i j}\right) \tag {4} \\ V _ {i} ^ {\prime} = V _ {i} + \sum_ {j \in \mathcal {N} _ {i}} e _ {v} (s _ {j}, V _ {j}, r _ {i j}, \vec {\beta} _ {i j}) \\ \end{array}
$$
In this equation,  $\mathcal{N}_i$  represents the set of neighbors of node  $i$ , and  $e$  is the interaction function,  $\vec{\beta}$  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:
$$
\begin{array}{l} 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) \\ 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 {5} \\ \end{array}
$$
where  $\phi$  denotes sequentially stacked perceptron layers and  $\eta$  denotes distance encoding. Specifically,  $\phi_{b}$  serves as a vector update layer as it operates on vector inputs. The distance encoding is calculated using the formula  $\eta (r_{ij}) = \phi_r\big(\chi (r_{ij})*\omega (r_{ij})\big)$ , where  $\chi (\cdot)$  denotes a basis function and  $\omega (\cdot)$  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:
$$
\mathrm {A W} (s, V) = \left\{\left(\phi_ {m} (\| \phi_ {v u} (V) \| \cup s)\right), \left(\mathrm {V A} \left(\phi_ {v u} (V) \odot \phi_ {v} (\| \phi_ {v u} (V) \| \cup s)\right) \right. \right\} \tag {6}
$$
where  $\phi_{\{m,v\}}$  denotes multi-layer perceptron (MLP) and  $\phi_{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  $\phi_v$ . This process presents bi-directional communication between different representations, thus enhancing the model's ability to learn and extract meaningful information from data.