Year: 2021

Abstract:
We consider the problem of generating 3D molecular geometries from scratch. While multiple methods have been developed for generating molecular graphs, generating 3D molecular geometries from scratch is largely under-explored. In this work, we propose G-SphereNet, a novel autoregressive flow model for generating 3D molecular geometries. G-SphereNet employs a flexible sequential generation scheme by placing atoms in 3D space step-by-step. Instead of generating 3D coordinates directly, we propose to determine 3D positions of atoms by generating distances, angles and torsion angles, thereby ensuring both invariance and equivariance properties. In addition, we propose to use spherical message passing and attention mechanism for conditional information extraction. Experimental results show that G-SphereNet outperforms previous methods on random molecular geometry generation and targeted molecule discovery tasks.

text: We consider the problem of generating 3D molecular geometries from scratch. While multiple methods have been developed for generating molecular graphs, generating 3D molecular geometries from scratch is largely under-explored. In this work, we propose G-SphereNet, a novel autoregressive flow model for generating 3D molecular geometries. G-SphereNet employs a flexible sequential generation scheme by placing atoms in 3D space step-by-step. Instead of generating 3D coordinates directly, we propose to determine 3D positions of atoms by generating distances, angles and torsion angles, thereby ensuring both invariance and equivariance properties. In addition, we propose to use spherical message passing and attention mechanism for conditional information extraction. Experimental results show that G-SphereNet outperforms previous methods on random molecular geometry generation and targeted molecule discovery tasks.

heading: 1 INTRODUCTION
text: Designing and synthesizing novel molecules with desirable properties is a challenging task in drug discovery (Wang et al., 2020; Stokes et al., 2020) and chemical science. The size of search space of all chemical molecules is estimated to be on the order of 1033 (Polishchuk et al., 2013), thereby making exhaustive search infeasible. In recent years, advances in machine learning methods have greatly accelerated the progress in this filed. Many studies represent molecules as 2D molecular graphs and propose to automatically generate molecular graphs and optimize molecular properties with deep generative models, such as variational auto-encoders (Kingma & Welling, 2014).
However, complete information of molecules cannot be obtained from 2D molecular graphs because 3D structures, also known as 3D molecular geometries, are critical in determining many molecular properties. 3D molecular geometries represent the 3D coordinates of atoms and are important for the accurate prediction of quantum properties (Schütt et al., 2017). Hence, we argue that generating 2D molecular graphs might not be the best way to identify novel molecules with certain desirable quantum properties. Instead, developing a generative model that can generate 3D molecular geometries from scratch is a promising solution to this problem. Currently, this area remains under-explored. Recently, a series of seminal studies (Xu et al., 2021b;a; Shi et al., 2021) have proposed to generate 3D molecular geometries from given 2D molecular graphs. These methods themselves do not generate novel molecules once the 2D molecular graphs are given.
In this work, we propose G-SphereNet, a Generative model for 3D molecular geometry generation from scratch inspired by SphereNet (Liu et al., 2021d). In G-SphereNet, 3D molecular geometries are generated by sequentially placing atoms in 3D space. The 3D positions of atoms are implicitly determined by generating distances, angles and torsion angles to ensure both invariance and equivariance properties. Our work is inspired by SphereNet, which use distances, angles and torsion angles to compute predictive representations of molecules. In addition, G-SphereNet employs SphereNet and attention mechanism to extract conditional information. Experimental results show that G-SphereNet can outperform prior methods on the 3D molecular geometry generation task.

heading: 2 BACKGROUND AND RELATED WORK

heading: 2.1 2D MOLECULAR GRAPH GENERATION
text: Recently, thanks to the advances of deep generative models, significant progress has been made in the problem of molecule design and generation. Some methods (Gómez-Bombarelli et al., 2018; Kusner et al., 2017; Dai et al., 2018; Liu et al., 2020) use sequence models to generate SMILES (Weininger, 1988) string representations of molecules. Other studies consider molecules as graphs, in which atoms and chemical bonds of the molecule are represented by nodes and edges, respectively. These studies either generate the node type and adjacency matrix of the graph (Simonovsky & Komodakis, 2018; Ma et al., 2018; De Cao & Kipf, 2018; Liu et al., 2021c), or form the molecular graph by sequentially adding nodes and edges (You et al., 2018; Shi* et al., 2020), or compose the molecule from a junction tree of molecular motifs (Jin et al., 2018).
However, these methods only generate the graph structures of molecules while crucial 3D molecular geometries are ignored. In other words, the 3D coordinates of atoms in the molecule are unknown. Hence, these generative models cannot distinguish spatial isomers, i.e., molecules with the same molecular graph but different 3D molecular geometries. In addition, 3D molecular geometries are required for the computation of some quantum properties of molecules, such as HOMO-LUMO gap (Hu et al., 2021; Liu et al., 2021a; Ying et al., 2021; Addanki et al., 2021; Xu et al., 2021c). Hence, these generation methods cannot be used when spatial isomers or quantum properties are needed.

