Empowering Active Learning for 3D Molecular Graphs with Geometric Graph Isomorphism

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:55:35

Abstract: Molecular learning is pivotal in many real-world applications, such as drug discovery. Supervised learning requires heavy human annotation, which is particularly challenging for molecular data, e.g., the commonly used density functional theory (DFT) is highly computationally expensive. Active learning (AL) automatically queries labels for most informative samples, thereby remarkably alleviating the annotation hurdle. In this paper, we present a principled AL paradigm for molecular learning, where we treat molecules as 3D molecular graphs. Specifically, we propose a new diversity sampling method to eliminate mutual redundancy built on distributions of 3D geometries. We first propose a set of new 3D graph isometries for 3D graph isomorphism analysis. Our method is provably at least as expressive as the Geometric Weisfeiler-Lehman (GWL) test. The moments of the distributions of the associated geometries are then extracted for efficient diversity computing. To ensure our AL paradigm selects samples with maximal uncertainties, we carefully design a Bayesian geometric graph neural network to compute uncertainties specifically for 3D molecular graphs. We pose active sampling as a quadratic programming (QP) problem using the proposed components. Experimental results demonstrate the effectiveness of our AL paradigm, as well as the proposed diversity and uncertainty methods.

Main paper: 
 introduction: 
 Molecular representation learning is essential for a variety of real-world applications, such as molecular design, drug discovery, material design, etc. In recent studies, molecules have been formulated as 3D graphs, based on the evidence that 3D spatial information is crucial to determine the properties of molecules (Liu et al., 2019; Townshend et al., 2019; Axelrod & Gomez-Bombarelli, 2020) . In a 3D graph, atoms are nodes, each of which is associated with the Cartesian coordinates in 3D space. A predefined cut-off distance can be used as a threshold to determine if there is an edge between two nodes in the 3D graph. With the advance of deep learning, 3D graph neural networks (GNNs) have been developed to learn from 3D molecular graph data (Schütt et al., 2017; Satorras et al., 2021; Gasteiger et al., 2020; Liu et al., 2021; Gasteiger et al., 2021; Thomas et al., 2018; Liao & Smidt, 2022) . These models are data-hungry and necessitate a large amount of annotated training data to attain good performance. However, annotation usually consumes excessive manpower, which is particularly challenging for molecules, e.g., the commonly used density functional theory (DFT) for molecular energy computing (Hohenberg & Kohn, 1964) is very expensive, inducing the complexity of O(n 3 e ), where n e is the number of electrons. As a concrete example, DFT can be hundreds of thousands of times slower than a reasonably good GNN for inference (Gilmer et al., 2017) .Active Learning (AL) algorithms automatically identify the salient and exemplar samples from large amounts of unlabeled data (Settles, 2009; Ren et al., 2021) . This tremendously reduces the human annotation effort in inducing a deep neural network, as only the few samples that are identified by the algorithm need to be labeled manually. Further, since the deep network gets trained on the representative samples from the underlying data population, it typically depicts better generalization capability than a passive learner, where the training data is selected at random. Deep AL has been used with remarkable success in a variety of applications, such as computer vision (Yoo & Kweon, 2019; Sinha et al., 2019) , natural language processing (Zhang et al., 2022) , medical diagnosis (Blanch et al., 2017) , and anomaly detection (Pimentel et al., 2020) among others. However, a principled AL algorithm for 3D molecular graph learning is currently lacking.In this paper, we propose a powerful AL paradigm for 3D molecular graphs. We formulate a criterion based on uncertainty and diversity, which ensures that the queried molecules are the ones where the graph learning model has maximal uncertainty about the labels, and are also mutually diverse to avoid duplicate sample queries. In particular, diversity computing for 3D graphs is challenging and the difficulties are twofold. Firstly, the AL pipeline requires to compute the difference between any two 3D molecular graphs, which could have different planar molecules (entangling different atom numbers, etc) in most cases. Secondly, the 3D shape (geometry) of a 3D graph should be captured completely for expressive geometric representations for accurate diversity computing. To tackle these challenges, we propose a novel diversity sampling method for 3D molecular graphs built on distributions of important 3D geometries. We particularly propose a set of new 3D graph isometrics for 3D graph isomorphism. Our geometric modeling method is provably more powerful than the Geometric Weisfeiler-Leman (GWL) test (Joshi et al., 2023) in distinguishing and representing 3D graph geometries. This indicates our approach sets an upper bound to the expressive power of any existing equivariant geometric models. The moments of the distributions of the associated geometries (e.g., reference distances, triangles) are extracted for accurate and efficient diversity computing. In addition, to ensure our AL paradigm selects samples with maximal uncertainties, we carefully design a Bayesian geometric graph neural network specifically for 3D graph uncertainty computing. Our method is shown to be effective and efficient based on a set of ground approximations. With our novel components, we pose the sample selection as a quadratic programming (QP) to identify exemplar molecules to be annotated. Our method is easy to implement and can be applied in conjunction with any 3D GNN architecture. We conduct extensive experiments on four properties of the QM9 dataset. Results demonstrate the effectiveness of our AL paradigm, as well as the proposed diversity and uncertainty methods. 
