A N A UTOREGRESSIVE F LOW M ODEL FOR 3D M OLECULAR G EOMETRY G ENERATION FROM S CRATCH
A BSTRACT
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 ﬂow model for generating 3D molecular geometries.
G-SphereNet employs a ﬂexible 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.
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.
Our code is publicly available as part of the DIG package ( https://github.com/divelab/DIG ).
I NTRODUCTION
Designing and synthesizing novel molecules with desirable properties is a challenging task in drug discovery (Wang et al., 2022; Stokes et al., 2020) and chemical science.
The size of search space of all chemical molecules is estimated to be on the order of 10 33 (Polishchuk et al., 2013), thereby making exhaustive search infeasible.
In recent years, advances in machine learning methods have greatly accelerated the progress in this ﬁled.
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 (Sch¨utt 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., 2022).
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.
B ACKGROUND AND R ELATED W ORK 2. 2D M OLECULAR G RAPH G ENERATION
Recently, thanks to the advances of deep generative models, signiﬁcant progress has been made in the problem of molecule design and generation.
Some methods (G´omez-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.
2. 3D M OLECULAR G EOMETRY G ENERATION FROM 2D I NFORMATION
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.
2. 3D M OLECULAR G EOMETRY G ENERATION FROM S CRATCH
In this work, we consider the problem of generating 3D molecular geometries from scratch.
Let
G = { G j } mj =1 be a set of 3D molecular geometries, and the function S ( G ) ∈
R computes a speciﬁc quantum property score of G .
We consider the two generation tasks deﬁned 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 E G ∼ 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 (Sch¨utt 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 & No´e, 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 ﬂow 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 deﬁnes the prior distribution in a subspace of the latent space to ensure translation invariance.
Different from these methods, our method employs a ﬂexible sequential generation pipeline based on autoregressive ﬂow models, which can capture the density of 3D molecular geometries more effectively.
2. F LOW M ODELS
A ﬂow model deﬁnes a parametric invertible mapping f θ : z ∈ R d → x ∈ R d , where the data point x and the latent variable z are both random variables.
Given that z is sampled from a known prior distribution p Z and f θ is invertible, we can compute the log-likelihood of x as
[Equation]
where J = ∂f − 1 θ ( x ) ∂x is the Jacobian matrix.
To train the ﬂow model on a given dataset X = { x i } mi =1 , the log-likelihoods of data points are computed from Eqn.
(1) and maximized by gradient descent.
Hence, tractable and cheap computation of det J is needed for efﬁcient training.
A common choice of f θ in most ﬂow models is the afﬁne coupling mapping (Dinh et al., 2014; 2016), in which case computing det J 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; K ¨ ohler 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 ﬂow 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 ﬂow 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.
M ETHODS
While autoregressive ﬂow models have been successfully applied to the molecular graph generation task, it remains unclear whether they are sufﬁciently 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 ﬂexible, effective, and efﬁcient sequential generation pipeline based on autoregressive ﬂow 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., 2022) 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 ﬁrst likelihood-based autoregressive generative model for 3D molecular geometry generation.
3. S EQUENTIAL G ENERATION
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 ∈ R n ×is the atom coordinate matrix.
Each row in the matrix
A is a one-hot vector, and A [ j, u ] =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 G 1 with one carbon atom at the origin point, and generate the complete geometry by adding a new atom at each step.
Speciﬁcally, at the i -th step, let the intermediate 3D molecular geometry generated from the previous i − 1 steps be G i
=
( A i , R i ) , which has i atoms.
The atom type a i ∈ { 0 , 1 } k of the new atom is generated by the generative model g a based on the latent variable z ai .
Afterwards, the generative model g r decides the 3D Cartesian coordinate r i ∈ R 3 of the new atom based on the latent variable z ri .
Both g a
and g r are autoregressive functions of intermediately generated geometries.
The overall sequential generation process can be described by the following equations:
[Equation]
[Figure]
3. G ENERATION WITH A UTOREGRESSIVE F LOW M ODELS
We employ autoregressive ﬂow models to generate the atom type a i 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
[Equation]
where U (0 , 1) is the uniform distribution over the interval (0 , 1) .
To generate a i , we ﬁrst sample the latent variable z ai ∈ R k from the standard Gaussian distribution N (0 , 1) , and then map z ai to ˜ a i by the afﬁne transformation as
[Equation]
where (cid:12) denotes the element-wise multiplication, the scale factor s ai and the shift factor t ai both depend on the conditional information extracted from the intermediate geometry G i =
( A i , R i ) .
Intuitively, ˜ a i should be invariant to any rigid transformation on R i , i.e., ˜ a
i should not change if we rotate or translate R i in 3D space.
Hence, we use a symmetry-invariant model to compute s ai and t ai from G i , which are described in Sec. 3.4 in detail.
After obtaining ˜ a i , a i can be computed by taking the argmax of ˜ a i , as a i = one-hot (arg max ˜ a i ) .
However, we cannot generate the 3D coordinate r i in the same way as the generation of the atom type a i .
Directly calculating r
i with the autoregressive ﬂow model, as in Eqn.
(4), neither satisﬁes the equivariance property of coordinates nor the invariance property of likehood.
First, it is easy to ﬁnd that if we rotate or translate R i , then r i needs to be rotated or translated correspondingly.
Formally, it means that for any orthogonal matrix Q ∈ R 3 × 3 and translation vector b ∈ R 3 , if r i =
g r
( z ri ;
A i , R i ) , we have
[Equation]
where 1 denotes a vector of all ones of length i .
If we directly compute r i by the autoregressive model similar to Eqn.
(4), i.e., r i =
s ri (cid:12) z ri + t ri , then to satisfy Eqn.
(5), the correctness of s
Qr
+ b i (cid:12) z ri +
[Equation]
Qr +
[Equation]
i = Q [ s ri (cid:12)
[Equation]
ri ] + Qt ri + b has to be ensured for any orthogonal matrix Q and translation vector b .
However, it is very hard to design a ﬂow model satisfying this complicated condition.
Second, the likelihood p ( r i | A
i , R i ) should be invariant to rotations and translations as they do not change the 3D structure.
In other words,
[Equation]
6) should be satisﬁed for any orthogonal matrix Q and translation vector b .
It follows from the changeof-variable theorem that p ( r i |
A i , R i ) = p ( z ri )  det ∂z ri ∂r
i  .
Let the latent variable corresponding to
Qr i + b be z
Qr
+ b i .
The afﬁne transformation in Eqn.
(4) does not admit any relationship between either p ( z ri ) and p ( z Qr + b i ) , or ∂z ri ∂r
i and ∂z Qr
+ b i ∂
( Qr i + 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 ﬁrst choose a focal atom among all i atoms in G i , which serves as the reference point for the location of the new atom.
Then, the model generates the distance d
i , the angle θ i , and the torsion angle ϕ i successively.
Speciﬁcally, assuming that the focal atom is the f -th atom of G i , the distance d i =
||
r
i − r f || 2 from the focal atom to the new atom is ﬁrst generated.
Then, if i ≥ 2 , the angle θ i ∈ [0 , π ] between the lines ( r f , r i ) and ( r f , r c ) is generated, where c is the atom closest to f in G i .
Finally, if i ≥ 3 , the torsion angle ϕ i ∈ [ − π, π ] between the planes formed by positions ( r f , r c , r i ) and ( r f , r c , r e ) is generated, where e is the atom closest to c but different from f in G i .
Similar to ˜ a i , d i , θ i , and ϕ i are generated as
[Equation]
where z di , z θi , z ϕi ∈ R are all latent variables sampled from standard Gaussian distributions, and the scale factors s di , s θi , s ϕi ∈ R and the shift factors t di , t θi , t ϕi ∈ R are functions of G i .
Afterwards, the coordinate r i is computed from the relative positional elements d
i , θ i , ϕ i and the coordinates r f , r c , r e .
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 classiﬁer.
Our proposed sequential generation method is related to SphereNet (Liu et al., 2022) 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.
3. D ISCUSSIONS
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 & No ´ e, 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 ﬂexible 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 efﬁcient and effective.
E-NFs maps latent variables to 3D coordinates of atoms by a ﬂow model.
Since coordinates are not translation invariant, E-NFs proposes to obtain latent variables by computationally expensive operations.
Speciﬁcally, E-NFs ﬁrst samples from the prior distribution deﬁned 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.
3. C ONDITIONAL I NFORMATION E XTRACTION
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., 2022), 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 e k and the feature of i -th node by v i , SphereNet initializes them with spherical basis functions and updates them to e (cid:48)
k
and v (cid:48)
i with spherical message passing as
[Equation]
[Equation]
Here, φ e , φ v are updating functions, ρ p → e , ρ e → v are aggregation functions, r k and s k are the sending and receiving nodes of the k -th edge, R s k are the coordinates of neighboring nodes of s k , E s k and
E i are the features of edges incident to node s k 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 G i , let the node embeddings computed from the SphereNet be { h i,j } i − 1 j =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 classiﬁcation scores are higher than 0.5.
If all the classiﬁcation scores outputted from the atom-wise classiﬁer 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 insufﬁcient 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 matrix Q , key matrix K , and value matrix V as inputs, and extract global information from inputs by the multi-head attention mechanism:
[Equation]
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 W Qi , W Ki , W Vi [Equation] and W O are all trainable parameters.
Let the node embedding matrix of G i be H
i =
[ h i, 0 , ..., h i,i − 1 ] T , s ai and t ai are computed as
[Equation]
where MH-ATT a is a multi-head attention network and MLP a is a multi-layer perceptron.
As for the scale and shift factors in Eqn.
[Equation] we ﬁrst multiply node embeddings with the atom type embedding vector to include the atom type information:
[Equation]
then compute them as
[Equation]
, i ≥ 2 ,
Con
[Equation]
H i (cid:17)(cid:17)(cid:17) , i ≥ 3 .
Here Embedding ( · ) is a lookup based embedding layer, MH-ATT d , MH-ATT θ , MH-ATT ϕ are all multi-head attention networks, and MLP d , 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.
3. T RAINING
To train the G-SphereNet model on a dataset, we ﬁrst 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.
[Table]
For a 3D molecular geometry G having n atoms ( n > 3) , we maximize its log-likelihood to train the G-SphereNet model.
Speciﬁcally, 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
[Equation]
12) where latent variables z ai , z di , z θi , z ϕi are computed by reversing the mapping in Eqn.
(4) and (7), and p Z a , p Z d , p Z θ , p Z ϕ are all standard Gaussian distributions.
Besides, the atom-wise classiﬁer 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 ﬁlled, otherwise is 0.
We describe the detailed training and generation algorithm of G-SphereNet in Appendix B.
E XPERIMENTS
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.
4. R ANDOM M OLECULAR G EOMETRY G ENERATION
Data.
For the random molecular geometry generation task, we evaluate G-SphereNet on the QM (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.
Speciﬁcally, all generated 3D molecular geometries are converted to molecular graphs by the method proposed in Kim & Kim (2015), and the Validity is deﬁned 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 ﬁnd 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
[Table]
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., 2022) 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 conﬁguration and other training details.
G-SphereNet 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 & No ´ e, 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.
4. T ARGETED M OLECULE D ISCOVERY
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 ﬁne-tune the G-SphereNet model that has been trained in Sec.
4.1 on the biased datasets.
Speciﬁcally, 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 Bohr 3 .
The G-SphereNet model is then ﬁne-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 ﬁne-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.
Speciﬁcally, we generate 1000 molecular geometries with the trained model, and ﬁlter 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 deﬁned as scores lower than 4.5 eV for the HOMO-LUMO gap and scores higher than 91 Bohr 3 for the isotropic polarizability, respectively.
We ﬁnd that E-NFs (Satorras et al., 2021a) fails to generate enough chemically valid molecular geometries and produce reliable results after ﬁne-tuning, so we only compare our G-SphereNet method with G-SchNet in this task.
[Table]
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 ﬁne-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
[Table]
4. A BLATION S TUDIES
In previous sections
[Table]
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
[Table]
and the consideration of distance, angle
[Table]
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 (Sch ¨ utt 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 classiﬁcation ﬁnding atoms that are not valence full ﬁlled.
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 beneﬁts of our method, we compare it with a G-SphereNet variant replacing our focal atom selection method by the one in Simm et al.
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.
C ONCLUSION
We propose G-SphereNet, a novel autoregressive ﬂow 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 ﬂexible and efﬁcient, 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.
R EPRODUCIBILITY S TATEMENT
We have included the source codes of our method in DIG (Liu et al., 2021b) library and provided experiment details in Appendix C for reproducing the results. A CKNOWLEDGMENTS
We thank Limei Wang for her help on explaining implementation details of the SphereNet model and Xuan Zhang for his guidance on using PySCF package.
This work was supported in part by
National Science Foundation grant IIS-
A P ROOF OF THE E QUIVARIANCE AND I NVARIANCE
In G-SphereNet, at the i -th generation step, we generate the relative distance d i , θ i , and ϕ i by the model as
[Equation]
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., 2022) model, hence we have
[Equation]
14) for any orthogonal matrix Q ∈ R 3 × 3 and translation vector b ∈ R 3 .
The coordinate r i of the new atom is then computed from d i , θ i , ϕ i and the coordinates r f , r c , r e .
Next, we show that this generation process satisﬁes the equivariance and invariance properties.
Theorem 1.
If r i =
g r
( z ri ; A i , R i ) , then we have Qr i +
b
=
g r (cid:0) z ri ;
A i , R i Q T + 1 b T (cid:1) for any orthogonal matrix Q ∈ R 3 × 3 and translation vector b ∈ R 3 .
Proof.
Here, we can get that z ri =
( z di , z θi , z ϕi ) , and
[Equation]
where the function h ( d, θ, ϕ ; r f , r c , r e ) is deﬁned as
[Equation]
In this equation, r e,cf is the coordinate of the projection of e on the line ( r f , r c ) , and r e,ϕ is the coordinate of e after rotating the plane ( r f , r c , r e ) along the line ( r f , r c ) by the torsion angle ϕ .
For any orthogonal matrix Q ∈ R 3 × 3 , because ||
Qr ||
=
r T Q T Qr = r T r = ||
r
|| 22 , we have
[Equation]
Combining the conclusions from both Eqn.
(14) and (15), we can easily get that
Qr i + b = g r (cid:0) z ri ; A i , R i Q T + 1 b T (cid:1) holds for any orthogonal matrix Q ∈ R 3 × 3 and translation vector b ∈ R 3 .
Theorem 2.
For any orthogonal matrix Q ∈ R 3 × 3 and translation vector b ∈ R 3 , we have p
[Equation]
Proof.
We can easily ﬁnd that the relative distance d i , angle θ i , and torsion angle ϕ i will not change if we transform r i and R i with the same Q and b .
In addition, from Eqn.
(14), we have
[Equation]
Hence, the corresponding latent variables z di , z θi , z ϕ i are invariant to the rotation and translation transformation.
Since p ( r i | A i , R i )
=
p Z d ( z d i |
A
i , R i ) p
Z θ ( z θ i |
A i , R i ) p Z ϕ ( z ϕ i |
A
i , R i ) , we can get that p (cid:0)
Qr i + b | A
i , R i Q T + 1 b T (cid:1) =
p ( r i |
A
i , R i ) is right for any orthogonal matrix Q ∈ R 3 × 3 and translation vector b ∈ R 3 .
B G ENERATION AND T RAINING A LGORITHM
1: Input: G-SphereNet model, latent distribution p Z a , p Z d , p Z θ , p Z ϕ , maximum number of atoms n 2: 3: Initialize molecular geometry G 1 with one carbon atom, whose coordinate is r 0 =
[0 , 0 , 0] 4: for i = 1 to n − 1 do 5:
z ai ∼ p
Z a 6:
Generate a i from z ai 7:
Get the candidate focal atom set by the atom-wise classifer.
8: if the candidate focal atom set is empty then 9:
Output G i 10: else 11:
Random select the focal atom f from the candidate focal atom set 12: end if 13: z d i ∼ p Z d , z
θ i ∼ p
Z θ (if i ≥ 2) , z ϕ i ∼ p
Z
ϕ
(if i ≥ 3 ) 14:
Generate d i , θ i (if i ≥ 2 ), ϕ i (if i ≥ 3 ) from z di , z θi , z ϕi 15:
Get r i from d i , θ i , ϕ i and r f .
Add a new node with type a i and coordinate r i to G i and set the updated geometry as G i + 17: end for 18: Output G n 1: Input:
2: 3: repeat 4:
Sample a batch
B graphs from M 5:
=6: for G G do 7:
Set the
G 8: Order the atoms in G 9: for i = 1 to n − 1 do 10:
Get a i , d i , θ i (if i ≥ 2 ), ϕ i (if i ≥ 3 ) 11: Get z ai , z di , z θi (if i ≥ 2 ),
[Equation]
[Equation] [Equation]
[Equation]
Add the binary cross entropy loss for the focal atom selection to L 16: end for 17: end for 18:
[Equation]
[Equation]
∇
[Equation]
20: until ω is converged 21: Output G-SphereNet model with parameter ω
C E XPERIMENT D ETAILS
Model conﬁguration.
For conditional information extraction, the input 3D molecular geometry is ﬁrst processed into a cutoff graph.
Speciﬁcally, 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.
We use the SphereNet (Liu et al., 2022) 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 ﬂow layers.
Such model conﬁguration 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.
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 ﬁne-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.
Speciﬁcally, we change the standard deviation of the Gaussian distribution to the temperature parameter.
We use 0.5 for sampling z a i ( i ≥ 1) , 0.3 for sampling z d
i ( 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 classiﬁcation score.
We consider the atoms whose scores are higher than 0.
as being classiﬁed 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 classiﬁer.
This atomwise classiﬁer is trained to classify whether an atom is not valence full ﬁlled so that it can be a candidate for next focal atom selection or not.
We believe it is not a hard binary classiﬁcation 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 ﬁnd that our method results in much higher chemical validity than the method in Simm et al.
D M ORE E XPERIMENT R ESULTS
[Figure]