heading: 2.2 3D MOLECULAR GEOMETRY GENERATION FROM 2D INFORMATION
text: Currently, some studies have proposed to generate 3D molecular geometries from the 2D conditional information of target molecules. Some methods (Simm et al., 2020; 2021) propose to generate 3D molecular geometries through minimizing the energy of atomic systems, in which the numbers and types of atoms are explicitly given in atom bags. Other methods (Mansimov et al., 2019; Simm & Hernandez-Lobato, 2020; Gogineni et al., 2020; Xu et al., 2021b;a; Shi et al., 2021; Ganea et al., 2021) propose to randomly sample multiple 3D molecular geometries from the molecular graph of the target molecule with deep generative models. These conditional generation methods assume that we are given the target molecules in the form of molecular graphs, but we do not know their 3D geometries. However, this assumption does not hold in the targeted molecule discovery problem. In this problem, we aim to discover novel molecules with good quantum properties, such as low HOMO-LUMO gaps. In other words, target molecules themselves are unknown and to be generated, and their 3D geometries are also needed because quantum properties cannot be accurately estimated only from molecular graphs. Hence, we argue that developing a method to generate 3D molecular geometries from scratch is more useful to this problem.

heading: 2.3 3D MOLECULAR GEOMETRY GENERATION FROM SCRATCH
text: In this work, we consider the problem of generating 3D molecular geometries from scratch. Let G = {Gj}mj=1 be a set of 3D molecular geometries, and the function S(G) ∈ R computes a specific quantum property score of G. We consider the two generation tasks defined as below.
• Learning a random generation model pθ(·) from G, so that the model can sample a valid 3D molecular geometry G with a high probability pθ(G).
• Learning a targeted molecule discovery model pθ(·) so as to maximize (or minimize) the expected quantum property score EG∼pθ [S(G)].
This problem is largely under explored and there are only a few studies attempting to solve it. GSchNet (Gebauer et al., 2019) uses an autoregressive model based on SchNet (Schütt et al., 2017) to sequentially generate the new atom and place it at the local grid point of a focal atom. On the other hand, EDMNet (Hoffmann & Noé, 2019) and 3DMolNet (Nesterov et al., 2020) generate pairwise distances between atoms with generative adversarial networks (GAN) (Goodfellow et al., 2014) and variational auto-encoders (VAE) (Kingma & Welling, 2014), respectively. Besides, E-NFs (Satorras et al., 2021a) proposes a geometry generation model by combining flow models (Rezende & Mohamed, 2015) with E(n) equivariant graph neural networks (Satorras et al., 2021b). It generates 3D coordinates of all atoms in a one-shot fashion and defines the prior distribution in a subspace of
the latent space to ensure translation invariance. Different from these methods, our method employs a flexible sequential generation pipeline based on autoregressive flow models, which can capture the density of 3D molecular geometries more effectively.

heading: 2.4 FLOW MODELS
text: A flow model defines a parametric invertible mapping fθ : z ∈ Rd → x ∈ Rd, where the data point x and the latent variable z are both random variables. Given that z is sampled from a known prior distribution pZ and fθ is invertible, we can compute the log-likelihood of x as
log pX(x) = log pZ ( f−1θ (x) ) + log |detJ | , (1)
where J = ∂f −1 θ (x)
∂x is the Jacobian matrix. To train the flow model on a given dataset X = {xi} m i=1,
the log-likelihoods of data points are computed from Eqn. (1) and maximized by gradient descent. Hence, tractable and cheap computation of detJ is needed for efficient training. A common choice of fθ in most flow models is the affine coupling mapping (Dinh et al., 2014; 2016), in which case computing detJ is very easy because J is an upper triangular matrix.
Flow models have been used in a variety of generation tasks (Rezende & Mohamed, 2015; Tran et al., 2019; Köhler et al., 2020). Compared with VAE and GAN, they allow for exact likelihood computation, and can model the density of data more accurately. Because of these advantages, many recent studies have used flow models in the molecule generation task. Some one-shot methods, including GraphNVP (Madhawa et al., 2019), GRF (Honda et al., 2019), and MoFlow (Zang & Wang, 2020), consider the node type and adjacency matrix as the generation targets. On the other hand, GraphAF (Shi* et al., 2020) and GraphDF (Luo et al., 2021) generate molecular graphs by generating nodes and edges sequentially through autoregressive flow models (Papamakarios et al., 2017). These models have stronger capacity to model graph structures than one-shot methods, and achieve state-of-the-art performance in the molecular graph generation task.

heading: 3 METHODS
text: While autoregressive flow models have been successfully applied to the molecular graph generation task, it remains unclear whether they are sufficiently powerful to model more complicated 3D molecular geometries. In this section, we present G-SphereNet, a novel 3D molecular geometry generation method. It adopts a flexible, effective, and efficient sequential generation pipeline based on autoregressive flow models, which can ensure the equivariance property of coordinates and the invariance property of likelihood simultaneously. In addition, expressive spherical message passing based graph neural networks (Liu et al., 2021d) and multi-head attention networks (Vaswani et al., 2017) are used in G-SphereNet to extract 3D conditional information for accurate generation. To the best of our knowledge, G-SphereNet is the first likelihood-based autoregressive generative model for 3D molecular geometry generation.