diversity computing for 3d molecular graphs: 
 In molecular AL tasks, diversity sampling is important for eliminating redundancy, thereby wisely leveraging the annotation budget. The model's capability of capturing the 3D shape diversity among molecules is crucial for informed sampling. A particular challenge lies in that, a diversity measure for two 3D molecules with different planar graphs is indispensable. Methods for diversity measures for 3D molecules with the same planar graph have been developed (Kumar & Zhang, 2018; Kearnes et al., 2016; Gfeller et al., 2013) , but a diversity method for two 3D molecules with different planar graphs (entailing different atoms, etc) is demanding. Inspired by the USR method (Ballester & Richards, 2007) , we propose a novel solution to achieve the goal from the distribution perspective. Generally, we develop a set of new isometrics for expressive representations of 3D molecular graphs, after which the distributions of geometries associated with the isometrics are obtained for diversity computing.As the first step, we introduce a set of new isometrics as a basis, aiming at expressive representations of 3D graphs. As we focus on 3D geometry of molecules in this section, for simplicity, we use 3D point clouds to illustrate our ideas. Let A = {a 1 , a 2 , ..., a n } and B = {f (a 1 ), f (a 2 ), ..., f (a n )} be two sets representing 3D point clouds. Here, each a i in A is associated with a positional vector ⃗ a i = (x ai , y ai , z ai ) in 3D space. f denotes a bijective mapping between A and B. Then, similarly, each point f (a i ) in B is associated with a positional vector ⃗ f (a i ) = (x f (ai) , y f (ai) , z f (ai) ). Two 3D point clouds A and B are said to have isometric transformation when, given a group SE(3), ∃ γ ∈ SE(3) such that A = γB. We further choose or compute a reference point (e.g., centroid) for each point cloud, denoted as r 1 and r 2 , respectively. Without loss of generality, we use a far to denote the farthest point from the reference point in point cloud A. Below, we will define three levels of isometrics, each of which fulfills an isometrical mapping between A and B. To satisfy any isometric, there needs to exist a bijective function f : A → B, such that h f (a) = h a for any node a ∈ A. Here, h f (a) and h a denote the node feature vectors for f (a) and a, respectively.Reference Distance Isometric: If there exists a collection of global group elementsγ i ∈ SE(3), such that (r 2 , f (a i )) = (γ i r 1 , γ i a i ) for each point a i ∈ A, A is reference distance isometrical to B.Reference distance isometric involves the Euclidean distance between any atom in the molecule and the predefined reference point. i ∈ SE(3), such that (r 2 , f (a far ) , f (a i )) = (γ i r 1 , γ i a far , γ i a i ) for each point a i ∈ A, A is triangular isomet- rical to B.With reference point r, we define the reference vector ⃗ v 0 as r pointing to the farthest point a far in a 3D molecule. Based on reference distance isometric, triangular isometric further involves the angle between ⃗ v 0 and any other vector pointing from r to any other point in the molecule, computed asθ k = cos -1 ⃗ v0•⃗ v k ∥⃗ v0∥∥⃗ v k ∥, where ⃗ v k denotes vectors originating from r and directed towards k th atoms in the molecule. The process is illustrated in part A of Fig. 1 . For a molecule with N nodes, we compute N -1 angles. Essentially, such angles provide insights into the spatial arrangement of atoms with respect to the pre-assigned reference vector.Cross Angular Isometric: If there exists a collection of global group elementsγ ij ∈ SE(3), such that (r 2 , f (a j ) , f (a i )) = (γ ij r 1 , γ ij a j , γ ij a i ) for all a i , a j ∈ A (i ̸ = j), A is cross angle isometrical to B.Beyond the angles in triangular isometric as well as based on reference distance isometric, cross angular isometric further considers angles formed by any two atoms in the molecule with respect to the reference vector as above. Specifically, for every pair of atoms i and j, a vector ⃗ v ij is formed from i to j. With the reference vector ⃗ v 0 , the cross angle is computed asα ij = cos -1 ⃗ v0•⃗ vij ∥⃗ v0∥∥⃗ vij ∥ .This approach, as depicted in part B of Fig. 1 , essentially reflects the torsion information (angles between two planes) in a molecular structure. For a molecule with N nodes, we compute N (N -1)/2 cross angles with the complexity of O(N 2 ). However, the commonly used way to compute torsion angles involves 3-hop neighborhood in a graph, inducing a much larger complexity of O(N 3 ).Notably, in this work, we propose the aforementioned three isometrics to define a complete isometry space, leading to the computing of the associated geometries (distances, angles) in Euclidean space for molecular diversity. There exist other studies to delineate isometry space, such as the first and second fundamental forms (Gallier, 2011) . Compared with existing methods, the advantages of our solution are twofold. Firstly, our method represents a straightforward solution that can be naturally fitted into existing pipelines for molecular similarity computing. The commonly used methods, such as the Ultrafast Shape Recognition (USR) algorithm (Ballester & Richards, 2007) and Gaussian overlay-based methods (Rush et al., 2005; Hawkins et al., 2007) , are all developed based on geometries (distances, angles) in Euclidean space. Secondly, our method is more efficient in capturing the essential geometric features. Taking the first and second fundamental forms as an example, these methods involve complex computations of a surface's differential properties, which could be expensive, especially for non-smooth surfaces that are common in molecular data. Specifically, estimating these forms accurately may lead to computational complexities up to O(N 3 ), depending on the technique used for approximating surface derivatives and curvatures. The proof of Theorem 1 can be found in Appendix A.2. Generally, we define three levels of isometrics for graph isomorphism. Reference distance isometric ensures that the Euclidean distance between each point and a predefined reference point is consistent in two different point clouds.Triangular isometric further manifests the spatial arrangement of atoms referring to the pre-assigned pivot. Built on Triangular isometric, cross angular isometric then reflects the spatial torsion angle information. An illustrative example for triangular isometric and cross angular isometric is also given in Fig. 2 . Clearly, cross angular isometric represents the strictest isometry among the three. In the following Sec. 2.1.2, we show that a designed geometric representation based on cross angular isometric can exhibit great expressive power.In this section, we aim to formally elucidate the expressive power of a geometric representation based on our developed isometrics in Sec. 2.1.1. Generally, we use GR to denote any geometric representation to capture molecular structures, and function ζ be the mapping function associated with GR. We explore the Geometric Weisfeiler-Leman (GWL) test (Joshi et al., 2023) , and then leverage GWL to measure the expressiveness of a GR. GWL test is an extension of the classic WL Test, enhancing its capabilities by incorporating both the topological structure of the graph and the geometric attributes of its vertices. Such an integration allows the GWL test especially apt for evaluating all 3D graph representation methods. Similar to the regular WL test, GWL test imposes an upper bound to the expressive power of 3D GNNs, i.e., if GWL test fails to distinguish two 3D graphs, then all existing 3D GNNs would also fail. See details of the GWL test in Appendix A.1.Let the relation A ≇ B indicate that two structures are not reference distance isometric in terms of their 3D geometries. We propose Theorem 2 as follows: Theorem 2. A geometric representation GR has the same expressive power as the Geometric Weisfeiler-Leman (GWL) Test in distinguishing non-isomorphic molecular structures if the associated mapping function ζ can map any two molecular structures A and B into two distinct representations(i.e., ζ (A) ̸ = ζ (B)) if and only if A ≇ B.The proof of Theorem 2 is provided in Appendix A.2. Next, we use GR ours to denote the geometric representation based on the isometrics developed in Sec. 2.1.1, From Theorem 1, we can easily conclude that cross angular isometric is the strictest isometric among the three. Naturally, we formulate GR ours as a set containing all reference distances, triangles, and cross angles in a 3D graph. We show GR ours adheres to the constraints stipulated in Theorem 2 through the Proposition below. The proof of Proposition 1 is provided in Appendix A.2. Based on both Theorem 2 and Proposition 1, it is evident that our geometric representation GR ours possesses an expressiveness at least as powerful as the GWL test. Moreover, through the proposition below, we rigorously show GR ours possesses greater expressive power than the GWL test. Proposition 2. GR ours is strictly more expressive than the Geometric Weisfeiler-Leman (GWL) test in distinguishing non-isometric point clouds.The proof of Proposition 2 can be found in Appendix A.2. In conclusion, the molecular geometric representation method GR ours developed in this work has the greater expressive power than the Geometric Weisfeiler-Leman (GWL), which indicates our diversity sampling method is accurate enough to capture the 3D shape diversity among different molecules. Notably, as mentioned before, GWL test sets the upper bound to any existing 3D GNNs. Apparently, our geometric representation method GR ours is provably more expressively powerful than any existing 3D GNN for learning geometric representations. Essentially, the three isometrics associated with GR ours define expressiveness at different levels. For example, as only considering distance information, a well pretrained SchNet is upper bounded by reference distance isometric (but not triangular isometric or cross angular isometric); as a more powerful model than SchNet, a well pretrained DimeNet is upper bounded by triangular isometric (but not cross angular isometric). Additionally, accurate geometric representation learning requires a perfect pretrained 3D GNN model, which is also hard to guarantee in practice. Our isomorphy study provides a model-agnostic diversity computing solution, avoiding the need of a perfectly pretrained 3D GNN model, as well as achieving a theoretically guaranteed upper bound of the expressiveness of all existing 3D GNN models.Based on the isomorphy study in Sec. 2.1.1, we obtain our geometric representation method GR ours and prove GR ours possesses greater expressive power than any existing 3D GNN models in Sec. 2.1.2. In this section, we aim to extract the distributions of the entangled three geometries in GR ours , including reference distances, triangles, and cross angles, for diversity computing. Fortunately, we have the theorem (Hall, 1983) implying that the sequence of translated moments can be used to determine the original distribution. Following the USR work (Ballester & Richards, 2007) , for each of the three aforementioned geometries, we also use four reference points to reflect the "translated" geometries; those are, the centroid that is denoted as ctd and computed by the mean position of all the atoms in the 3D molecule, the point closest to the centroid (denoted as cst), the point farthest from the centroid (denoted as fct), and the point farthest from fct (denoted as ftf). For each reference point, we use a set of moments, including mean, variance, skewness, and kurtosis, which describe a distribution from different angles, e.g., skewness indicates the asymmetry and kurtosis describes the tailedness of a distribution. Detailed formulae for these moments can be found in the Appendix A.3. Notably, we compute these translated moments for all three entangled geometries as above. Eventually, we obtain summarized representations of distributions over geometries of 3D graphs, capturing essential characteristics of a molecule's shape.We use cross angles as an example to describe the final distributional vector. For a molecule with N atoms, as shown in Fig. 1 , we can obtain a set of cross angles[α ref ij ] N (N -1)/2i̸ =j,0<i,j<N for a reference point (e.g., ctd). After applying statistical moments as an approximation, we can obtain a 4-dimensional vector--→ M ca ref = [m ca ref , v ca ref , s ca ref , k ca ref ],where the four elements denote the mean, variance, skewness, and kurtosis for this reference point, respectively. We perform a similar process for all four reference points mentioned above. By doing this, we can obtain four 4-dimensional vectors including--→ M ca ctd , --→ M ca cst ,--→ M ca fct , and --→ M ca ftf , which are then concatenated together, resulting in the final 16-dimensional vector to represent the distribution of cross angles. We repeat the similar process for reference distances and triangles, and then all three corresponding 16-dimensional vectors are further concatenated as a 48-dimensional distributional vector to represent the geometric information of the input molecule. The 48-dimensional distributional vectors are then used to compute the diversity matrix as in Sec. 2.1.3. Finally, for any two molecules n 1 and n 2 in the dataset with N molecules, we perform the inner product on their distributional vectors to achieve the similarity, and then use 1similarity to obtain the final value D n1n2 as the diversity measure between them. Finally, a matrix D ∈ ℜ N ×N is obtained, which contains the diversity between every pair of molecules.In Sec. 2.1, we develop an effective method for diversity computing among different 3D molecular graphs. However, as mentioned in Sec. 1 and Sec. 3, in addition to selecting diverse molecules, it is important to select molecules where the model has maximal prediction uncertainty about the labels, so as to append maximal information to the model. To this end, we develop a principled pipeline to compute uncertainties for 3D molecular graphs in this section. Uncertainty qualification is well studied in planar graph analysis (Hirschfeld et al., 2020) , but an effective and principled paradigm for 3D molecular graphs is currently lacking. Additionally, existing methods, such as Bayesian neural networks (BNNs) (Lampinen & Vehtari, 2001; Titterington, 2004; Goan & Fookes, 2020) and deep model ensemble methods (Lakshminarayanan et al., 2017; Huang et al., 2017) , are excessively computationally expensive, limiting their capacity in 3D graph analyses. In this work, we develop an effective and efficient method, known as Bayesian geometric graph neural network (BGGNN), that takes a 3D graph as input and produces the demanding properties as well as uncertainty values, e.g., mean and variance.Formally, a 3D graph is represented as G = (V, E, P ), where V denotes the set of vertices (atoms), E denotes for the set of edges (bonds), and P stands for the set of Cartesian coordinates for each atoms. A 3D molecular graph is associated with a set of properties, denoted as O. Recently, researchers have developed 3D GNNs, such as SchNet (Schütt et al., 2017) , DimeNet (Gasteiger et al., 2020) , SphereNet (Liu et al., 2021) , and GemNet (Gasteiger et al., 2021) , for 3D graph representation learning. The likelihood of a 3D GNN can be represented as p 3DGNN (O | G, w), where 3DGNN indicates any existing 3D GNN and w denotes the set of parameters of the used 3D GNN. We also use p 3DGNN (w) to represent the prior distribution for the parameters. Assume we collect a new input and output pair, denoted as g * and o * . Then based on the conventional Bayesian theorem, Bayesian inference for this new output o * is given by pEQUATIONwhere R n is the whole space of n parameters in 3DGNN. It's infeasible to perform the above integration on R n due to prohibitive computational cost. To tackle this, the variational inference method is introduced to approximate p 3DGNN (O | G, w) with the parameterized q θ (w) through minimizing the Kullback-Leibler (KL) divergence between these two distributions. After applying Bayesian theorem once more, the minimization objective becomesEQUATIONTo completely avoid the integration over the whole parameter space, the MC-dropout method (Gal & Ghahramani, 2016; Srivastava et al., 2014) is further used in our BGGNN. Specifically, it employes the Monte-Carlo estimator (Gal et al., 2016; Gal & Ghahramani, 2016) to approximate the integration by performing summation over the sampled models. In practice, researchers implement an MC-dropout network by using dropout as the network's regularization (Gal & Ghahramani, 2016) . Following this, we propose to insert dropout layers after the linear layers in our used 3DGNN as an effective yet efficient estimation of Bayesian inference.Now as we have obtained the variational predictive distribution of a new output with q θ (w), we can easily compute the predictive mean and variance of this distribution. For the molecular property prediction tasks, after we sample N outputs from the same input, the heteroscedastic predictive uncertainty is then given byEQUATIONwhere ô * n is the n th sampled output and σ 2 n is the variance that is the same among all the data samples. By doing this, we can obtain an uncertainty value (variance) for each molecule. Additionally, built on a 3D GNN, our BGGNN can faithfully produce a set of molecular properties O.Practically, any of the existing 3D GNN can be used as the backbone network for property prediction and uncertainty computing. In this study, we employ SphereNet (Liu et al., 2021) as our 3DGNN, owing to its completeness in incorporating 3D geometric information. We apply dropout layers onto the linear layers of SphereNet for Bayesian inference in our BGGNN. To allow more accurate AL selections, we particularly employ the concrete dropout with a learnable dropout rate (Gal et al., 2017) in our BGGNN. Overall, our method is shown to be an effective and efficient paradigm for 3D graph uncertainty computing, as further empirically demonstrated in Sec. 4.EQUATIONA schematic diagram of our active sampling framework is depicted in Fig. 6 and described in A.4 in Appendix. Specifically, in Sec. 2.1, we obtain the matrix D ∈ ℜ N ×N containing the mutual diversity between every pair of unlabeled molecules, where N is the number of unlabeled molecules. In Sec. 2.2, we employ our designed BGGNN to achieve the vector r ∈ ℜ N ×1 quantifying the prediction uncertainty score of each unlabeled molecule. In the AL setting, our objective is to select a batch of k unlabeled molecules (where k is the pre-determined query batch size) with high prediction uncertainty and high mutual diversity among them. Let z ∈ {0, 1} N ×1 be a binary vector with N entries which denotes whether the unlabeled molecule x i will be included in the batch (z i = 1) or not (z i = 0). The molecule selection can thus be posed as the following optimization problem as in Eq. 4, where λ is a weight parameter governing the relative importance of the two terms. This is a standard quadratic programming (QP) problem; we relax the integer constraints into continuous constraints and solve the problem using an off-the-shelf QP solver. In this work, we employ the widely used Operator Splitting Quadratic Program (OSQP) (Stellato et al., 2020) to solve the QP problem in Eq. 4. We then apply a greedy approach to project the continuous solution back to the binary space, where the k highest entries of the continuous solution vector are set to 1 and the remaining to 0. Such an approach is commonly used to convert continuous solutions obtained from a QP solver to binary solutions in AL Chattopadhyay et al. (2013) ; Wang & Ye (2013) . Notably, the predictions in main tasks (e.g., molecular properties) are produced by our BGGNN built on SphereNet as in Sec. 2.2.Assume A and B are Cross Angular Isometric, which means for all a i , a j ∈ A with i ̸ = j, there exists a collection of global group elementsγ ij ∈ SE(3) such that (r 2 , f (a j ), f (a i )) = (γ ij r 1 , γ ij a j , γ ij a i ).From this assumption, for any point a i ∈ A, there exists a corresponding γ i ∈ SE(3) (which is one of the γ ij where j corresponds to the farthest point a far ) such that (r 2 , f (a far ), f (a i )) = (γ i r 1 , γ i a far , γ i a i ). This is precisely the condition for A and B to be Triangular Isometric. Hence, if A and B are Cross Angular Isometric, then A and B are also Triangular Isometric. Proof. 1. GWL Distinguishable implies GR Distinguishable: Suppose two molecular structures, A and B, are distinguishable by the GWL Test after k iterations. That is to say, the multiset of node colors remains the same through iteration 0 to k -1 but becomes different in the k th iteration. According to the GWL Test in the Graph Compasison step, there exist a 0 ∈ A and f (a 2. GR Distinguishable implies GWL Distinguishable: Suppose that, at the k th iteration, two molecular structures A and B are distinguishable by the representation method GR. This implies that the mappings ζ (A) and ζ (B) are different by being non-Reference Distance Isometric according to the constraints on GR in the theorem. So we are able to find at least a 0 ∈ A, such that (r 2 , f (a 0 )) ̸ = (γr 1 , γa 0 ) holds for ∀γ ∈ SE(3) which can be colored differently by the GWL Test when g (t) i encodes geometrical information of its neighbors. So GWL Test is also able to distinguish between A and B. Proposition 2. GR ours is strictly more expressive than the Geometric Weisfeiler-Leman (GWL) test in distinguishing non-isometric point clouds.EQUATIONProof. First we prove that GR ours satisfies the two constraints as below.1. GR ours can distinguish any two point clouds A and Bby cross angles; that is,ζ (A) ̸ = ζ (B),if A and B are reference distance isometric but are non cross angular isometric.2. The mapping function ζ associated with GR ours is injective.Then we would prove that with the constraints satisfied, GR ours would be more expressive than GWL Test.1. GR ours satisfies the two constraints:1. A and B are reference distance isometric implies that there exists a collection of global group elements γ i ∈ SE(3) such that (r 2 , f (a i )) = (γ i r 1 , γ i a i ) for each point a i ∈ A.Meanwhile, A and B are non cross angular isometric implies that ∃ a i , a j ∈ A (i ̸ = j), such that (r 2 , f (a j ) , f (a i )) ̸ = (γr 1 , γa j , γa i ) holds for ∀γ ∈ SE(3). Given GR ours introduced in Sec. 1. For a given iteration k in the GWL Test, the primary cause of failure arises when, despite the point clouds not being entirely isometric, there exist local structures or geometrical features that are sufficiently similar. In such case, they might be assigned the same node color partition even after the k th iteration.2. When the color update scheme in GWL Test depends only on local area scalar features which are the same between two point clouds, the test will thus assign the two point clouds the same node color partition.3. For GWL Test without a pre-assigned iteration number, it may also fail as a result of Hash collision. Different inputs may be mapped to the same Hash value and thus result in the same node color partition between different point clouds.These are the cases that GWL Test may fail to distinguish two distinct point clouds. Since there exist point clouds that GWL Test cannot distinguish while GR ours can, we can draw the conclusion that GR ours is more expressive than the GWL Test. 