heading: 3.1 SEQUENTIAL GENERATION
text: Let k be the number of atom types. We use a 3D point cloud G = (A,R) to represent the 3D geometry of a molecule with n atoms, where A ∈ {0, 1}n×k is the atom type matrix and R ∈ Rn×3 is the atom coordinate matrix. Each row in the matrix A is a one-hot vector, and A[j, u] = 1 represents that the j-th atom has type u. The 3-dimensional row vector at the j-th row of the matrix R represents the 3D Cartesian coordinate of the j-th atom.
We consider the generation of 3D molecular geometries as a sequential decision process. We start from a molecular geometry G1 with one carbon atom at the origin point, and generate the complete geometry by adding a new atom at each step. Specifically, at the i-th step, let the intermediate 3D molecular geometry generated from the previous i − 1 steps be Gi = (Ai, Ri), which has i atoms. The atom type ai ∈ {0, 1}k of the new atom is generated by the generative model ga based on the latent variable zai . Afterwards, the generative model g
r decides the 3D Cartesian coordinate ri ∈ R3 of the new atom based on the latent variable zri . Both ga and gr are autoregressive functions of intermediately generated geometries. The overall sequential generation process can be described by the following equations:
ai = g a (zai ;Ai, Ri) , ri = g r (zri ;Ai, Ri) , i ≥ 1. (2)

heading: 3.2 GENERATION WITH AUTOREGRESSIVE FLOW MODELS
text: We employ autoregressive flow models to generate the atom type ai of the new atom at each step. Since atom types are discrete numbers, we adopt the dequantization method (You et al., 2018; Shi* et al., 2020) to convert them into continuous numbers by adding real-valued noise as
ãi = ai + u, u ∼ U(0, 1)k, i ≥ 1, (3)
where U(0, 1) is the uniform distribution over the interval (0, 1). To generate ai, we first sample the latent variable zai ∈ Rk from the standard Gaussian distribution N (0, 1), and then map zai to ãi by the affine transformation as
ãi = s a i zai + tai , i ≥ 1, (4)
where denotes the element-wise multiplication, the scale factor sai and the shift factor tai both depend on the conditional information extracted from the intermediate geometry Gi = (Ai, Ri). Intuitively, ãi should be invariant to any rigid transformation on Ri, i.e., ãi should not change if we rotate or translate Ri in 3D space. Hence, we use a symmetry-invariant model to compute sai and t a i from Gi, which are described in Sec. 3.4 in detail. After obtaining ãi, ai can be computed by taking the argmax of ãi, as ai = one-hot(argmax ãi).
However, we cannot generate the 3D coordinate ri in the same way as the generation of the atom type ai. Directly calculating ri with the autoregressive flow model, as in Eqn. (4), neither satisfies the equivariance property of coordinates nor the invariance property of likehood. First, it is easy to find that if we rotate or translate Ri, then ri needs to be rotated or translated correspondingly. Formally, it means that for any orthogonal matrix Q ∈ R3×3 and translation vector b ∈ R3, if ri = g r(zri ;Ai, Ri), we have
Qri + b = g r ( zri ;Ai, RiQ T + 1bT ) , (5)
where 1 denotes a vector of all ones of length i. If we directly compute ri by the autoregressive model similar to Eqn. (4), i.e., ri = sri zri + tri , then to satisfy Eqn. (5), the correctness of sQr+bi zri + t Qr+b i = Q[s r i zri ] + Qtri + b has to be ensured for any orthogonal matrix Q and translation vector b. However, it is very hard to design a flow model satisfying this complicated condition. Second, the likelihood p(ri|Ai, Ri) should be invariant to rotations and translations as they do not change the 3D structure. In other words,
p ( Qri + b|Ai, RiQT + 1bT ) = p (ri|Ai, Ri) (6)
should be satisfied for any orthogonal matrix Q and translation vector b. It follows from the changeof-variable theorem that p (ri|Ai, Ri) = p (zri ) det∂zri∂ri . Let the latent variable corresponding to Qri + b be z Qr+b i . The affine transformation in Eqn. (4) does not admit any relationship between either p(zri ) and p(z Qr+b i ), or ∂zri ∂ri and ∂z Qr+b i ∂(Qri+b) , so Eqn. (6) is not guaranteed to hold. Hence, we cannot guarantee the invariance property.
Given these limitations, we instead propose to determine the 3D relative position of the new atom by symmetry-invariant elements. Similar to G-SchNet (Gebauer et al., 2019), we first choose a focal atom among all i atoms in Gi, which serves as the reference point for the location of the new atom. Then, the model generates the distance di, the angle θi, and the torsion angle ϕi successively.
Specifically, assuming that the focal atom is the f -th atom of Gi, the distance di = ||ri− rf ||2 from the focal atom to the new atom is first generated. Then, if i ≥ 2, the angle θi ∈ [0, π] between the lines (rf , ri) and (rf , rc) is generated, where c is the atom closest to f in Gi. Finally, if i ≥ 3, the torsion angle ϕi ∈ [−π, π] between the planes formed by positions (rf , rc, ri) and (rf , rc, re) is generated, where e is the atom closest to c but different from f in Gi. Similar to ãi, di, θi, and ϕi are generated as
di = s d i zdi + tdi , i ≥ 1,
θi = s θ i zθi + tθi , i ≥ 2, ϕi = s ϕ i z ϕ i + t ϕ i , i ≥ 3,
(7)
where zdi , z θ i , z ϕ i ∈ R are all latent variables sampled from standard Gaussian distributions, and the scale factors sdi , s θ i , s ϕ i ∈ R and the shift factors tdi , tθi , t ϕ i ∈ R are functions of Gi. Afterwards, the coordinate ri is computed from the relative positional elements di, θi, ϕi and the coordinates rf , rc, re. We show that such a process of placing the new atom in 3D space can strictly satisfy the conditions in Eqn. (5) and (6) in Appendix A, thereby satisfying both invariance and equivariance properties. The sequential generation process repeats until either the maximum number of atoms is reached, or no atom can be chosen as the focal atom by the atom-wise classifier. Our proposed sequential generation method is related to SphereNet (Liu et al., 2021d) in that our generation targets are the 3D information used by SphereNet to extract features, so we name our method as G-sphereNet. An illustration of the overall generation process in G-SphereNet is given in Figure 1.

heading: 3.3 DISCUSSIONS
text: We argue that our proposed G-SphereNet method has many advantages over previous 3D molecular geometry generation methods. First, it is easier for G-SphereNet to generate valid geometries theoretically because the exact 3D coordinate of each atom can always be obtained. However, EDMNet (Hoffmann & Noé, 2019) and 3DMolNet (Nesterov et al., 2020) generate the pairwise distances of atoms in the form of distance matrices. There is no theoretical guarantee that the generated matrices are always valid Euclidean distance matrices, or correspond to coordinates in 3D space (Dokmanic et al., 2015). Second, the generation of 3D positions in G-SphereNet is more flexible than that in G-SchNet (Gebauer et al., 2019). In G-SchNet, the new atom has to be placed at one of the candidate grid points circling around the focal atom, but it can be placed at any relative position of the focal atom in G-SphereNet. Third, compared with E-NFs (Satorras et al., 2021a), G-SphereNet is more efficient and effective. E-NFs maps latent variables to 3D coordinates of atoms by a flow model. Since coordinates are not translation invariant, E-NFs proposes to obtain latent variables by computationally expensive operations. Specifically, E-NFs first samples from the prior distribution defined in a subspace of the latent space, then maps the sampled variable to the latent variable by a linear projection. In contrast, G-SphereNet obtain 3D positions by generating distances, angles and torsion angles, which are naturally translation invariant. Hence, G-SphereNet can avoid the complicated operations of E-NFs. In addition, E-NFs generates the coordinates of all atoms in the geometry at a time, while G-SphereNet obtains the coordinate of one atom at a time. Though G-SphereNet may be slower, we argue that the sequential generation fashion helps the model to capture the dependency between atoms and the density of geometries more effectively. Experimental results also demonstrate that G-SphereNet can generate much more valid molecular geometries than E-NFs.

heading: 3.4 CONDITIONAL INFORMATION EXTRACTION
text: As we have mentioned in Sec. 3.2, generating the atom type and 3D position of the new atom requires capturing conditional information from the intermediate geometry at each step. Desirable conditional information should incorporate comprehensive 3D structural features of the geometry, and be invariant to any rigid transformation. To achieve this goal, we propose to capture conditional information using SphereNet (Liu et al., 2021d), an advanced 3D graph neural network model. SphereNet considers the input molecular geometry as a directional cutoff graph. Denoting the feature of k-th edge by ek and the feature of i-th node by vi, SphereNet initializes them with spherical basis functions and updates them to e′k and v ′ i with spherical message passing as
e′k = φ e (ek,vrk ,vsk ,Esk , ρ p→e (Rsk)) , v′i = φ v (vi, ρ e→v(Ei)) . (8)
Here, φe, φv are updating functions, ρp→e, ρe→v are aggregation functions, rk and sk are the sending and receiving nodes of the k-th edge, Rsk are the coordinates of neighboring nodes of sk, Esk and Ei are the features of edges incident to node sk and i, separately. SphereNet has powerful 3D structural feature extraction ability, and achieves good performance in multiple quantum property prediction tasks. Considering these advantages, we use the SphereNet model as the backbone feature extractor to capture conditional information from the intermediate molecular geometry.
For the input molecular geometry Gi, let the node embeddings computed from the SphereNet be {hi,j}i−1j=0. To select the focal atom, we use an atom-wise multi-layer perceptron (MLP) taking the corresponding node embedding as input, and randomly choose the focal atom f from atoms whose classification scores are higher than 0.5. If all the classification scores outputted from the atom-wise classifier are lower than 0.5, then no atom can be chosen as the focal atom and the sequential generation process terminates. Afterwards, the scale and shift factors in Eqn. (4) and (7) are computed. We ever tried computing them using solely the node embeddings of the reference points f , c, and e. However, our experiments show that it frequently causes incorrect placements of new atoms in 3D space. We think it is because node embeddings only contain local 3D information, which is insufficient for the accurate generation of the 3D positions of new atoms.
To tackle the above issue, we propose to augment node embeddings with global features extracted by a multi-head attention network (Vaswani et al., 2017). Formally, a multi-head attention network takes the query matrixQ, key matrixK, and value matrix V as inputs, and extract global information from inputs by the multi-head attention mechanism:
Qi = QW Q i , Ki = KW K i , Vi = VW V i , ATTi = softmax
( QiK
T i√ p
) Vi, 1 ≤ i ≤ o,
MH-ATT(Q,K, V ) = Con (ATT1, ...,ATTo)WO, (9)
where Con(·) denotes the concatenation operation, p is the size of the second dimension of K, o is the number of attention heads, the matrices WQi , W K i , W V i , and W
O are all trainable parameters. Let the node embedding matrix of Gi be Hi = [hi,0, ..., hi,i−1] T , sai and t a i are computed as
Con(sai , t a i ) = MLP
a ( Con ( hi,f ,MH-ATTa ( hTi,f , Hi, Hi ))) , i ≥ 1, (10)
where MH-ATTa is a multi-head attention network and MLPa is a multi-layer perceptron. As for the scale and shift factors in Eqn. (7), we first multiply node embeddings with the atom type embedding vector to include the atom type information:
hai = Embedding(ai), h̃i,j = h a i hi,j , H̃i = [ h̃i,0, ..., h̃i,i−1 ]T , i ≥ 1, 0 ≤ j ≤ i− 1,
then compute them as
Con(sdi , t d i ) = MLP
d ( Con ( h̃i,f ,MH-ATTd ( h̃Ti,f , H̃i, H̃i ))) , i ≥ 1,
Con(sθi , t θ i ) = MLP
θ ( Con ( h̃i,f , h̃i,c,MH-ATTθ ( h̃Ti,f , H̃i, H̃i ))) , i ≥ 2,
Con(sϕi , t ϕ i ) = MLP
ϕ ( Con ( h̃i,f , h̃i,c, h̃i,e,MH-ATTϕ ( h̃Ti,f , H̃i, H̃i ))) , i ≥ 3.
(11)
Here Embedding(·) is a lookup based embedding layer, MH-ATTd,MH-ATTθ,MH-ATTϕ are all multi-head attention networks, and MLPd,MLPθ,MLPϕ are all multi-layer perceptrons. The use of the multi-head attention networks helps extract more comprehensive 3D conditional information and is demonstrated to improve the generation performance a lot in our ablation study.