related work: 
 3.1 ACTIVE LEARNING Active Learning (AL) is a well-researched problem in the machine learning community (Settles, 2009) . There exist two commonly used strategies for AL sampling. Uncertainty based sampling queries unlabeled samples with the highest prediction uncertainties for annotation. Diversity/representativeness based sampling aims to select the subset that can well represent the entire data distribution. A full review of the two AL sampling methods is provided in Appendix A.5.Molecular shape similarity plays a pivotal role in drug discovery and virtual screening of compounds (Kumar & Zhang, 2018; Murgueitio et al., 2012; Shang et al., 2017) . Methods predominantly fall into several categories (Kumar & Zhang, 2018) , including descriptor-based methods (Schreyer & Blundell, 2012; Cannon et al., 2008; Li et al., 2016; Armstrong et al., 2009; Zhou et al., 2010) , atom-centered gaussian-based methods (Haque & Pande, 2010; de Lima & Nascimento, 2013; Yan et al., 2013) , surface-based methods (Hofbauer et al., 2004; Mavridis et al., 2007; Cai et al., 2012; Karaboga et al., 2013; Venkatraman et al., 2009; Sael et al., 2008) , etc. Descriptor-based methods are notably represented by the Ultrafast Shape Recognition (USR) algorithm (Ballester & Richards, 2007) , which uses statistical moments of the distance distribution to characterize molecular shapes. Gaussian overlay-based methods, with ROCS (Rush et al., 2005; Hawkins et al., 2007) being the most commonly used one, evaluate the maximum volume overlap between two molecules. Surface-based methods typically employ shape signatures (Zauhar et al., 2013) or shape histograms to delineate molecular surfaces for shape similarity assessment. Despite the progress, a principled and theoretically ground similarity method for 3D molecular graphs is currently lacking.We perform experiments on four properties of the QM9 (Ramakrishnan et al., 2014) benchmark dataset: mu, alpha, homo, and lumo. These properties have continuous values, making the prediction problem a regression task. We randomly divide the training set of 110, 000 molecules into three splits of size 25, 000 each. From each split, we randomly select 5, 000 molecules as the initial labeled set and the remaining 20, 000 molecules as the unlabeled set. In each AL iteration, we query 1, 500 molecules from the unlabeled set, which are labeled and appended to the labeled set. The model's performance is evaluated on a held-out validation set containing 10, 000 molecules. We save the best-performing model on the validation set and report its performance on the test set containing 10, 831 molecules. The process is repeated for 7 AL iterations, which is taken as the stopping criterion. The final results are averaged over the three splits to rule out the effects of randomness. The value of λ in Eq. 4 is taken as 1, for the QP problem. The Mean Absolute Error (MAE) is used as the evaluation metric in this work. Implementation Details: We use SphereNet (Liu et al., 2021)  
active learning performance: 
 The active learning performance results are depicted in Fig. 3 . In each graph, the x-axis denotes the iteration number and the y-axis denotes the MAE on the test set. Our analysis revealed that Evidential Uncertainty depicted the worst performance and furnished significantly high error values for all the four properties, which obscured the difference in performance among the other methods in the plots.For better interpretation and understanding, we exclude the Evidential Uncertainty method from the plots here and present the results with this baseline in Sec. A.6 of the Appendix. The other baseline methods depict more or less similar performance, with Coreset marginally outperforming the other baselines. Our method comprehensively outperforms all the baselines. At any given AL iteration, it consistently attains a lower MAE compared to all the baselines. It also attains the least MAE after all the AL iterations, for all the four properties studied. We also conducted statistical tests of significance using paired t-test to assess whether the improvement in performance achieved by our method is statistically significant. For this purpose, we compared the average MAE achieved by our method against each of the baselines individually. The results are reported in Table 1 ; each entry in the table denotes the p-value of the paired t-test between our method and the corresponding baseline (denoted in the columns) for the property studied (denoted in the rows). From the table, we note that the improvement in performance achieved by our method is statistically significant (p < 0.05) compared to all the baselines, consistently for all the four properties studied. These results unanimously corroborate the promise and potential of the proposed active sampling method to tremendously reduce the annotation cost in inducing a robust 3D graph neural network for molecular property prediction. The goal of this experiment is to study the effect of query budget (batch size) on the AL performance.The results on the mu property for query budgets 1, 000, 1, 500 and 2, 000 are depicted in Fig. 4 . Since Evidential Uncertainty depicted much worse performance than all the methods, it was excluded from this comparison. Our framework once again outperforms all the baselines consistently for all the query budgets and attains the least MAE after all the AL iterations. As before, we conducted a paired t-test to estimate the statistical significance of the obtained performance improvement. The results are presented in Sec. A.7 of the Appendix. From the p-values, we conclude that the error values furnished by our method are statistically significantly better (p < 0.05) than all the baselines, consistently for all the query budgets.This shows the robustness of our framework to the query budget. These results are particularly significant from a practical standpoint as the available query budget in a real-world application is dependent on time, resources, and other constraints of the application.We conduct ablation studies to examine the power of our diversity computing method, as it is our primary contribution in this research. We perform experiments on the mu and lumo properties from two aspects. Firstly, we compare our framework with only the diversity term in Eq. 4 against Coreset, the state-of-the-art diversity-based AL technique. The results are reported in Fig. 5 , from which we note that the diversity component of our framework consistently furnishes much lower MAE values than Coreset over all the AL iterations, for both properties. Secondly, we also conducted experiments where we compared the performance of our overall framework (using both uncertainty and diversity) against the baseline where only the uncertainty term in Eq. ( 4) was used for active sampling. The results revealed that removing the diversity term adversely affected the performance of our framework. A paired t-test revealed that the improvement in performance achieved by our diversity component is statistically significant (p < 0.05) for both these properties (p = 0.0001 for mu and p = 0.04 for lumo). These results show the usefulness of the proposed diversity metric in developing an AL framework to train a 3D GNN for molecular property prediction. In this experiment, we analyze the computation time of all the methods studied in this paper. The average time taken to query a batch of unlabeled samples and update the SphereNet model (one active learning iteration) are shown in Table 2 . For fair comparison, all the methods were run on the same NVIDIA RTX A4500 20GB GPU. We note that Random, Learning Loss and Evidential Uncertainty all have similar computation time. Coreset depicts the highest computation time, as it needs to solve a mixed integer programming (MIP) problem. The computation time of our framework is marginally less than Coreset, and approximately double that of the other three. Our analysis revealed that solving the QP problem is the main bottleneck of our method. Hence, we explored a strategy to improve the computation overhead of solving the QP problem by running it in the GPU (instead of the CPU) using the parallel implementation of the alternating direction method of multipliers, as detailed in Schubiger et al. (2020) . This results in a substantial reduction of the computation time, as depicted in Table 2 (Ours (Fast QP)). The AL performance using the fast QP solver is very similar to that obtained using the original solver, and is presented in Fig. 8 in Sec. A.8 of the Appendix. Apparently, the performance studies in both Sec. 4.2 and Sec. A.8 show that our framework is much more accurate than these baselines. Given the large margin of performance improvement, we think the efficiency of our methods is acceptable. As solving QP problems is not the main technical novelty of this work, we will continue exploring faster QP solvers from literature for engineering purposes as future work.The procedure terminates when colors do not change from the previous iteration or a predetermined maximum number of iterations is reached. 5. Graph Comparison: Finally, two geometric graphs G and H are geometrically nonisomorphic if there exists some iteration t for which the sets of colors of their vertices are not equal, i.e., c EQUATIONγ i ∈ SE(3) such that (r 2 , f (a far ) , f (a i )) = (γ i r 1 , γ i a far , γ i a i ) for each point a i ∈ A.From the above assumption, it is clear that for each point a i ∈ A, there exists a corresponding γ i ∈ SE(3) such that (r 2 , f (a i )) = (γ i r 1 , γ i a i ) which is the condition for A and B to be Reference Distance Isometric. Hence, the statement is proved. 