heading: 3.5 TRAINING
text: To train the G-SphereNet model on a dataset, we first need to split each individual molecular geometry in the dataset into a trajectory of atom addition steps. In other words, the generation order of atoms in the geometry and all corresponding focal atoms need to be determined. Since the generated atom is supposed to be placed in the local region of the focal atom during generation, we propose to obtain the training trajectory by applying Prim’s algorithm on the geometry. This procedure can
ensure that the sampled focal atom is always the nearest neighbor of the new atom among all atoms in the intermediate geometry.
For a 3D molecular geometry G having n atoms (n > 3), we maximize its log-likelihood to train the G-SphereNet model. Specifically, we obtain the generation targets, i.e., the atom type and 3D position of the atom to be generated at each step, then compute the log-likelihood of G as
log p(G) = n−1∑ i=1 [ log pZa(z a i ) + log pZd(z d i ) ] + n−1∑ i=2 log pZθ (z θ i ) + n−1∑ i=3 log pZϕ(z ϕ i ), (12)
where latent variables zai , z d i , z θ i , z ϕ i are computed by reversing the mapping in Eqn. (4) and (7), and pZa , pZd , pZθ , pZϕ are all standard Gaussian distributions. Besides, the atom-wise classifier in the G-SphereNet model, which is used for the focal atom selection, is trained by the binary cross entropy loss. The ground-truth label is 1 if the atom is not valence full filled, otherwise is 0. We describe the detailed training and generation algorithm of G-SphereNet in Appendix B.

heading: 4 EXPERIMENTS
text: In this section, we evaluate the proposed G-SphereNet method on the random molecular geometry generation task and the targeted molecule discovery task described in Sec. 2.3. We show that in these tasks, G-SphereNet can outperform previous 3D molecular geometry methods, including GSchNet (Gebauer et al., 2019) and E-NFs (Satorras et al., 2021a). Note that we do not compare with recent methods (Xu et al., 2021b;a; Shi et al., 2021) of generating 3D molecular geometries from 2D information because they cannot do targeted molecule discovery (Sec. 2.2). In addition, we conduct extensive ablation studies to evaluate the advantages of some designs in G-SphereNet method.

heading: 4.1 RANDOM MOLECULAR GEOMETRY GENERATION
text: Data. For the random molecular geometry generation task, we evaluate G-SphereNet on the QM9 (Ramakrishnan et al., 2014) dataset. The QM9 dataset provides over 130k molecules and their corresponding 3D molecular geometries computed by density functional theory (DFT). We randomly select 100k 3D molecular geometries as the training data and 10k 3D molecular geometries as the validation data. For fair comparison, the models of our G-SphereNet method and all other methods are trained with the same data split.
Setup. We use the chemical validity percentage (Validity in short) to evaluate the generation accuracy of G-SphereNet. Specifically, all generated 3D molecular geometries are converted to molecular graphs by the method proposed in Kim & Kim (2015), and the Validity is defined as the percentage of molecular graphs which do not violate chemical valency rules. In addition to the chemical validity, we also evaluate the 3D structural accuracy of the generated molecular geometries. We ever tried to follow G-SchNet (Gebauer et al., 2019) to compute the aligned coordinate differences between the generated geometries and their relaxed ones, but we find the relaxing process involves the expensive computation based on DFT, which takes hours for a single molecular geometry. Hence, we propose to evaluate by the Maximum Mean Discrepancy (MMD) (Gretton et al., 2012) distances of bond length distributions. Formally, for a certain type of bond, we obtain its length distribution in the generated geometries and in the geometries of the dataset, separately, and compute the statistical discrepancy between the two length distributions with the MMD distance. We compute the MMD
distances on six most frequently appeared types of chemical bonds, including carbon-carbon single bonds (C-C), carbon-nitrogen single bonds (C-N), carbon-oxygen single bonds (C-O), hydrogencarbon single bonds (H-C), hydrogen-nitrogen single bonds (H-N), and hydrogen-oxygen single bonds (H-O). All metrics are computed from 10,000 generated molecular geometries.
The implementation of the SphereNet (Liu et al., 2021d) model used for condition information extraction is based on the code of DIG (Liu et al., 2021b) package. We use Adam (Kingma & Ba, 2015) optimizer to train the G-SphereNet model for 100 epochs, with a batch size of 64 and a learning rate of 0.0001. See Appendix C for model configuration and other training details. GSphereNet is compared with G-SchNet (Gebauer et al., 2019) and E-NFs (Satorras et al., 2021a) in terms of Validity and MMD distances, and we run the code provided by authors to obtain the results of two baseline methods. We do not compare with EDMNet (Hoffmann & Noé, 2019) or 3DMolNet (Nesterov et al., 2020) because the authors of 3DMolNet do not provide their implementation, and the EDMNet model cannot be trained on molecular geometries with variable numbers of atoms, which inhibits a fair comparison with G-SphereNet.
Results. We present the performance of different methods in Table 1. Our G-SphereNet method achieves the highest Validity of 88.18%, while E-NFs achieves a much lower Validity of 39.77%. The good performance strongly demonstrates that the sequential fashion of G-SphereNet helps the model to capture the dependency between atoms and learn the underlying chemical rules of molecular geometries more effectively. In addition, compared with G-SchNet, our G-SphereNet achieves lower MMD distances for 4 types of chemical bonds, which shows that our method can model the 3D structural distribution of molecular geometries more accurately. We visualize some molecular geometries generated by G-SphereNet in Figure 2 of Appendix D.