conclusion and future work: 
 In this paper, we present a novel active learning framework with the goal of reducing the annotation cost for learning from 3D molecules represented as 3D graphs. The sample selection is posed as a QP problem, which selects samples with high mutual diversity and high uncertainty. Novel diversity and uncertainty components are proposed specifically for 3D graphs, with strong empirical results demonstrating the promise and potential of our AL framework. As part of future work, we plan to apply our methods to molecular analysis where much more accurate but expensive annotation is required, such as computing ground states of molecular systems using the Schrödinger equation. DFT calculations are widely used in many practical applications but still involve approximations, as Schrödinger equation is prohibitively expensive and its use is limited in very small molecules. Our AL pipeline is anticipated to unleash greater potential in such extreme-scale applications. Additionally, given AL needs several interactions with each requiring the model is well trained, we test our methods on the commonly used but medium-scale QM9 dataset in this work. Even though we think the empirical studies are sufficient to support our theory, we still plan to test the scalability of our methods on large-scale molecule datasets, such as OC20 (Chanussot et al., 2021) , in the future. 
 
 Appendix: A.1 GEOMETRIC WEISFEILER-LEMAN (GWL) TESTFor the Geometric Weisfeiler-Leman (GWL) test, consider a graph G with its set of vertices represented as V(G) and its set of edges as E(G). A vertex in graph G is denoted by i, and N i signifies the set of vertices adjacent to i. The color of vertex i at iteration t is given by c (t)i , and the geometric object for vertex i at iteration t is represented by g (t)i . The procedure for the GWL test is as follows:1. Initialization: Each vertex i is assigned an initial color c (0) i and a geometric object g (0)i , typically based on its local property or geometric attributes. 2. Iterative Aggregation: For each iteration t ≥ 1, the geometric object of each vertex i is updated to aggregate geometric information from its t-hop neighborhood, represented as g(t)i , which includes the colors and geometric objects from the previous iteration of vertex i and its neighbors. 3. Color Update: The color of each vertex i at iteration t is computed by aggregating the geometric information around vertex i using a G-orbit injective and G-invariant function, denoted by I-HASH, i.e., ci := I -HASH (t) g (t) i.The equations that we used for calculating four moments are as follows.The mean, often referred to as the average, represents the sum of all data points divided by the number of data points and is given byEQUATIONA.4 SCHEMATIC DIAGRAM OF OUR FRAMEWORK A schematic diagram of our active sampling framework is depicted in Fig. 6 . We are given a labeled training set L, an unlabeled set U and a query budget k for each active learning iteration. The SphereNet model is first trained on the labeled set L. In the second step, the trained model is applied on the unlabeled set to compute a prediction uncertainty of each unlabeled molecule, which is used to populate the uncertainty vector r; the diversity matrix D is also computed in this step where D(i, j) is the diversity between unlabeled molecules x i and x j . Next, the QP problem is solved to select k unlabeled molecules for annotation. These molecules are removed from the unlabeled set U and appended to the labeled set L. The active sampling process is continued iteratively until some stopping criterion is satisfied (taken as 7 iterations in our work). Note that, computing the diversity matrix D in Step 3 needs to be executed just once for the whole process. Once we have the initial D, as more and more samples are queried through AL, we keep deleting the corresponding rows and columns from D to derive the updated matrix. (Settles, 2009) . Uncertainty sampling is an important strategy for AL, where unlabeled samples with the highest prediction uncertainties are queried for annotation. Several techniques have been explored to compute the uncertainty, such as Shannon's entropy (Guo & Schuurmans, 2007; Li & Guo, 2013) , the distance of a sample from the separating hyperplane for SVM classifiers (Tong & Koller, 2001) , the disagreement among a committee of classifiers regarding the label of a sample (Freund et al., 1997; Gilad-Bachrach et al., 2005) , among others (Hoi et al., 2006; HOI et al., 2008; Guo & Greiner, 2007; Freytag et al., 2014) . With the advent of deep learning, Deep AL has attracted significant research attention (Hino, 2020; Ren et al., 2021) , Entropy-based methods are developed as well (Wang & Shang, 2014; Ranganathan et al., 2017) . Yoo & Kweon (2019) cascaded a task-agnostic loss learning module that queries samples with the highest predicted loss values. Huang et al. (2021) proposed a strategy based on temporal output discrepancy. Techniques based on adversarial training have also been explored (Sinha et al., 2019; Mayer & Timofte, 2020; Zhang et al., 2020; Zhu & Bento, 2017) . Bayesian neural networks (BNNs) (Lampinen & Vehtari, 2001; Titterington, 2004; Goan & Fookes, 2020) and deep model ensemble (Lakshminarayanan et al., 2017; Huang et al., 2017) generally achieve superior performance but may induce excessive computational cost.Diversity/representativeness based AL sampling has also been exploited. A core-set sampling technique proposed by Sener & Savarese (2017) queries a batch of samples such that a model trained on the queried subset is competitive for the remaining data samples. Diversity sampling has also been exploited in the context of Bayesian neural networks (Kirsch et al., 2019) . Buchert et al. (2023) uses 


Decision: This paper presents an active learning method for 3D molecular graphs, where the diversity of molecules are measured according to 3D graph isomorphism analysis. Experimental results on two data sets prove the effectiveness of the proposed approach. 

In general, the reviewers like the novelty of the proposed approach. There are a few aspects that the authors need to address before publication: (1) the claims are too strong; (2) more baselines should be included, e.g., 2D-graph based approaches.

Reviews: review: paper_summary: paper_summary: This paper describes a way of sampling 3D graphs for active learning on
molecules, leveraging isometries.
strength_weakness: strength_weakness: I think the main weakness of this work is that it needs stronger connections
to the actual chemistry.

* The actual elements are important. Two molecules can have similar
geometries, yet very different electronic interactions that give rise to very
different properties.

* The paper would benefit from explaining the relationship between the
different isometries and actual molecules and chemical properties.

* The space of molecules is very large.  The tests hold back part of known
datasets as unlabeled molecules.  But the most useful new labels may not be
in the QM9 dataset, but may be in some other dataset, or even a novel
molecule.  Thus, a more interesting algorithm would be one that identified a
molecule X as as important new molecule to add, allowing some search for it.
Of course, this would raise a number of other important issues, such as
generating the new molecule X, etc.

* The paper does not test against 2-D methods.  Not sure why.  Perhaps they
perform well.
questions: questions: * How does sampling the graphs relate to sampling the interactions that are
responsible for properties?

* Is chirality considered?

* Why do the graphs in Figure 3 not start at the same loss?  Before any active
learning has occured, shouldn't that the situation?

* Are the other AL techniques also designed to be used in this iterative manner?
If not, I think a better baseline test would be to use the same number of
samples.
limitations: limitations: The paper does not adequately address limitations. I would imagine one would be that it does not consider the actual elements of each atom. In other words, two molecules with similar geometries are considered similar, even though they could be very different chemically.

What about periodic structures?

Also, what if there is not a pre-existing source of valid, but unlabeled molecules? A common situation is that there are only a small number of molecules known at all.
paper_summary: paper_summary: The paper proposes an active learning scheme for molecular property prediction using uncertainty estimates from Dropout Monte Carlo and diversity metrics. In each active learning iteration, molecules are selected by maximizing uncertainty and diversity in the batch by solving a quadratic programming problem. The authors propose a 48-dimensional vector that encodes statistical moments of reference distances, angles and cross-angles for 4 reference points. Experiments on QM9 and MD17 datasets are performed to benchmark the effectiveness of the proposed method against several baselines.
strength_weakness: strength_weakness: 1.	Many of the claims in the paper are too bold, neglecting several important related works in the field:
     - There are several active learning applications for 3D GNNs already in the literature, especially in the neural network potential literature:  https://www.nature.com/articles/s41467-021-21376-0 ,   https://www.nature.com/articles/s41524-023-01104-6 
     - In the review of alternative 3D structural descriptors, the most common methods such as SOAP ( https://journals.aps.org/prb/abstract/10.1103/PhysRevB.87.184115  ), or ACE ( https://journals.aps.org/prb/abstract/10.1103/PhysRevB.99.014104 ) are neglected. 
     - Unlike claimed by the authors, the proposed Dropout Monte Carlo scheme for active learning in molecular property prediction with 3D GNNs is not novel, but has been proposed previously: https://www.nature.com/articles/s41524-024-01277-8 

2.	The considered baselines might not be the most relevant to the method and do not allow to attribute the origin of the empirical performance increase given that neither Coreset nor Learning loss uses a Bayesian uncertainty. BatchBALD would be a much more relevant baseline given that the same uncertainty estimates can be used. This would also allow to investigate the effect of the quadratic programming formulation against the greedy clustering-based approach of BarchBALD. It would also allow to investigate the benefit of the proposed diversity metric vector against, for example, classical descriptors such as SOAP. This would allow to clarify the underlying source of the outperformance: Is it the Bayesian UQ estimate, the quadratic programming formulation of the proposed novel structure descriptor vector?

3.	The quadratic scaling of the computational cost of the method is concerning. The pool sizes considered in the experimental section are very small (15,000) in the field of 3D molecular property prediction, where datasets can easily span several million conformations (PubChemQC, ANI-2x, GEOM). Given that there is already a computational overhead visible over competing AL methods, this difference will only increase for realistically-size datasets.
questions: questions: Does the 48-dimensional vector contain any information about the atom-species involved in distances, angles and cross angles or is this metric purely geometry-based?
limitations: limitations: The authors discuss some limitations of the proposed approach, including the scaling to large datasets such as OC20. However, the scaling to large atom sizes could be discussed more as requiring further research given that it is non-obvious that the statistical moments measured at 4 reference points are sufficient to distinguish protein-size geometries or large crystal structures. Additionally, it should be stated more clearly that the proven power of the proposed descriptors exceeding the GWL test only holds before compressing the descriptors into the 48-dimensional vector.
paper_summary: paper_summary: This paper introduces a principled active learning (AL) paradigm tailored for molecular learning. The proposed AL approach aims to alleviate the hurdle of human annotation by automatically querying labels for the most informative samples. The authors treat molecules as 3D molecular graphs and they introduce a set of new 3D graph isometries for 3D graph isomorphism analysis, which are shown to be as expressive as the Geometric Weisfeiler-Lehman (GWL) test. To ensure the selection of samples with maximal uncertainties, the authors design a Bayesian geometric graph neural network specifically for 3D molecular graphs. Active sampling is formulated as a quadratic programming (QP) problem integrating these components. Experimental results demonstrate the effectiveness of the proposed AL paradigm, highlighting the advantages of the diversity and uncertainty methods introduced.
strength_weakness: strength_weakness: 1. **Overly Strong Claims**: Some of the claims made in the paper are quite strong and require more careful justification. For detailed points of concern, please refer to the Questions section. Providing additional evidence or more nuanced discussions could strengthen these claims.

2. **Need for More Detailed Experimental Insights**: While the experimental section is comprehensive, it could benefit from additional insights and details. Specific experiments would be clearer and more informative with further explanation. Providing more context and interpretation of the results would enhance the reader's understanding and the overall impact of the findings.
questions: questions: 1. **Clarification Needed on Line 175**: The sentence starting in line 175 needs to be more carefully justified. Specifically, how can you assert that SchNet is upper bounded by distance isometry but not by triangular and cross-angular isometries? Please provide a more detailed explanation or supporting evidence for this claim.

2. **Generalization in Section 4.2**: It is not clear how the experiment in Section 4.2 demonstrates generalization. Could you provide more information and insight on the choice of graphs used in this section? Additionally, explaining the rationale behind the experimental setup and how it supports your conclusions would be helpful.

3. **Complete Results for Section 4.4**: Please report all the results from Section 4.4.
limitations: limitations: Limitations addressed
paper_summary: paper_summary: This paper proposes a principled AL paradigm to alleviate the annotation hurdle of 3D molecular graphs. It introduces a novel diversity component for 3D molecular graphs, which is provably at least as expressive as the GWL test. Furthermore, the authors develop an effective and efficient pipeline to compute uncertainties for 3D molecular graphs rooted in Bayesian inference.
strength_weakness: strength_weakness: + The paper lacks some necessary explanations, such as those for "USR" and "GSL," making it difficult to understand.
+ Experiments were conducted only on the QM9 and MD17 datasets, so the effectiveness and efficiency on larger datasets remain unknown.
questions: questions: + The authors formulate a criterion based on uncertainty and diversity. How do these two factors respectively impact performance?
limitations: limitations: Please see the weaknesses part.

Comments: title: Follow up response to Reviewer yNCZ
Completely agree that 3D is important
We are eagerly awaiting your response
Thank you.
The reviewer-author discussion stage is ending soon
The discussion period is closing in 32 hours
We are looking forward to your response
Reply to Reviewer UZqH (Part III)
Reply to Reviewer UZqH (Part II)
Reply to Reviewer UZqH (Part I)
References to Rebuttal
comment: Thanks for your comment and discussion! While we highly respect your thought to maintain the rating, we still hope you can reconsider if other aspects in your initial set of comments have been clarified and addressed. 

----------------------------------
Regarding this new comment, we provide clarifications below:


1) **We do consider elements in $\mathbf{G}$ in line 227.** In $\mathbf{G}=(V, E, P)$, $V$ denotes the set of elements (atom types), **as mentioned in the same line**. Particularly, this is a widely adopted setting in mainstream 3D GNN models for molecular learning [1,2,3,4].

2) A typical active learning process includes two stages: the selection stage (selecting the most informative molecules to annotate, to reduce the annotation budget), and representation learning (like molecular property prediction and molecular dynamic simulation). **Elements are considered in both stages.**

3) The selection stage (Eq. (4) in the paper, lines 266-272) combines two strategies (diversity and uncertainty) to select the most informative molecules. **Elements are considered in the uncertainty part.** Following classical geometric descriptors, diversity focuses on geometric descriptors. **Hence, the overall selection stage (Eq. (4), where *r* denotes uncertainty scores for molecules) considers elements.**

4) For the example you gave - two molecules with the same graph but different elements, $\mathbf{G}$ in line 227 would be different, thus uncertainty would be very different, through Eq. (4)  the informativeness of these two molecules are totally different. 


5) Last but not least, for our diversity descriptor, we included a comparison with a well-known geometric descriptor in chemistry, the SOAP descriptor [5, 6, 7], which produces descriptors that describe local atomic environments using spherical harmonics and radial basis functions. **SOAP considers both geometric information and elements (species).** Our results reveal that our diversity component outperforms SOAP; our overall method (using both diversity and uncertainty as in Eq. (4)) also outperforms the SOAP descriptors (using both SOAP diversity and uncertainty). This outperformance can be attributed to the local nature of SOAP descriptors. Our work is among the first to consider a global 3D geometric descriptor for molecular learning, combined with an uncertainty component that considers chemical contexts (elements etc), enabling more accurate and robust quantification of the informativeness of an unseen (by the GNN) molecule. 