heading: 4.2 TARGETED MOLECULE DISCOVERY
text: Setup. In the targeted molecule discovery task, we aim to maximize or minimize the expected quantum property score. We conduct two targeted molecule discovery experiments, namely minimizing the HOMO-LUMO gap and maximize the isotropic polarizability. Following G-SchNet (Gebauer et al., 2019), we fine-tune the G-SphereNet model that has been trained in Sec. 4.1 on the biased datasets. Specifically, from the QM9 dataset, we collect all molecular geometries whose HOMOLUMO gaps are lower than 4.5 eV and all molecular geometries whose isotropic polarizabilities are higher than 91 Bohr3. The G-SphereNet model is then fine-tuned on these two biased datasets so as to generate molecular geometries with low HOMO-LUMO gaps or high isotropic polarizabilities, respectively. Details about the model fine-tuning process are summarized in Appendix C.
In this task, we evaluate the performance by three statistic metrics over the quantum property scores of generated molecular geometries. Specifically, we generate 1000 molecular geometries with the trained model, and filter out the geometries that are not chemically valid. Afterwards, the PySCF (Sun et al., 2018; 2020) package is used to compute the quantum property scores of valid molecular geometries. The performance is then evaluated by three statistic metrics over these quantum property scores. Formally, we calculate the mean and the optimal value over all property scores, and the percentage of property scores falling into the good region (good percentage in short). The good region is defined as scores lower than 4.5 eV for the HOMO-LUMO gap and scores higher than 91 Bohr3 for the isotropic polarizability, respectively. We find that E-NFs (Satorras et al., 2021a) fails to generate enough chemically valid molecular geometries and produce reliable results after fine-tuning, so we only compare our G-SphereNet method with G-SchNet in this task.
Results. Results of targeted molecule discovery for two quantum properties are summarized in Table 2. For both properties, our G-SphereNet can outperform G-SchNet in all metrics, showing that G-SphereNet can generate more molecular geometries with good properties. Since the two methods use the same pretraining and fine-tuning pipeline, we argue that the better performance of G-SphereNet indicates its stronger ability to search molecular geometries with desirable properties. We illustrate some generated molecular geometries with good properties in Figure 3 of Appendix D.

heading: 4.3 ABLATION STUDIES
text: In previous sections, we have demonstrated the effectiveness of our G-SphereNet method on two 3D molecular geometry generation tasks. However, it is unclear whether some designs in our method, such as the use of global features extracted by multi-head attention networks, can indeed lead to good performance. Hence, we conduct extensive ablation studies justifying the use of both local and global features, and the consideration of distance, angle, and torsion information in the 3D conditional information extraction of G-SphereNet. In addition, we study the effects of different focal atom selection methods. In each ablation study, different variants of G-SphereNet models are trained with the setting in Sec. 4.1 and evaluated by the Validity metric. Table 3 show all the results of ablation studies.
Ablation on local and global feature. We compare with G-SphereNet variants which only use local features, i.e., node embeddings extracted by SphereNet, and only use global features, i.e., outputs from multi-head attention networks) in Eqn. (10) and (11). Results in Table 3(a) show that using only local or global features both achieves a worse performance.
Ablation on 3D information. To show the advantages of using comprehensive 3D information, we replace the SphereNet by SchNet (Schütt et al., 2017) considering only distance information, and DimeNet++ (Klicpera et al., 2020) considering only distances and angles, respectively. As presented in Table 3(b), missing partial 3D information leads to performance degradation.
Ablation on focal atom selection. In G-SphereNet, an atom-wise MLP and sigmoid function performs the atom-wise binary classification finding atoms that are not valence full filled. Then the focal atom is randomly selected from those atoms. In contrast, Simm et al. (2020) proposes to directly select the focal atom by an MLP and softmax function. We discuss the differences between two methods in detail in Appendix C. To demonstrate the benefits of our method, we compare it with a G-SphereNet variant replacing our focal atom selection method by the one in Simm et al. (2020). We denote our method by Sigmoid and the method in Simm et al. (2020) by Softmax. As shown in Table 3(c), our Sigmoid method can achieve much better performance than the Softmax method.