&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; We herein present the results. The reported values are Mean Absolute Error (MAE, lower is better).

&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; *Abbreviations:* D= Diversity Only, B= Uncertainty + Diversity

| Iteration |  |*mu* | |  |  |*lumo*| | |
|----------|------|-|---|---|----|----|--|--|
|   | SOAP D| Our D |SOAP B| Our B |&#124;   SOAP D| Our D |SOAP B | Our B |
| 1 | 2154|  1769|2057| 1741 |   &#124;   1016 |    890 | 1013 | 876
| 2 | 1901|  1576|1877| 1550 |   &#124; 921 |    805  |902|  799
| 3 | 1732|  1440|1721 |1412 |   &#124; 839 |    740 | 838|  732
| 4 | 1701|  1352|1539| 1315 |   &#124; 797 |    692 |791 |  677
| 5 | 1587|  1225|1456| 1205 |   &#124; 759 |    680 |744|   660
| 6 | 1414|  1157|1345| 1121 |   &#124; 699 |     629 |711|  615
| 7 | 1322|  1092|1280| 1072 |    &#124; 667 |     604 |681|  594

&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; The *p-values* in the table below further demonstrate that our proposed method (using both diversity and uncertainty) significantly improves the performance of the selection strategy compared to the SOAP descriptor (using both SOAP diversity and uncertainty).

||*mu*|*lumo*|
|--| ---| --- |
|*p-value*| 4.20 x 10E-6| 2.51 x 10E-6|

[1] K.S., et al. SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. NIPS'17

[2] J. G., et al. Neural message passing for quantum chemistry. ICML'17

[3] Y. L., et al. Spherical message passing for 3D molecular graphs. ICLR'22

[4]  J. G., et al. GemNet: Universal Directional Graph Neural Networks for Molecules. NeurIPS'21

[5] S.D., et al. Comparing molecules and solids across structural and alchemical space. Physical Chemistry Chemical Physics

[6] M.J., et al.Machine learning hydrogen adsorption on nanoclusters through structural descriptors. npj Computational Materials

[7] A. B., el al. On representing chemical environments. Physical Review B - Condensed Matter and Materials Physics

-----------------
Hope these clarifications can address your concern. If you have further questions, don't hesitate to let us know, and we are always ready to discuss them in detail.
Thank you for the rebuttal. I completely agree that 3D geometry is crucial. But it is not clear to me why 2 molecular graphs with different elements is not considered diverse, but rather uncertain. To a chemist, those are two completely different molecules, even though the graph, with elements removed, is the same. In other words, the formalism for graphs, as far as I can tell, described on line 227, does not include the elements.

After careful thought, I think this work has promise, but will stay with my rating.
Dear Reviewer yNCZ,

Thanks again for your insightful comments, which we believe will improve the clarity of our work. Regarding your concerns, we've made efforts to address them in detail in our rebuttal. As you are the only reviewer who has not responded to our initial rebuttal, we sincerely hope you can check it at your earliest convenience. 

Since the discussion period is approaching its end, we hope you can let us know if we have addressed your critical points and reconsider the score if we have. Meanwhile, we welcome any additional questions and wish to discuss them in detail.

Sincerely,

Authors
Dear Reviewer Abqb,

Thank you for your response and for increasing your score.
We are happy to know that your concerns have been addressed.

Sincerely,

Authors
Thank you for your thoughtful response. My concern has been addressed. I will increase the score from 4 to 5.
Dear Reviewer Abqb,

Thank you for your comments! As **the reviewer-author discussion stage is ending in less than 30 hours**, we kindly remind you that we have provided a detailed rebuttal (both the individual rebuttal and global rebuttal) to address your concerns. We hope our clarifications and additional experiments (in the global rebuttal PDF file) have resolved the concerns to your satisfaction.  

 If you believe your concerns have been adequately addressed, we kindly ask you to reconsider your scores. We are more than willing to discuss them in detail if you have additional questions.

Sincerely,

Authors
Dear Reviewers yNCZ,

Thank you for your insightful comments, which have been very helpful in improving the clarity of our work. As the **reviewer-author discussion stage is ending in 32 hours**, we kindly remind you that we have provided a detailed response to address your concerns. We hope our clarifications have resolved the issues to your satisfaction.

If you believe your concerns have been adequately addressed, we kindly ask you to reconsider your scores. However, if there are any aspects that still need further clarification, please let us know, and we are more than willing to discuss them in detail.

Sincerely,

Authors
Dear Reviewer JSdU,

Thank you so much for your constructive comments, which are very helpful in improving the clarity of our paper.
Also thanks for acknowledging that we've addressed your concerns and raising your score.

As we cannot revise the paper for now, we'll definitely include these new results either in the main paper or in the Appendix in the next version.

Sincerely,

Authors
Dear Reviewer UZqH,

Thanks a lot for acknowledging our clarifications/experiments and raising the score.
Your comments have been constructive and helpful in improving our work. We will integrate our discussions and the new results in the next version of our paper.

Sincerely,

Authors,
The additional results with regards to SOAP descriptors are much appreciated and support the value of the proposed diversity-based descriptor.

Tuning down these bold statements and discussing the contribution properly in the context of the existing literature (as seems to be the case in light of the recent answers) further improves the paper.

In light of this, I've increased my score.
I appreciate the authors' responses. Including all the baselines from the experiments in Section 4.4 would enhance the completeness of the paper, and I strongly encourage adding them to the appendix. Given that my concerns have been satisfactorily addressed, I am willing to increase my score.
Dear Reviewers,

Thank you for your insightful and constructive comments, which have been very helpful in improving our work. As the reviewer-author discussion stage is ending soon, we kindly remind you that we have provided a detailed response to address your concerns (both individual response and global response). We hope our clarifications, along with the additional experimental results and efforts we've included, have resolved the issues to your satisfaction.

If you believe your concerns have been adequately addressed, we kindly ask you to reconsider your scores. However, if there are any aspects that still need further clarification, please let us know, and we are more than willing to discuss them in detail.

Sincerely,

Authors
In Part I of the reply, we present the results of using SOAP descriptors to compute the diversity. We managed to conduct additional experiments using SOAP descriptors to compute the diversity matrix and supplemented it with the uncertainty component for a complete comparison with our method.


We herein present the full results. The reported values are Mean Absolute Error(MAE).

*Abbreviations:* D= Diversity Only, B= Uncertainty + Diversity

| Iteration |  |*mu* | |  |  |*lumo*| | |
|----------|------|-|---|---|----|----|--|--|
|   | SOAP D| Our D |SOAP B| Our B |&#124;   SOAP D| Our D |SOAP B | Our B |
| 1 | 2154|  1769|2057| 1741 |   &#124;   1016 |    890 | 1013 | 876
| 2 | 1901|  1576|1877| 1550 |   &#124; 921 |    805  |902|  799
| 3 | 1732|  1440|1721 |1412 |   &#124; 839 |    740 | 838|  732
| 4 | 1701|  1352|1539| 1315 |   &#124; 797 |    692 |791 |  677
| 5 | 1587|  1225|1456| 1205 |   &#124; 759 |    680 |744|   660
| 6 | 1414|  1157|1345| 1121 |   &#124; 699 |     629 |711|  615
| 7 | 1322|  1092|1280| 1072 |    &#124; 667 |     604 |681|  594

The *p-values* in the table below further demonstrate that our proposed method (using both diversity and uncertainty) significantly improves the performance of the selection strategy compared to the SOAP descriptor (using both SOAP diversity and uncertainty).

||*mu*|*lumo*|
|--| ---| --- |
|*p-value*| 4.20 x 10E-6| 2.51 x 10E-6|

It is evident that our method (using both diversity and uncertainty) outperforms the SOAP descriptors (using both SOAP diversity and uncertainty). As mentioned earlier, global SOAP descriptors may have limited capabilities in capturing essential global geometric information, leading to suboptimal performance. We will include the full results of both BatchBALD and SOAP in our paper.
> $O(n^2)$ might become problematic for larger systems in case the cost starts to dominate the cost of the backbone.