heading: 5 CONCLUSION
text: We propose G-SphereNet, a novel autoregressive flow model for 3D molecular geometry generation from scratch. G-SphereNet employs a sequential generation pipeline, in which the 3D positions of atoms are obtained through generating the relative distances, angles, and torsion angles. It is flexible and efficient, and can ensure both equivariance and invariance properties. In addition, spherical message passing and attention mechanism are used to extract conditional information during sequential generation. We empirically demonstrate that, compared with previous methods, our G-SphereNet method models the distribution of 3D molecular geometries more accurately, and has stronger capacity to search molecules with good properties. In the future, we will apply our G-SphereNet to more complicated 3D structures, such as proteins and many-body particle systems.
REPRODUCIBILITY STATEMENT
We have provided the source codes of our method in supplementary materials and experiment details in Appendix C for reproducing the results.

heading: A PROOF OF THE EQUIVARIANCE AND INVARIANCE
text: In G-SphereNet, at the i-th generation step, we generate the relative distance di, θi, and ϕi by the model as di, θi, ϕi = g ( zdi , z θ i , z ϕ i ;Ai, Ri ) . (13) Without loss of generality, here we assume i ≥ 3. Note that g is symmetry-invariant because we extract features using symmetry-invariant SphereNet (Liu et al., 2021d) model, hence we have
di, θi, ϕi = g ( zdi , z θ i , z ϕ i ;Ai, RiQ T + 1bT )
(14)
for any orthogonal matrix Q ∈ R3×3 and translation vector b ∈ R3. The coordinate ri of the new atom is then computed from di, θi, ϕi and the coordinates rf , rc, re. Next, we show that this generation process satisfies the equivariance and invariance properties. Theorem 1. If ri = gr(zri ;Ai, Ri), then we have Qri + b = gr ( zri ;Ai, RiQ T + 1bT )
for any orthogonal matrix Q ∈ R3×3 and translation vector b ∈ R3.
Proof. Here, we can get that zri = (z d i , z θ i , z ϕ i ), and gr(zri ;Ai, Ri) = h ( g ( zdi , z θ i , z ϕ i ;Ai, Ri ) ; rf , rc, re ) ,
where the function h(d, θ, ϕ; rf , rc, re) is defined as
h(d, θ, ϕ; rf , rc, re) = rf + d cos θ(rc − rf ) ||rc − rf ||22 + d sin θ(re,ϕ − re,cf ) ||re,ϕ − re,cf ||22 .
In this equation, re,cf is the coordinate of the projection of e on the line (rf , rc), and re,ϕ is the coordinate of e after rotating the plane (rf , rc, re) along the line (rf , rc) by the torsion angle ϕ.
For any orthogonal matrix Q ∈ R3×3 , because ||Qr||22 = rTQTQr = rT r = ||r||22, we have
h(d, θ, ϕ;Qrf + b,Qrc + b,Qre + b)
=Qrf + b+ d cos θ(Qrc −Qrf ) ||Qrc −Qrf ||22 + d sin θ(Qre,ϕ −Qre,cf ) ||Qre,ϕ −Qre,cf ||22 =Qrf + b+Q d cos θ(rc − rf ) ||rc − rf ||22 +Q d sin θ(re,ϕ − re,cf ) ||re,ϕ − re,cf ||22
=Q [ rf +
d cos θ(rc − rf ) ||rc − rf ||22 + d sin θ(re,ϕ − re,cf ) ||re,ϕ − re,cf ||22
] + b
=Qh(d, θ, ϕ; rf , rc, re) + b.
(15)
Combining the conclusions from both Eqn. (14) and (15), we can easily get that Qri + b = gr ( zri ;Ai, RiQ T + 1bT )
holds for any orthogonal matrix Q ∈ R3×3 and translation vector b ∈ R3. Theorem 2. For any orthogonal matrix Q ∈ R3×3 and translation vector b ∈ R3, we have p ( Qri + b|Ai, RiQT + 1bT ) = p (ri|Ai, Ri).
Proof. We can easily find that the relative distance di, angle θi, and torsion angle ϕi will not change if we transform ri and Ri with the same Q and b. In addition, from Eqn. (14), we have
zdi , z θ i , z ϕ i = g −1(di, θi, ϕi;Ai, RiQ T + 1bT ). (16)
Hence, the corresponding latent variables zdi , z θ i , z ϕ i are invariant to the rotation and translation transformation. Since p(ri|Ai, Ri) = pZd(zdi |Ai, Ri)pZθ (zθi |Ai, Ri)pZϕ(z ϕ i |Ai, Ri), we can get
that p ( Qri + b|Ai, RiQT + 1bT ) = p (ri|Ai, Ri) is right for any orthogonal matrix Q ∈ R3×3 and translation vector b ∈ R3.