We refer to the case in which we assume an unbounded cutoff distance and a body-order of 2 (which is the most basic case, such as SchNet[2] that only uses relative distances), the complexity of GNN is $O(n^2)$. In general, $d< n$ with a practical cut-off distance and the model complexity is $O(d^{k-1} n)$, where $k$ is the body order. Asymptotically, we agree that our method might dominate the computational complexity of GNNs, but in real implementations, the constant term in the complexity of GNNs might be very large compared to that of our methods because of the large amount of calculations in the network, including linear and nonlinear transformations, feature aggregation and passing, etc.. Therefore, this domination will occur only with a reasonably large number of $n$. In this work, we focus on molecules, so the number of atoms ($n$) is usually small. For other scientific data like proteins, $n$ is large, but this is out of the scope of this work. 
However, we acknowledge that for large-scale datasets that contain several million conformations, the difference in complexity will increase, and we will address this in our discussion of limitations.

But, most importantly, as an active learning approach, our primary focus is on minimizing the costs associated with performing annotation, rather than the computational costs tied to GNNs. These annotation costs can vary, including computational expenses like those incurred from DFT calculations ($O(n^3)$), costs from wet lab experiments, or even the time and expertise required from specialists. In many cases, we prefer to allocate more computational resources to active learning, as this investment can lead to more efficient and cost-effective labeling.

> This should be noted in the limitations given that 2 molecules with similar geometries...

By combining both uncertainty and diversity, we take into consideration both the chemical and geometric contexts. If we only consider the 48-dimensional vector, then, as you mentioned, the scenario you described would be a potential limitation. We acknowledge this concern and will include it in the limitations section.

[1] Hyperactive learning for data-driven interatomic potentials. npj Comput Mater

[2] Quantum-chemical insights from deep tensor neural networks. Nature Communications
Thank you for your prompt reply. Your insights are invaluable to us.

> These works do train 3D GNNs via active learning...

First, we apologize for any confusion caused by our language. When we said that they do not specifically consider 3D geometry information, we meant that the active learning primarily pertains to 3D geometries, as in our diversity metric.

The scopes of the mentioned two methods [1,2] are slightly different, but for the active selection part, both methods sample configurations in sub-regions where the machine learning models have maximal **uncertainty** about for quantum mechanics (e.g., DFT) calculations, to be added to training data. In performing similar tasks (e.g., in our MD17 Benzene experiments), their methods share some similarities with our uncertainty component. However, their methods specifically targets non-equilibrium state modeling, while our method is more general, applicable to molecular systems in both equilibrium (e.g., QM9 experiments) and non-equilibrium states, covering various quantum properties and molecular dynamics simulation tasks.

Moreover, our approach incorporates both uncertainty and diversity. While 3D information is considered in active learning approaches such as [1, 2], it is not as central as in our diversity computations. Usually, the models (3D GNNs) play an important role in uncertainty quantification; since many node features, other than geometries, are fed as inputs to the model, uncertainty quantification may or may not prioritize 3D geometric information. Our diversity component is laser-focused on 3D geometric information. While pure uncertainty-based methods focus solely on the most uncertain samples, our method balances exploration and exploitation, leading to more comprehensive and reliable sample selection. Experiments in Sec. 4.4 in our paper and Fig. 2 in the global rebuttal PDF show that, our method consistently outperforms the individual use of uncertainty for all active learning iterations.

We understand your point of view that these methods also consider 3D information, although not as specific as our method (especially the diversity component). We will make sure to include these references, discuss our focus in further detail, and modify our sentences to avoid bold claims, as you mentioned in your first reply.

> It is unfortunate that there are no results on this given that the proposed 48-dim vector seems to be the biggest novelty of the paper, so a comparison to the clustering capabilities to standard approaches such as SOAP (acknowledging the fact that another step is required to obtain a global descriptor from the local descriptors) would be highly valuable.

Unfortunately, we did not have enough time to complete the experiments during the rebuttal period. However, during this discussion period, we managed to perform the experiments using the SOAP descriptor and compare its clustering capabilities to our 48-dimensional vector. We first obtain the local SOAP descriptors and then aggregate them to obtain a global descriptor for a molecule. Then, we apply the same setting as mentioned in section 2.3 of the paper to obtain the results.

We herein present the results. The reported values are Mean Absolute Error (MAE) values.

*Abbreviations:*
D= Diversity Only, B= Uncertainty + Diversity

| Iteration |  |*mu* |  |  |*lumo*| |
|----------|------|-|---|----|----|--|
|   | SOAP D| Our D | Our B |&#124;   SOAP D| Our D | Our B |
| 1 | 2154|  1769| 1741 |   &#124;   1016 |    890  | 876
| 2 | 1901|  1576| 1550 |   &#124; 921 |    805  | 799
| 3 | 1732|  1440| 1412 |   &#124; 839 |    740 |  732
| 4 | 1701|  1352| 1315 |   &#124; 797 |    692 |  677
| 5 | 1587|  1225| 1205 |   &#124; 759 |    680 |  660
| 6 | 1414|  1157| 1121 |   &#124; 699 |     629 | 615
| 7 | 1322|  1092| 1072 |    &#124; 667 |     604 | 594

*p-values* in the table below show that our proposed 48-dimensional vector significantly improves the performance of the selection strategy over the SOAP descriptor.

||*mu*|*lumo*|
|--| ---| --- |
|*p-value*| 3.24 x 10E-6| 2.29 x 10E-5|

It can be clearly observed that our method (whether using diversity alone or both diversity and uncertainty) outperforms SOAP descriptors. This outperformance may be attributed to the local nature of SOAP descriptors; aggregation to obtain a global descriptor can be limiting, as it may not fully capture the intricate geometric variations. Meanwhile, another important aspect to consider is permutation invariance. Through aggregation, we can ensure such symmetry. However, it is challenging and potentially a future work to determine how to effectively design a global descriptor based on SOAP that fulfills permutation invariance. We will include the results and discussion of both BatchBALD and SOAP in our paper.
> The first two active learning (AL) papers do not specifically consider 3D geometry information

These works do train 3D GNNs via active learning. They do not perform any diversity-based selection, but the models and datasets are inherently 3D. When the authors claim that their model uses the atomic species via the uncertainty-metric, so do these works consider the 3D structure via the uncertainty metric.

> SOAP descriptors

It is unfortunate that there are no results on this given that the proposed 48-dim vector seems to be the biggest novelty of the paper, so a comparison to the clustering capabilities to standard approaches such as SOAP (acknowledging the fact that another step is required to obtain a global descriptor from the local descriptors) would be highly valuable.

> existing mainstream 3D GNNs have the complexity of O(n^2)

(Assuming n is the number of atoms in the system) Then these 3D GNN backbones have have complexities of O(d^p n), i.e. are linear in the number of atoms, times the average number of neighbors d (which is a constant that is depending on the cut_off and the density of the system, but independent of n). If the proposed method is indeed O(n^2), this might become problematic for larger systems in case the cost starts to dominate the cost of the backbone.

> Q1: 48-dim vector: It is purely geometry-based

This should be noted in the limitations given that 2 molecules with similar geometries, but different atom species tend to behave very differently, even though the 48-dim vector treats them very similarly.
References

[1] J.K., et al. Directional message passing for molecular graphs. ICLR'20

[2] Y. L., et al. Spherical message passing for 3D molecular graphs. ICLR'22

[3] J. G., et al. GemNet: Universal Directional Graph Neural Networks for Molecules. NeurIPS'21

[4] M.G., et al. GPU acceleration of admm for large-scale quadratic programming. Journal of Parallel and Distributed Computing

[5] S. C., et al. Active batch selection via convex relaxations with guaranteed solution bounds. TPAMI

[6] P. J B., et al. Ultrafast shape recognition to search compound databases for similar molecular shapes. Journal of computational chemistry

[7] Hall P. A distribution is completely determined by its translated moments. Zeitschrift für Wahrscheinlichkeitstheorie und Verwandte Gebiete.