heading: B GENERATION AND TRAINING ALGORITHM
text: Algorithm 1 Generation Algorithm of G-SphereNet 1: Input: G-SphereNet model, latent distribution pZa , pZd , pZθ , pZϕ , maximum number of atoms n
2: 3: Initialize molecular geometry G1 with one carbon atom, whose coordinate is r0 = [0, 0, 0] 4: for i = 1 to n− 1 do 5: zai ∼ pZa 6: Generate ai from zai 7: Get the candidate focal atom set by the atom-wise classifer. 8: if the candidate focal atom set is empty then 9: Output Gi
10: else 11: Random select the focal atom f from the candidate focal atom set 12: end if 13: zdi ∼ pZd , zθi ∼ pZθ (if i ≥ 2), z ϕ i ∼ pZϕ (if i ≥ 3) 14: Generate di, θi (if i ≥ 2), ϕi (if i ≥ 3) from zdi , zθi , z ϕ i 15: Get ri from di, θi, ϕi and rf . 16: Add a new node with type ai and coordinate ri to Gi and set the updated geometry as Gi+1 17: end for 18: Output Gn
Algorithm 2 Training Algorithm of G-SphereNet 1: Input: Molecular geometry datasetM, G-SphereNet model with trainable parameter ω, latent
distribution pZa , pZd , pZθ , pZϕ , learning rate α, batch size B 2: 3: repeat 4: Sample a batch of B molecular graphs G fromM 5: L = 0 6: for G ∈ G do 7: Set n as the number of atoms in G 8: Order the atoms in G 9: for i = 1 to n− 1 do
10: Get ai, di, θi (if i ≥ 2), ϕi (if i ≥ 3) 11: Get zai , z d i , z θ i (if i ≥ 2), z ϕ i (if i ≥ 3) 12: L = L− log pZa(zai )− log pZd(zdi ) 13: L = L− log pZθ (zθi ) if i ≥ 2 14: L = L− log pZϕ(z ϕ i ) if i ≥ 3 15: Add the binary cross entropy loss for the focal atom selection to L 16: end for 17: end for 18: L = LB 19: ω = ω − α∇ωL 20: until ω is converged 21: Output G-SphereNet model with parameter ω

heading: C EXPERIMENT DETAILS
text: Model configuration. For conditional information extraction, the input 3D molecular geometry is first processed into a cutoff graph. Specifically, any two nodes whose distances are lower than 5.0 Å are connected. The node features are initialized to the one-hot vectors of atom types and the edge features are initialized by spherical basis functions as in Liu et al. (2021d). We use the SphereNet (Liu et al., 2021d) model with 4 layers to extract features from the input geometry, where the input embedding size is set to 64 and output embedding size is set to 256. Afterwards, global features are extracted by the multi-head attention network with 4 attention heads. In addition, we employ 6 flow layers. Such model configuration is used for all experiments.
Training and generation details. In the random molecular geometry generation task, the GSphereNet model is trained with Adam optimizer for 100 epochs, where the learning rate is 0.0001 and the batch size is 64. We report the results corresponding to the epoch with the best validation loss. In the target molecule discovery task, the model is fine-tuned with a learning rate of 0.0001, a batch size of 32. The number of training epochs is 40 for the HOMO-LUMO gap and 80 for the isotropic polarizability. During generation, we use temperature parameters in the prior Gaussian distributions. Specifically, we change the standard deviation of the Gaussian distribution to the temperature parameter. We use 0.5 for sampling zai (i ≥ 1), 0.3 for sampling zdi (i ≥ 1), 0.4 for sampling zθi (i ≥ 2), and 1.0 for sampling z ϕ i (i ≥ 3).
Focal node selection. In G-SphereNet, the atom-wise MLP takes the feature of each atom as input, and outputs a binary classification score. We consider the atoms whose scores are higher than 0.5 as being classified to candidate focal atoms, and select the exact focal atom from candidate focal atoms. In our experiments, we do not observe any unstable training for this classifier. This atomwise classifier is trained to classify whether an atom is not valence full filled so that it can be a candidate for next focal atom selection or not. We believe it is not a hard binary classification task so there is no unstable training around the threshold and the training is not sensitive to the threshold choice. Differently, in Simm et al. (2020), the MLP and softmax function directly outputs the probability of being selected to the focal atom for each atom. This MLP is trained by using the sampled focal atom as targets. From the results in Table 3(c), we can clearly find that our method results in much higher chemical validity than the method in Simm et al. (2020).

heading: D MORE EXPERIMENT RESULTS

Review summaries:
The paper proposes a new way to directly generate 3D molecular geometries from scratch. They output the atoms one by one, including the type of the atom and the related position from one chosen focal atom. They use a SphereNet (Liu et al. 2021d) - a spherical message-passing network and an additional multi-head attention network (Vaswani et al. 2017) to process the existed information in the partial generation results. They have shown in 2 different tasks that they are better than previous methods.

The paper introduces a new generative model for 3D molecular geometry, taking the form of a graph embedded in 3D where each node is an atom, and each edge is a bond. Differently from existing methods, the proposed pipeline attempts to generate the 3D molecule from scratch, instead of starting, e.g., from a 2D molecular graph as other methods do. The pipeline is sequential and is based on autoregressive flow models, which have been previously used for the generation of 2D molecular graphs, but not for the present task. The paper convincingly shows that the proposed model works better than competing approaches for the same task.

The paper focuses on a generative model for 3D molecules that generates full 3D molecular gemoetry, instead of generating a molecular graph and then generating the 3D molecular geometry conditional on the graph. The method using an autoregressive flow model, which has been applied to molecular graph generation but not full geometry. This enables some new applications such as molecular design to target a geometry-dependent objective function.  

The paper addresses the problem of generating 3D molecular geometries. The proposed G-SphereNet takes SphereNet as the backbone to generate 3D molecular geometries from scratch. The generation is done in a step-by-step fashion, with each step generating a (distance, angle, torsion) triplet, conditioning on features extracted by an attention mechanism.