TensorVAE: A Direct Generative Model for Molecular Conformation Generation driven by Novel Feature Engineering
Authors: Hongyang Yu
Hongjiang Yu
Venue: Submitted to ICLR 2023 (2023)
Keywords: generative model
feature engineering
molecular conformation generation

Abstract: Efficient generation of 3D conformations of a molecule from its 2D graph is a key challenge in in-silico drug discovery. Deep learning (DL) based generative modelling has recently become a potent tool to tackling this challenge. However, many existing DL-based methods are either indirect-leveraging inter-atomic distances or direct-but requiring complex feature transformation or numerous sampling steps to generate conformations. In this work, we propose a simple model abbreviated TensorVAE capable of generating conformations directly from a 2D molecular graph in a single step. The main novelty of the proposed method is focused on feature engineering. We develop a novel encoding and feature extraction mechanism relying solely on standard convolution operation to generate token-like feature vector for each atom. These feature vectors are then transformed through standard transformer encoders under a conditional Variational Auto Encoder framework for learning to generate conformations directly. We show through experiments on two benchmark datasets that with intuitive and sensible feature engineering, a relatively simple and standard model can provide promising generative capability rivalling recent state-of-the-art models employing more sophisticated and specialized generative architecture.

Under review as a conference paper at ICLR 2023
TENSORVAE: A DIRECT GENERATIVE MODEL FOR MOLECULAR CONFORMATION GENERATION DRIVEN BY NOVEL FEATURE ENGINEERING
Anonymous authors Paper under double-bling blind review
ABSTRACT
Efficient generation of 3D conformations of molecule from its 2D graph is key challenge in-silico drug discovery. Deep learning (DL) based generative mod- eling has recently become potent tool to tackling this challenge. However, many existing DL-based methods are either indirect-leveraging inter-atomic distances direct-but requiring numerous sampling steps generate conformations. In this work, we propose a simple model abbreviated Tensor capable of generat- ing conformations directly from 2D) molecular graph in single step. The main novelty of the proposed method is focused on feature engineering. We develop novel encoding and feature extraction mechanism relying solely on standard convolution operation to generate token-like feature vector for each atom. These feature vectors are then transformed through standard transformer encoders un- der a conditional Variational Autoencoder framework for generating conforma- tions directly. We show through experiments on two benchmark datasets that with intuitive feature engineering, a relatively simple and standard model can provide promising generative capability rivalling recent statc-of-the- models employing more sophisticated and specialized generative architecture.
INTRODUCTION
Recent advance in deep learning has enabled sigmificant progress computational drug design (Chen et al.) 2018). Particularly, capable graph-based generative models have been proposed to
generate valid 2D graph representation of novel drug-like molecules (Honda et al.) 2019 Mah-
mood et al., 2021 Yu & Yu 2022). and there is an increasing interest on extending these methods to generating 3D molecular structures which are essential for structured- based drug discovery (Li al. 2021; Simm et al. 2021; Gebauer et al., 2022). A stable 3D structure conformation of molecule specified by the 3D Cartesian coordinates of all its atoms. Traditional molecular dy- namics statistical mechanie driven Monte Carlo methods are computationally expensive, making them unviable for generating 3d molecular structures at scale (Hawkins) (017). In this regard, deep leaming(DD_)-basedd generative methods have become attractive alternative.
based generative methods may be broadly classified into three categories: distance-based, reconstruction-bassee and direct methods. The main goal of distance-based methods is learn- ing probability distribution over the inter-atomic distances. During inference, distance matrices sampled from the learned distribution and converted valid 3D conformations through post- processing algorithms. Two representative methods of this category include GraphDG (Simm & Hemándezz-obaul 2019) and CGCF (Xu et al. (021a). An advantage of modeling distance is its roto-translation invariance property-an important inductive bias for molecular geometry modeling Köhler al. (2020). Additional virtual edges and their distances between 2nd and 3rd neighbors are often introduced to constrain bond angles and dihedral angles crucial to generating a valid con- formation. However, Luo et al. [202I) have argued that these additional bonds are still inadequate to capture structural relationship between distant atoms. To alleviate this issue, DGSM (Luo et al.) 2021 proposed to add higher-order virtual bonds between atoms in an expanded neighborhood re- gion. Another weakness of the distance- based methods is the error accumulation problem; random noise in the predicted distance can be exaggerated by an Euclidean Distance Geometry algorithm, leading to generation of inaccurate conformations (Xu al. 2022 2021b).
3nd

---

Under review as a conference paper at ICLR 2023
To address the above weaknesses, reconstruction- methods directly model distribution over 3D coordinates. Their main idea is to reconstruct valid conformations from distorted coordinates. GeoDiff (Xu et al., 2022) and Uni-Mol (Zhou et al., 2022) pioneering studies in this respect. Though sharing similar idea, they differ in the process of transforming corrupted coordinates to sta- ble conformations. While GeoDiff adapts reverse diffusion process (Sohl-Dickstein et al.) 2013), Uni-Mol treats conformation reconstruction as an optimization problem. Despite their promising performance, both methods require designing of task-specific and complex coordinate transfor- mation methods. This is to ensure the transformation is into-translation or -Eu3iaaquv.riian. To achieve this, GeoDiff proposed a specialized other hand, Uni-Mol accomplished the same by combining a task-specific adaption of transformer (aaswani et al. 2017) inspired by the AlphaFold's Evoformer (Jumper et al. (021)) with another spe- cialized equivariant prediction head (Satorras et al. (202T). Furthermore, GeoDiff requires numerous diffusing steps to attain satisfactory generative performance which can be time consuming.
-quivvaaan Markov transition kernel. On the
CVGAE (Mansimov et al. 2019) and DMCG (Zhu et al. 2022) have attempted to resolve the gen- crative efficiency issue by developing models that can produce a valid conformation directly from 2D molecular graph in single sampling step. Regrettably, the performance of CVGAE is sig- nificantly worse than its distance-based counterparts mainly due to the use of inferior graph neural network for information aggregation (Zhu et al. 2022). DMCG aimed to improve the performance of its predecessor by using a more sophisticated graph neural network and a loss function invariant to symmetric permutation molecular substructures. Although DMCG achieved superior perfor- mance, acquiring such loss function requires enumerating all permutations of molecular graph, which can become computationally expensive for long-sequence molecules.
Regardless of their category, a common recipe of success for these models can be distilled to devel- oping model architecture with ever increasing sophistication and complexity. There is little attention on input feature engineering. In this work, we forgo building specialized model architecture but in- stead focus on intuitive input feature engineering. We propose encode molecular graph using fully-connected and symmetric tensor. For preliminary information aggregation, we run rect- angle kernel filter through the tensor in a ID convolution manner. This operation has a profound implication; with a filter size of 3, the information from two immediate neighbors as well all their connected atoms can be aggregated onto the focal atom in a single operation. It also generates token- like feature vector per atom which can be directly consumed by standard transformer encoder for further information aggregation.
The generative framework follows the standard conditional variational autoencoder (CVAE) setup- We start with building two input tensors with one encoding only the 2D molecular graph and the other also encoding 3D coordinate and distance. Both tensors go through the same feature engi- neering step and the generated feature vectors are fed through two separate transformer encoders. The output of these two encoders are then combined in an intuitive way form the input for an- other transformer encoder for generating conformation directly. The complete generative model is abbreviated TensorVAE.
In summary, the proposed method has three main advantages. (1) Direct and Efficient, generating conformation direcity from 2D molecular graph single step. (2) Simple, not requiring task- sepecific design of neural network architecture, relying only simple convolution and off-the-shelf transformer architecture; (3) Easy to implement, no custom module required as both PyTorch and TensorFle Flow offer ready-to-use convolution and transformer implementation. These advantages translate directly to excellent practicality of the TensorVAB method. We demonstrate through ex- tensive experiments two benchmark datasets that the proposed TensorVAE, despite its simplicity, can perform competitively against 18 recent State-of-the-are methods for conformation generation and molecular property prediction.
2 METHOD
2. PRELIMINARIES
Problem Definition, We formulate molecular conformation generation as a conditional generation task. Given set of molecular graphs G and their corresponding i.i.d conformations R, the goal
2

---

Under review as a conference paper at ICLR 2023
is to train generative model that approximates the Boltzman distribution, and from which valid conformation conditioned on a molecular graph can be easily sampled in a single step.
Story Line. In the ensuing sections, breakdown the formulation of the proposed method in three novel ideas. We first introduce how the tensor encoding method is derived. Based on the tensor input, propose "naive" model that applies convolutional neural network directly the tensor to generate distance matrix, from which conformation are obtained through an Euclidean Distance Geometry algorithm. We also show that such model can already provide comparable performance to several advanced distance-based methods. Subsequently, we demonstrate how token-like feature vector can be generated from the input tensor by using ID convolution operation. Finally, we elab- orate on how to combine all the components together under a CVAE framework to arrive at the final generative model.
2.2 INPUT TENSOR GRAPH
Graph neural network (GNN) is popular feature extraction backbone for DL-based molecular formation generation. The input of GNN in this case is composed of three components, including atom features, edge features and adjacency matrix. Atom and edge features normally pass through separate embedding steps before being fed to the GNN. Adjacency matrix is then used to determine neighboring atoms for layer-wise information aggregation. Although bond features aggregated onto atom features and vice versa, these two features are maintained separately throughout the mes- sage passing layers (Gilmer et al. 2017; Satorras et al., 202T). Instead of having separate inputs, our first simple idea is to combine them into single input. Specifically, we add an additional dimension to the adjacency matrix, making a tensor, similar to that used in computer vision task.
The diagonal section of the tensor holds the atom features.
We consider three types of atom fea- tures comprising atom type, charge and chirality. Each feature is one- hot encoded and they are stacked to- gether to form a single atom feature vector. There are two variants of the atom feature vector corresponding to two input tensors for the two coders of the CVAE: an encoder con- ditioned only on graph (for which the tensor is referred to the G tensor) and the other conditioned on both graph and coordinates (for which the tensor is referred
Atom type encodling
1
charge encoding
Single
C
Double bond
C
chirality eticodiigg
coord
BUB
DD
Bond stereochem
the GDR ten- sor). For the GDR tenosr, every atom feature vector has three additional channels incorporating the 3D coor- dinate of the respective atom, and a distance channel filled with zeros.
X
0
540
0.5
0
feature vectors summed Coordinate channels zeroed Stacked with bond feature vector
Bond type
ring size
Normalized Bond length
Euclidean
distance
Benzene ring tensor graph example. Note that the values in the feature vector and its dimension are for demon- stration purpose only. We explain how they are determined in Sec
Figure 1:
The off-diagonal section holds the bond features. The considered bond features are bond type. bond stereo- chemistry type. associated ring size and normalized bond length. A virtual bond is also included in the bond type. The normalized bond length is calculated as edge length (1 for direct neighbor, 2 for 2nd neighbor, etc.) divided by the longest chain length. It is worth noting that all high-order virtual bonds share the same virtual bond type: they only differ in their normalized bond length. To construct bond feature vector, we first sum the atom feature vectors of the related atoms. This new vector is then stacked with one-hot encoded bond type vector, normalized bond length, and one-hot encoded ring size vector to become the bond feature vector.
There are also two variants of the bond feature vector. For the G tensor, coordinate and distance channles excluded from all bond feature vectors. For the GDR tensor, to match the size of the
3

---

Under review as a conference paper at ICLR 2023
atom feature vector, every bond feature vector has three more coordinate channels filled with Os, and an additional distance channel holding the Euclidean distance between two connected atoms. This bond feature vector is obtained for all atom pairs, making the proposed tensor fully-connected and symmetric. Despite being referred differently, the structure of the bond feature vector and atom
feature vector are the same. For both types in the GDR tensor, there are 5 blocks (3 blocks for G tensor) of channels stacked in the exact same order as following. Therefore, convolution over the input tensor is uniform for both on-diagonal and off-digonal channels.
atom feature channels (atom type, charge, and chirality);
atom coordinate channels (coordinate channels are excluded the tensor);
bond type feature channels (bond type and normalized bond length);
Euclidean distance channel (pairwise distance channel is excluded in G tensor);
other bond type feature channels (bond stereo-chem type and bond ring size)
An example input tensor graph of the benzene ring is illustrated in Fig Having obtained the tensor representation, a naive way of building generative model is to apply convolutional neural network directly on the tensor, and train it to predict a distribution over the inter-atomic distances. We utilize a standard UNet (Ronneberger et al.) 2015) structure to map the input tensor to a probability distribution over distance matrix containing all pair-wise Euclidean distances. Distance matrices are then sampled and converted to valid conformations following the same method presented in GraphDG (Simm& & Hernández-Lobato 2019). We refer to this model as the NaiveUNet, More details of the NaiveUNet can be tound in Sec A.3 and a further explaination of its poor performance can be found in Sec
Despite its naive nature, this model achieves a mean coverage (COV) score of 52. 14 1.48% and a mean matching (MAT) score of 0.4322 on the GEOM-Drugs dataset, already comparable to several more complex distance-basee baselines, shown in See 3.2 However, there two major issues to this approach. First, with a small kernel size (3 x used in the UNet), it takes many convolution layers achieve information aggregation between atoms that are far apart; it does not take full advantage of high-order bonds already made available in the input tensor. Secondly, the output size grows quadratically with the number of atoms, as compared to only linear growth in reconstruction-basee or direct generation methods. The solution to the first issue is rather simple, obtained by increasing the kernel size expand its "field of view". On the other hand, solving the second issue requires elevating the naive two-step generative model to direct one.
2.3 EXTENDED KERNEL AND ATTENTION MECHANISM
We observe that every row or column of the proposed tensor contains global information of focal atom and all of its connected atoms (by both chemical and virtual bond). This motivates our second main idea which is to extend the length of the kernel to the length of the tensor graph while keeping the width unaltered. This idea has profound implication; information from the immediate neigh- bors, all their connected atoms, and all the bond features can be aggregated onto the focal atom in a single convolution operation. In contrast, achieving the same aggregation may require many layers of propagation for the naive model and other GNN-based models. A direct consequence of this modification is that only ID convolution is permitted. With multiple kernels being applied simulta- neously, each stride of these kernels generates a feature vector for a single atom. An illustration of the ID convolution operation is shown in Fig 2
3 Wemel
Ni3 Dome/ttm!
CSFR
Figure 2: Extending kernel and ID convolution. We further observe that the generated feature vectors resemble the token-like feature vectors used in language modeling. This observation combined with the proven success of attention mechanism
4

---

Under review as a conference paper at ICLR 2023
in other related work leads to the selection of transformer architecture the backbone of our gen- erative model. A significant advantage of using transformer's 's self-attention mechanism is, similar to the extended kemel,, it enables global information aggregation from and for all atoms. It also eliminates the need to maintain separated atom and bond features at each step of feature transforma- tion. We present further insight and more detailed analysis of the adavantage of this input feature engineering in Sec 1
2.4 PUTTING EVERYTHING TOGETHER
Conditional variational autoencoder framework. We aim at obtaining generative model G) that approximates the Boltzmann distribution through Maximum Likelihood Estimation. Particularly, given set of molecular graphs G' and their respective ground-truth conformations Rt we wish to maximize the following objective.
log pa (RIG) = log p(2) po (R]z, G) dz
A molecular graph can have many random conformations. We assume this randomness is driven by a latent random variable z ~ p (z), where p (=) is a known distribution e.g. a standard normal distribution. As pe (R|z, G) often modeled by complex function e.g. a deep neural network, evaluation of the integral in iq[Tis intractable. Instead. we resort to the same techniques proposed
in the original VAE (Kingma & Welling 2013) to establish a tractable lower bound for Eq []
log (R|G) 2 (=)...) [log pe (R]2, G)] -DL (gw (2|R,G) (z)]
(2)
where Dxz is the Kullback-Leibler divergence and qw (z) R, G) variational approximation of the true posterior p (z) R, G). We assume p (=) = N (0, I) and qu (=|Rt G¹) is a diagonal Gaus- sian distribution whose means and standard deviations are modeled by a transformer encoder. The input of this transformer encoder is the proposed tensor containing both the coordinate and dis- tance information. We denote this tensor the GDR tensor. On the other hand, Pe (R]2, 2, G) is further decomposed into two parts: decoder pe; (R| (G)) for predicting conformation directly and another encoder da, (G) for encoding the 2D molecular graph. The input tensor for an, (G) is ab- sent of coordinate and distance information, and is therefore denoted the G tensor. Both encoders share the same standard transformer encoder structure. However, there is minor modification to the transformer structure for the decoder. Specifically, the Query, Key matrices for the first multi-head attention layer are computed based on the output vectors of dg, (G). and the Value matrices come directly from the reparameterization of the output of que (z|R, G), as 2 fuw + Ene, where pa is sampled from N (0, I). CVAE We also show the illustration of the modified multi-head attention in Fig 3b
and En are the predicted mean and standard deviation respectively. We present the complete picture how the two encoders and the decoder are arranged framework in Fig 3a
Intuition behind the modified attention. There are multiple ways to join together the output of the two encoders to form the input to the final decoder. Popular methods include stacking addi- tion. We tried both these methods with unsatisfactory performance. We notice that, due to direct stacking addition of the sampled output of qu onto the output of da, attention weights computed in the first layer of the decoder are easily overwhelmed by random noise of the sampled values, and become almost indiscernible]] This leads to ineffective information aggregation which is then further cascaded through the remaining attention layers. Intuitively, in the first attention layer, the attention weights dictating how much influence an atom exerts on the other should predominantly be determined by the graph structure, and remain stable for the same molecule. Further, attention weights are computed by Query and Key matrices. Therefore, these two matrices should stay stable for the same graph. This motivates our third and final main idea; that is, we compute Query and Key matrices only from the output (hf,... ..., h of dø and attribute the variation in conformation to the Value matrices which are directly sampled from {21..., 2N } ~ qw- The resultant information aggregation is much more meaningful and each output vector corresponding to an individual atom carries distinct features, facilitating information aggregation of the ensuing attention layers.
Imagine mixture model with randomly varying mixture weights.
5

---

Under review as a conference paper at ICLR 2023
Po: (R (0))
Robscb afgemmen
L-2 tramsformerr encoder
Reconstruction
Nomn
KL Regularization loss
Modifted Vruid heedd
Scaled bet product Adamsion
Nor
Feed Formard
transfomer encader
Nom
Lower
(b) Modified multi-head attention
GAR
Te (G)
qu (z) G)
(a) Variational AutoEncoder framework
Figure 3: Tensor VAE model
Roto-translation invariant loss. Following ConfVAE Xu et 2021b), we formulate the recon- struction loss as.
Raj A ( R). )
L (R) = log Po (R|2, G) -- 1=1
(3)
where A (+) is function aligning the predicted conformation R onto the reference conformation R. We choose Kabsch algorithm (contributors) 2022) as the alignment method which translates its corresponding ground-truth before loss computa- tion. This makes the reconstruction loss roto-translation invariant. Finally, the KL-loss component DEL k fo R2) G) (=)] does not involve any coordinate. Therefore, the objective function defined in Eq is roto-translation invariant.
and rotates the predicted conformation
Direct conformation generation at inference time. To generate single conformation, first construct the G tensor of a molecular graph and obtain a single latent sample {=1,...Zy] from standard diagonal Gaussian distribution. The G tensor passed through de, encoder produce {hf (hf.-..115) 1/ which is then combined with the latent sample via the modified multi-head attention mechanism. The output of this modified attention layer further goes through T- standard attention layers be transformed the final conformation. The entire generation process depends only 2D molecular graph, and requires a single sampling step and single pass of the Tensor VAE model.
3 EXPERIMENT
In this section, we first elaborate the implementation details of the TensorVAE model includ- ing determining the size of the input tensors, network architecture and how the entire framework is trained end-to-end. We then present conformation generation experiment results of the proposed TensorVAE on two benchmark data-sets, GEOM-QMÍ and GEOM-Drugs. These results are com- pared to those of 11 state-of-the-ar baselines. In addition conformation generation, present

---

Under review as a conference paper at ICLR 2023
further experiment molecular property prediction in A.c where the performance of the pro- posed method is compared against 7 more state-of-the-art baselines on the MolecularNet benchmark (Wu et al. 2018).
3.1
EXPERIMENT SETUP
Dataset. Following existing work (Luo et al. 2021 Shi et al. 2021 Xu et al. 2021b 2022 Zhou
al. 2022). we utilize the GEOM data-set for evaluating the performance of the proposed Tensor- VAE. GEOM contains 37 million energy and statistical weight annotated molecular conformations corresponding to 450,000 molecules (Axeirod & Kómez-Bombane divided into two constituent datasets, Drugs and QM9. The Drugs dataset covers 317,000 median- sized molecules averaging 44.4 number of atoms. The QM9 dataset contains 133,000 smaller molecules averaging only 18 atoms. We randomly select 40000 molecules from each dataset to form the training set. For each molecule, we choose the top most likely conformations. This results in 200,000 training conformations for each train set. For validation set, randomly sample 2,500 conformations for both Drugs and QM9 experiments. Finally, for testing, following (Shi et al.) 2021 Xu al. 2022), We randomly select 200 molecules each with more than 50 and less than 500 annotated conformations from QM9, and another 200 with more than 50 and less than 100 annotated conformations from D
2022). This dataset further
Determining input tensor graph size. We conduct a basic data analysis on the entire Drugs dataset to determine the 98.5th percentile of the number of atoms be 69, and the percentage of molecules having more than 69 atoms and with more than 50 but less than 100 conformations is only 0.19%. Accordingly, we set the size of the input tensor to 69 x 69 for Drugs experiment. On the other hand, we use the maximum number of atoms 30 for QM9 experiment. The channel features for the input tensor include atom types, atom charge, atom chirality, bond type, bond stereo-chemistry and bond in in-ring size. For the GDR tensor, also include 3D coordinate channels and the computed distance channel. The resulting channel depth 50 for GDR tensor and 46 for tensor. The detailed information of these features and their encoding method is listed in Sec A.4
Implementation details. We implement the proposed TensorVAE using Tensorflow 2.3.1. All three transformer encoders of Tensor VAE follow the standard Tensorflow implementation in https: //www. twnsor feensorloow... low. httgg//ee//ttttttto/oooo//o heads and latent dimension of 256. Both QM9 and Drugs experiments share the architecture and syyee--praameer configuration. We present the detailed training hyperparameter configuration in Scc|A.2
a s/trrtrra rans former sstraasffrrerrr All of them have 4 layers,
network
Evaluation metrics. We adopt the widely accepted coverage score (COV) and matching score (MAT) (Shi al., 202T) evaluate the performance of the proposed TensorVAB model. These two scores are computed as.
COV CCg C,)
{ RMMSD
RR R) SoVREC,)
(4)
C,
MAT (C (C,,
Σ REC,
min RMSD R,
(5)
where C. is the set of generated conformations and C, is the corresponding reference set. The size of Cy is twice of that of C., as for every molecule, we follow (Xu et al. 2022) to generate twice the number of conformations set to 0.5A for QM9 and 1.25A for Drugs respectively (Shi et al. 2021 RMSD stands for the root-mean-square -eaan-square deviation between R and R, and is computed using the Get Best method in the RDKit (Riniker & Landrum| 2015) package. While COV score measures the ability of a model in generating diverse conformations cover all reference conformations, MAT score measures how well the generated conformations match the ground-truth. A good generative model should have high COV score and low MAT score.
that of reference conformations. is predefined threshold and is
Ranked by their Boltzmann weight. TThis limit the number of conformations for testing molecules is taken directly from https: // which is also followed by all other compared methods.
gi tithub. tittb..com//eeeerpplllaaa

---

Under review as a conference paper at ICLR 2023
Baselines. We compare the performance of the proposed TensorVAE model those of 1 classical RDKit method; 5 distance- based methods including GraphDG, CGCF, ConfVAE, ConfGF and DGSM; reconstrcctionn based methods including GeoDiff and Uni-Mol; 4 direct methods including CVGAF, GeoMol, DMCG and its symmetric permutation variant. The detailed information of the molecular property prediction baselines presented in A.7
3.2 RESULTS AND DISCUSSION
The COV and MAT scores for all compared methods on both QM9 and Drugs datasets are presented Tab[[] All experiments follow the same data generation configuration in (Shi et al.) (021). Additionally, we have conducted three 3 ablation studies on the input feature engineering method in Sec(A. to demonstrate why 1D convolution with x kernel necessary to achieve good generative performance.
Table 1: Performance comparison between Tensor and 11 other SOTAs on GEOM dataset.
QM9
Drugs
Models
COV
(%) f
MAT (Å) 4
COV
(%) f
MAT
(Å) 4
Mean
Median
Mean
Median
Mean
Median
Mean
Median
RDkit
83.26
90.78
0.3447
0.2935
60.91
65.70
1,2026
1.1252
CVGAE
0.09
0.00
1.6713
1.6088
0.00
0.00
3.0702
2.9937
GraphDG
73.33
84.21
0.4245
0.3973
8.27
0.00
1.9722
1.9845
CGCF
78.05
82.48
0.4219
0.3900
53.96
57.06
1.2487
1.2247
ConfVAE
80.42
85.31
0.4066
0.3891
53.14
53.98
1.2392
1.2447
ConfGF
88.49
94.13
0.2673
0.2685
62.15
70.93
1.1629
1.1596
GeoMol
71.26
72.00
0.3731
0.3731
67.16
71.71
1.0875
1.0586
DGSM
91.49
95.92
0.2139
0.2137
78.73
94.39
1.0154
0.9980
GeoDiff
92.65
95.75
0.2016
0,2006
88,45
97.09
0.8651
0.8598
DMCG
94.98
98.47
0.2365
0.2312
91.27
100
0.8287
0.7908
Uni-Mol
97.95
100
0.1831
0.1659
91.91
100
0.7863
0.7794
98.11
100
0.1970
0.1926
94.91
100
0.7789
0.7585
TensorVAE!
+0.25
+0
+0.0016
+0.0027
+0.35
+0
+0.0027
+0.0076
97.11
100
0.2041
0.1920
93.34
99.90
0.8074
0.7927
TensorVAE2
+0.31
+0
+0.0046
+0.007
+1.17
+0.31
+0.0135
#0.01 86
"Bold font indicates best result. Results for RdKit, CVGAE, GraphDG, CGCF, ConfGF are taken from (Shi at. 2021 results for ConfVAE and GeoDiff are taken from (Xu et al.) (2222): all other results are taken from (Zhou et al. 2022): TensorVAE' results and standard deviations are obtained by running 10 experiements each with different random seed single 200 testing molecules set. Tensor results and standard deviations are obtained by running experiements each with a different random seed as well different set of 200
testing molccules.
In general, distance-basee based methods except for DGSM and ConfGF have relatively poor performance compared to that of the classic RDKit. It has been argued that the the performance of RDkit is facilitated by an additional empirical force field (FF) (Hllgren 1996) optimization. Subsequently, CGCF and ConfVAE showed that with FF optimization, they can outperform RDkit. We also show the performance of the proposed TensorVAE with FF optimization compared to 5 other methods also employing FF optimization in Tab in A.c where TensorVAE outperforms all of them with significant margin. However, this additional step further introduces complexity to the already complex two-stage generative model. ConfGF attempted rectify this weakness by simulating a pseudo gradient-basse force field. Such force field can be utilized in an annealed Langevin dynam- ies sampling to sequentially guide atom positions to a valid conformation. DGSM further improves the performance of ConfGF by using dynamically constructed graph structure that is able to model long range atom interaction. Despite being posed as a direct method, both DGSM and ConfGF still need to compute atom distance as an intermediate step. Noticeably, DGSM also requires dynami- cally changing the graph structure for every sampling step. As shown in Tab DGSM and ConGF outperform RDKit by significant margin in both experiments. Nevertheless, their main weakness lies in the fact they they require numerous sampling steps to attain desirable performance. It has been reported in the ensuing work (Xu et al. 2022), that it took ConfGF approximately 8500 sec-
8

---

Under review as a conference paper at ICLR 2023
onds to fully decode 200 QM9 molecules and staggering 11500 seconds for decoding 200 Drugs molecules. In constrast, TensorVAE takes only 62 seconds using a single Xeon 8163 CPU core to decode 200 QM9 molecules, and 128 seconds for 200 Drug molecules.
A main goal of GeoDiff is to be free of the dependence on distance. To achieve this, transformations are directly applied to 3D conformation, through which, random distorted conformation can be sequentially denoised to a valid conformation. An essential requirement for such transformation is that it needs be SE(3) equivariant. This necessitates designing of sophisticated equivariant transition kernel. Despite its claim of not involving computing distance as an intermediate step, the equivariant graph field network (Satorras al. 202I) used still relies on distance-like quantity
| - = for every step of transformation. GeoDiff produces promising performance on both datasets. Unfortunately, needs numerous diffusion steps (T' 5000) for generating conformations, requiring approximately the same amount of decoding time as ConfGF (Xu et al. 2022).
Uni-Mol circumvents this issue by reformulating molecular generation problem as an optimization problem. The structure of Uni-Mol is inspired by the Evoformer proposed in the Alphafold (Jumper al.) 202T) which also considers "atom-pain interaction. This requires maintaining a pair interac- tion matrix (N x N) through every attention layer. In addition, Uni-Mol has significantly larger transformer structure as compared ours, featuring 15 attention layers, 64 attention heads and la- tent dimension size of 512. The Uni-Mol model is also first pretrained molecules, each with 10 conformations), and then fine-tuned on the GEOM dataset for conformation generation. Despite using much smaller and also standard transformer model without pretraining, the proposed TensorVAL outperforms Uni-Mol on the Drugs dataset. Further, we show that the pro- posed model also performs competitively against Uni-Mol a molecular property prediction task in Sec A.7 without any pretraining, validating the effectiveness of the input feature engineering.
much larger dataset (19c
CVGAE is the first method proposed to generate molecular conformation directly from 2D molec- ular graph. However, it yields the worst performance among the compared methods. DMCG at- tempted to revitalize the same framework by adapting a more advanced graph neural network struc- ture combining GATv2 (Brody et al.) (021) with GN block (Battaglia et al., 2018). Noticeably, similar Uni-Mol, it also maintains a N X N bond feature matrix throughout all layers. In con- strast, we only use a standard transformer encoder architecture. Another major modification pro- posed in DMCG is introducing the permutation invariance of symmetric molecular substructures to the RMSD loss. Achieving this invariance requires enumerating all possible permutations which can become expensive for large molecules. The proposed Tensor VAE outperforms DMCG with a significant margin. Finally, some samples of the TensorVAB generated conformations are shown in Sec
4 CONCLUSION
We develop TensorVAE, a simple yet powerful model that is able to generate 3D conformation di- rectly from a 2D molecular graph. Unlike many existing work focusing on designing complex neural network structure, we focus on developing novel input feature engineering techniques. We decom- pose these techniques into three main ideas, and explain how one idea naturally evolves the next. We first propose a tensor representation of a molecular graph. Then, we demonstrate that sliding a rectangle kernel through this tensor in an ID convolution manner can achieve complete information aggregation. Finally, we present the complete CVAE-based framework featuring 2 transformer- based encoders and another mansformer-based decoder, and propose a novel modification the first multi-head attention layer of the decoder to enable sensible integration of the output of the other two encoders. We show through extensive experiments that with intuitive feature engineering, simple and standard model architecture can provide competitive performance compared to 18 recent state- of-the-art models. For future work, we plan to extend our method and methodology to tackling the challenging protein structure prediction.
9

---

Under review as a conference paper at ICLR 2023
5 REPRODUCIBILITY STATEMENT
We did not introduce nor have used any task-specific neural network archiecture. The results pre- sented in this study can be straightforwardly reproduced using publically available datasets and ready-to-use implementation of convolution operation and Transformer from either PyTorch or TensorFlow. Specifically, we the same transformer implementation found in https: haalstttansoorerr To compute RMSD, we first zero-center both the predicted and the ground-truth conformations. Then, we obtain the optimal rotation matrix using the Kabsch algorithm implementation found in httpss://er. wikipedia wwo/oorrii..... Kaabschaahhhhaaaa gor thm] is straightforward to implement this algo- rithm in python function. Please note that if TensorFlow Flow is used, this python function needs to wrapped inside a f ptp.ppy-uuucs.e- on, and then followed by t stop-gradi ent to prevent gradient update. Finally, we rotate the predicted conformation onto its ground-truth for cal- culating the RMSD. The above details together with the already presented hyper-parameter setting should suffice to reproduce results.
7/www. www..tensofffooo
REFERENCES
Simon Axelrod and Rafael for property prediction and molecular generation. Sciemtific Data, 9()):185, 2022, doi: 10. 10.1038/ 2115977022-01288-4
Geom, energy-annotated molecular conformations
1038/ htt.:///........///////110000000
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanzhez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 180.00166, 2018.
Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? arXiv preprint arXiv:2105. 14491, 2021.
Hongming Chen, Ola Engkvist, Yinhai Wang, Marcus Olivecrona, and Thomas Blaschke. The rise of deep learning in drug discovery. Drug discovery today, 1241-1224 2018.
Wikipedia contributors. Kabsch algorithm Wikipedia, the free encyclopedia. https: htmloorithmmooldd
7/en // www.wwdiiiiec... a. www.///gg//iin....
1082253417 2022. [Online; accessed 19-October-2022).
Xiaomin Fang. Lihang Liu, Jieqiong Lei, Donglong He, Shanzhuo Zhang, Jingbo Zhou, Fan Wang, Hua Wu, and Haifeng Wang. Geometry-enhancee molecular representation learning for property prediction, Nature Machine intelligence, 4(2):127-1344 2022.
Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, and Lawrence Carin. Cyclical annealing schedule: A simple approach to mitigating kl vanishing. arXiv preprint 10X1v:1903.10145, 2019.
Niklas WA Gebauer, Michael Gastegger, Stefaan SP Hessmann, Klaus-Rober Müller, and Kristof Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. Nature communications, 13(1):1-11, 2022.
Justin Gilmer, Samuel Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp.
PMLR, 2017.
Thomas A Halgren. Merck molecular force field. extension of mmff94 using experimental data, additional computational data, and empirical rules. Journal of Computational Chemistry, 17(5-6): 616-641, 1996.
Paul CD Hawkins. Conformation generation: the state of the art. Journal of chemical information and modeling, 178)):747-1755, 2017.
Shion Honda, Hirotaka Akita, Katsuhiko Ishiguro, Toshiki Nakanishi, and Kenta Oono. Graph residual flow for molecular graph generation. arXiv preprint arXiv:1909) 13521, 2019.
10

---

Under review as a conference paper at ICLR 2023
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre- pre-training graph neural networks. arXiv preprint arXiv:1905.1222 12265, 2019.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(78733:583-588) 2021.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
/1412. 6980
Diederik Kingma and Max Welling. Auto-encoding variational bayes.
arXiv preprint
1rXiv:1312.6114, 2013.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. International conference on machine learning, pp. 5361-5370. PMLR, 2020.
Yibo Li, Jianfeng Pei, and Luhua Lai. Structure-based de novo drug design using 3d deep generative models. Chemical science, 1264115:3664-13555 2021.
Shengchao Liu, Mehmet F Demirel, and Yingyu Liang. N-gram graph: Simple unsupervised repre- sentation for graphs, with applications molecules. Advances in neural information processing systems, 32, 2019.
Shitong Luo, Chence Shi, Minkai Xu, and Jian Tang. Predicting molecular conformation via dy- namic graph score matching. Advances Neural Information Processing Systems, 34:19784- 19795, 2021.
Omar Mahmood, Elman Mansimov, Richard Bonneau, and Kyunghyun Cho. Masked graph model- ing for molecule generation. Nature communications, 12(1):1-12, 2021.
Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. Molecular geometry pre- diction using a deep generative graph neural network. Scientific reports, 9(1):1-13, 2019.
Sereina Riniker and Gregory Landrum. Better informed distance geometry: using what know to improve conformation generation. Journal of chemical information and modeling, 55(12): 2015.
Yu Rong,
Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on large-scale molecular data. Advances in Neural Information Processing Systems, 33: 12559-12257, 2020.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference Medical image computing and computer- assisted intervention, pp. 234-241. Springer, 2015.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In International conference on machine learning, pp. 9323 9332. PMLR, 2021.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular confor- mation generation. In International Conference on Machine Learning, pp. 9558-9568. PMLR, 2021.
Gregor N. Simm, Robert Pinsler, Gábor Csányi, and José Miguel Hernández-Lobato Symmetry- aware actor-critic for 3d molecular design. In International Conference on Learning Representa- tions, 2021. URL hhtttp:::/ppereviiiiet net wwtp:ntoooommmdddee//.rr..m JPExx>)
Gregor NC Simm and José Miguel Hemnández-Lohato. A generative model for molecular distance geometry. arXiv preprint arXiv: 11099999.11499 2019.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Leam- ing, pp. 2256-2265. PMLR, 2015.
11

---

Under review as a conference paper at ICLR 2023
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learn- ing. Chemical science, 9(2):513-530, 2018.
Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li, Zhao- jun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, et al. Pushing the boundaries of molecular representation for drug discovery with the graph attention mechanism. Journal of medicinal chemistry, 63(16))8749-877) 8760, 2019.
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. arXiv preprint arXiv:2102 102-40 2021a.
Minkai Xu, Wujie Wang, Shitong Luo, Chence Shi, Yoshua Bengio, Rafael
and
Jian Tang. An end-to-end framework for molecular conformation generation via bilevel program- ming. In International Conference on Machine Learning, pp. 11537-11547 PMLR, 2021b.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geo- metric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02200 02923, 2022.
Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman- Perez, Timothy Hopper. Brian Kelley, Miriam Mathea, et al. Analyzing learned molecular repre- sentations for property prediction. Journal of chemical information and modeling, 59(8):3370- 3388, 2019.
Hongyang K Yu and Hongjiang C Yu. Powerful molecule generation with simple convnet. Bioin- formatics, 38( 38(13)33438-3444, 2022.
Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. 2022.
Jinhua Zhu, Yingee Xia, Chang Liu, Lijun Wu, Shufang Xie, Tong Wang, Yusong Wang, Wengang Zhou, Tao Qin, Houqiang Li, al. Direct molecular conformation generation. arXiv preprint arXiv:2002.01336, 2022.
12

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 GLOBAL INFORMATION AGGREGATION BEYOND THE N HOP
A geometric interpretation of GNN's message passing layer it aggregates information between atoms (and their bond) that are 1-hop away. With L layers, information from atoms that are L-hop apart can aggregated. Here, we define global information aggregation as the N aggrega- tion with N being the total number of atoms, where each atom is able to aggregate information from its farthest neighbour.
It is worth noting that for fully-connected GNN, 1-hop message passing can already achieve this global information aggregation. Transformer's self-attention can be considered as a type of fully- connected GNN. However, vanilla transformer only aggregate features from each token/atom; edge features are not included, they needed be incorporated somehow through additional inputs (e.g. the pair interaction matrix of Uni-Mol). The primary reason motivating the creation of the fully-connected tensor representation we want each generated token contain both atom and bond features, such that can eliminate the pair interaction bond matrix. To achieve this, fill each column of the fully-connected tensor with;
focal atom features;
chemical and virtual bond features indicating how the focal atom is connected all other
atoms;
atom features of all connected atoms, since for each cell (except for cell of the focal atom) we atom features of both the connected atom and the focal atom.
In fact, running N x kernel filter on the proposed tensor is conceptually similar to achieving global information aggregation with a fully-connected GNN. By increasing kernel width to 3, the aggregation window also includes global information from two immediate neighbours. This type of information aggregation extends far beyond just 4 -hop.
More interestingly, when multiple kernels applied simultaneously to the same N 3 C region, each kernel is free to choose whichever group of atom//onn features to attend to depending on its kernel weights. This resembles the multi-head attention mechanism of transformer, where each keeadl(head) contributes to portion of the generated feature token. We believe the effective global information aggregation driven by these two (tenor representation 1D Conv) simple yet intuitive ideas is the main reason why the proposed TensorVAB achieves SOTA with much less number of parameters.
A.2 TRAINING HYPERPARAMETERS
Training is conducted on single Tesla V100 GPU. We follow similar learning rate schedule, shown by Eq.3 of the original Transformer paper (Vaswani et al. 2017) but with dmodel = 9612. This results in maximum learning rate of 1.6e To tackle the notorious issue of KL vanishing Fuu et al. 2019). we set a minimum KL. weight of 1e-4 and double it every 62.5e iterations until maximum weight of 0.0256 is reached. We select Adam optimizer (Kingma & Bal 2015) for training. We present some interesting observations of the raining/validatinn this setup in A.5 For both experiments, the Tensor VAE is trained for l iterations with a batch size of 128. The implementation details of NNiveUNet is explained in Sec
curve corresponding to
13

---

Under review as a conference paper at ICLR 2023
A.3 NAIVEUNET MODEL ARCHITECTURE
Encoder
Decoder
Masked ourgur due symmetry
Feature daytt
Skip connection
CUR
154
121
03
Encoder Black
Decoder Block
Figure 4: Naive UNet model. N = 69
We train the above NaiveUNet on the Drugs dataset for 30 epochs with a constant learning rate of 1e and batch size of 32. We follow the same method presented in GraphDGi (Simm & Hernández- Lobato 2019) to convert the predicted distance matrix to conformation.
A.4 ATOM AND BOND FEATURES
We list the atom features and bond features together with the encoding method used to construct the proposed tensor in Tab
Table 2: Atom and bond features used to construct input tensor.
Feature name
Feature value
Encoding method
Atom type
H. C. N. O, S. CI, Br. P, Na, B, Si, Se, K, Bi -2, 0, 2,
one-hot
Atom charge
one-hot
Atom chirality
Unspecified, Tetrahedral.CW Tetrahedral. CCC, Other Single, Double, Triple, Aromatic, Virtual
one-hot
Bond type Normalized bond length
one-hot
real-value
Bond stereochem
StereoNone, StereoAny, StereoZ StereoF, StereoCIS, StereoTrans 3 10
one-hot
Bond in-ring size
one-hot
Coordinate (3 channels)
real-value
Pair wise atom distance
real-value
14

---

Under review as a conference paper at ICLR 2023
A.5 TRAINING AND VALIDATION CURVE
We present the train and validation plots for KL and reconstruction loss based on Drugs dataset in Fig 5a and Fig 5b respectively. Both plots are based
initial KL weight of le- doubling every 62.5k iterations (40 epochs). While KL validation loss reached 18.29 after le iterations (640 epochs), the reconstruction/MMSDD loss reached 0.64À at the end of training. During the first epochs of training. model learning focused on reducing the KL loss due to it orders of magnitude larger than the RMSD loss. We were expecting this trend to continue for while until both losses converge roughly the same range. However, much to our surprise, the model seemed to find way to drastically reduce RMSD loss much earlier by leveraging the information from the GDR encoder; it learned to "cheat" by directly reversing coordinate information embedded in the output of GDR encoder back to the original conformation. The RMSD loss dropped to as low as 0.08Å. On the other hand, the KL loss climbed almost 800, signaling signifcant divergence from standard normal distribution. At this stage, output of the GDR encoder contains informative features of the original 3D coordinates. With the KL loss weight increasing, it becomes more difficult for the model to cheat since training is forcing the output of GDR encoder to conform to a standard uninformative Gaussian distribution. The KL loss started drop while the RMSD loss remained steady, indicating increasing reliance on the output of G encoder for reconstructing the conformation. As the output of GDR encoder becomes less informative, the model learned to rely almost entirely on the aggregated feature from the G encoder to decode conformation.
We attempted initiate the training with a much larger initial KL weight (1e prevent "cheat- ing" from begining. However, this quickly led to the notorious KL vanishing issue (Fu et al. (019)). We figure that "cheating" actually beneficial that it reduces learning difficulty particularly for the decoder; its weights tuned easy training task, simply reversing what GDR encoder has done. In other words, the tuned weights the decoder already hold crucial information on how to decode highly informative input features. As KL weight increases, model learning shifts to make the output of G encoder more informative. Also, this maybe an easier learning task as the RMSD loss is already very low (back-prop propaaarpaa of this loss contributes little to weight update); instead, model learning primarily focuses on optimizing the KL loss. This two-stage iterative loss optimization is much easier than optimizing both losses simultaneously throughout the training process.
3
800
300
0.7
0.2
20
0.06
100
200
300
400
500
600
100
200
300
400
500
600
(a) KL loss.
(b) RMSD loss.
Figure 5: Training and validation plots for Drugs dataset. Orange line: Train; Blue line: Validation
15

---

Under review as a conference paper at ICLR 2023
A.6 GENERATION PERFORMANCE ON DRUGS DATASET WITH FORCE FIELD OPTIMIZATION
Table 3: Performance comparison between methods with FF optimization.
COV
MAT
Method
Mean
Median
Mean
Median
CVGAE + FF
83.08
95.21
0,9829
0.9177
GraphDG + FF
84.68
93,94
0.9129
0.9090
CGCF FF
92.28
98.15
0.7740
0.7338
ConfVAE FF
91.88
100
0.7634
0.7312
GeoDiff + FF
92.27
100
0.7618
0.7340
96.15
100
0.6723
0.6605
TensorVAE + FF
+0.34
+0
+0.0023
+0,0064
*Results for CVGAE, GraghDG, CGCF, ConfVAE and GcoDiff
2021b); Standard deviations are obtained by repeating experiments 10 times each with a different random seed on a test set with 14396 conformations.
taken from (Xu et al.
16

---

Under review as a conference paper at ICLR 2023
A.7 MOELCULAR PROPERTY PREDICTION RESULTS
We conduct further experiment on molecular property prediction to demonstrate the effectiveness of the proposed feature engineering method. Following Uni-Mol (Zhou et al., 2022) and GEM (Fang et al.] )022), we report property prediction result on the MolecularNet QM9 regression task lhttps: / meluulennt. wwt.oorgaaaasssaa/aeaaa
The goal of this task is to estimate homo, lumo, and homo- Lomoo-lumoooooo gap properties of molecules in the QM9 dataset based on their molecular structure. This task is different from the conformation generation experiment in that both 3D co- ordinates and 2D molecular graph of molecule are supplied as input to prediction model. The accuarcy of property prediction depends on how well a model can extract and aggregate such input information among atoms.
Similar to Uni-Mol, we adapt the proposed GDR encoder to this regression task by changing its prediction head. Specifically, 3 (with only attention layers) and add an additional mean pooling layer, which is then followed by a linear layer for property prediction. To obtain training data, we follow the same data train-val- test splif in Uni-Mol and GEM and standardize the output property data. This results 106362 train samples, 13299 val samples and 13356 test samples. We train the adapted model for 300 epochs with a batch size of 128. The learning rate schedule is the same as TensorVAE, We report the mean average error(MAA) over all the test samples.
the same GDR transformer encoder structure presented in
The result of the adapted model is compared to those of 7 other models including;
D-MPNN (Yang al., 2019). AttentiveFP (Xiong et al. 2019) and GEM which are GNN based models without pretraining;
N-Gram (Luu et al., 2019), PretrainingGNN (Hu et al. 2019) and GROVER (Rong et al., 2020) with pretraining. In particular, GROVER integrates GNN into a Transformer ar- chitecture, and there are two variants with different model capacity, GROVER and GROVERiargei
and finally, three variants of Uni-Mol including one without pretraining, another without using the N x N pair representation matrix, and the complete version.
The MAE for all compared methods are summaried in Tab The proposed method produces a com- petitive performance, only underperforming the complete Uni-Mol setup with pretraining and also considering pair representation. This experiment demonstrates that the proposed feature engineering method is very effective global information aggregation.
Table 4: Property prediction result comparison based on MolecularNet QM9 benchmark.
Method
MAE
D-MPNN
0.00814 (0.00001
AttentiveFP
0.00812 (0.00001)
N-Gram
0.00964 (0.00031)
PretrainGNN
0.00922 (0.00004)
GROVER base
0.00984 (0.00055)
GROVER large
0,00986 (0.00025)
GEM
0.00746 (0.00001)
Uni-Mol w/o pair representation
0.00573 (0.00004)
Uni-Mol w/o pretraining
0.00653 (0.00040)
Uni-Mol
0.00467 (0.00004)
GDR encoder (ours)
0.00553 (0.00012)
All results are taken from (Zhou et al. 2022). Values in paren-
thesis are standard deviation obtained by repeating experiments
times.
"The data split with ratio of using scaffold spliting while considering chirality
17

---

Under review as a conference paper at ICLR 2023
A.8 ABLATION STUDIES
In this section, we further demonstrate the effectiveness and necessity of running an 1D convolution with N x 3 kernels over the proposed input tensor through 3 ablation studies.
Why is 1D convolution necessary. We have shown a model based on a x33 kernel in S.3 called UNet Here, we provide a more detailed analysis of why NaiveUNet produces unsatisfactory result. The primary reason for this poor performance is the "field of view" of a conventional d x d (d < N) kernel only sees a partial connection pattern of focal atom. In comparison, N kernel's "field of view" encompasses the complete connection pattern of focal atom. We further observe that when applying 3 x 3 kernel filter the top left region of the proposed tensor, its field of view only includes focal atom, its two neighboring atoms and how the focal atom is connected to them. There are two main disadvantages associated with this. Firstly, it only achieves I-hop information aggregation. Secondly when the 3 kernel moves to an off-diagonal part of the tensor, where most connections are virtual bonds (as atoms of a molecule are often sparsely connected), information aggregation occurs mostly between atoms that are not chemically connected and is therefore less meaningful than that on the diagonal part of the tensor. For these two reasons, the NaiveUNet's performance the GEOM Drugs dataset is the worst as shown in Tab
What happens if we remove all virtual bonds. Notice that if we remove all the virtual bonds in each column and still run N x 3 kernel through the tensor, its "field of view" is "2-hop atomic- environment" (because the focal atom can "see" how neighhoring atoms are chemically connected to all their direct neighbors). Another observation is that after removing all virtual bonds, each column does not correspond to a fully-connected GNN. Therefore it no longer enables a global information aggregation. The conformation generation results of this variant of Tensor on Drugs dataset is shown as as Tensor VAE ablal in table below. It is observed that due to local-only information aggregation as result of removing all virtual bonds (and related atom features), the performance is worse than the complete TensorVAB version.
What happens if a N x 1 kernel used. The final ablation study concerns with using a N x kernel with a smaller "field of view" as compared to that of N x 3 kernel. Its performance on Drugs dataset is shown as TensorVAE abla2 in Tab It performs slightly better than the ablation removing all virtual bonds. The reason is that though its field of view is smaller, it still achieves a global information aggregation for the focal atom. Nevertheless, it underperforms the complete TensorVAB version due to smaller "field of view" for information aggregation.
Table 5: Performance comparison among models with different input feature engineering setupon GEOM Drugs dataset
COV
MAT Median
Method
Mean
Median
Mean
NNiv
52.14 + 1.48
51.69 + 1.17
1.4322 + 0.0247
1.3861 + 0.0173
Tenso ablal
90.72 + 1.54
99.53 + 0.64
0.8748 + 0.0161
0.8619 + 0.0214
Tenso abla2
91.04 + 1.21
99.74 0.42
0.8706 + 0.0131
0.8561 + 0.0204
TensorVAE
93.34 + 0.35
99.90 0.31
0.8074 + 0.0135
0.7927 + 0.0186
*The standard deviations for all ablation studies are obtained by testing on 2000 testing molecules.
18

---

Under review as a conference paper at ICLR 2023
A.9 SAMPLES OF GENERATED CONFORMATIONS
Ground Truth
Generated
Ground Truth
Generated
Ground Truth
Generated
Ground Truth
Generated
Figure 6: Generated samples by the TensorVAE
19


Reviews:
[{"review_id": "X5-SMtbbb8", "reviewer": "Reviewer_4X5j", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper illustrates a method for how to train a neural network to predict 3D conformers of small molecules starting from their 2D structures. Building on previous works such as CVGAE, the authors propose a few tweaks which are meant to improve the models’ performance: 1) use a single tensor to represent both atom and bond information, and featurize the single tensor into atom-level vectors using a 1D convolution operation on the tensor; and 2) use a transformer architecture and combine 2D and 3D representations during training using a variant of cross attention. The authors demonstrate that with their method, they can achieve performance comparable to other similar SOTA methods.", "summary_of_the_review": "The paper (other than the grammatical issues and typos) is well written, and flows nicely. It is informative and the figures are generally very enlightening. The benchmarks are decently well done, as well. However, the paper demonstrates limited novelty in the technical space, as the proposed tweaks to existing work are either questionable or small. The result of these tweaks also do not seem to lend to significantly improved performance beyond the existing methods.", "strength_and_weaknesses": "The paper presents the authors’ main ideas in a largely clear manner, and it introduces and mostly justifies the choices made. The implementation details are fairly clear. The authors also present an interesting analysis on the loss curves which show a tradeoff between the reconstruction and variational losses, which is very appreciated. The authors are relatively thorough in their benchmarks with other works.\n\nThere were, however, some areas where the paper makes some design decisions which might not be optimal:\n### Convolving over the input tensor with both atom and bond features can be problematic\nThe input tensor that is constructed is very heterogeneous, because it contains both atom feature vectors and bond feature vectors (side note: presumably, the atom feature vector is padded with 0s in order to make the two vectors equal size). Because the tensor is structured such that the atom features are along the diagonal, this means that a convolution operation over the tensor is not uniform. That is, consider an $N \\times 3 \\times 50$ convolutional kernel. The entry at $(0, 1, 0)$ will be multiplied against an atom feature in the first sliding window, and then it will be multiplied against a bond feature (which is twice the magnitude of the atom feature) in the next sliding window. This leads to the same entries in the kernel to experience entries in the tensor that are of varying magnitudes and meanings, which is suboptimal.\n### Convolution is not particularly natural for this kind of input\nFirstly, the input tensor is symmetric about the diagonal so the convolution sees redundant information. Combined with the issue brought up above, the redundancy structure is different for every sliding window (i.e. sometimes entries near the top and bottom are redundant, and other times entries near the middle are redundant).\n\nSecondly, because the molecules are variable sizes, the input tensor will need to be zero padded to a maximal size. This means that the same kernel will also experience zeros when there is padding. This can also lead to confusion because zero is also a meaningful value in the input tensor (which is mostly one-hot encoded)\n### Kabsch alignment shouldn’t be necessary\nThe use of Kabsch alignment (or alignment in general) should not be necessary. Transformers retain the ordering of the input vectors/tokens. That is, the $i$th input token is mapped to the $i$th output token through the attention mechanism. This means that the use of alignment should not be needed, as each atom/token can be mapped to an output token.\n### Variation of cross attention is used without much justification\nThis is more of a minor point, but cross attention typically combines the key and value of one sequence with the query of another sequence. This work computes the key and query of the 2D tokens with the value of the 3D tokens. This is somewhat distinct from the traditional cross-attention mechanism. It would be good to have a justification for why this might work better than traditional cross attention.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "WSpHd-Ah8sP", "reviewer": "Reviewer_n6uh", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose TensorVAE, a relatively simple model for generating 3D conformations from 2D molecular graphs. TensorVAE employs 1) a unique feature engineering step that represents each molecule as a tensor (with or without 3D coordinates and distances), and 2) a VAE with two transformer encoders, one that encodes the graph (using the tensor input without 3D information) and the approximate posterior that produces latents from the 3D information, and a transformer decoder for the likelihood, where keys and queries come from the 2D graph representation, and values are the latents from the 3D posterior encoder. The loss is a standard roto-translation invariant loss.  The authors show that, using standard transformer architectures and training procedures, TensorVAE performs comparably to the best current methods at conformation generation using the GEOM dataset, and argue that TensorVAE is much simpler than the comparable models due to superior feature engineering.", "summary_of_the_review": "The paper extends existing methods and molecular featurizations to achieve near state-of-the-art results for conformation generation. The paper spends a lot of time arguing for the superiority of the feature engineering used, but it is not at all clear that the results depend on that featurization. The authors could greatly improve this paper by demonstrating the value of the feature engineering beyond performance on benchmarks through ablation studies (e.g. removing aspects of the features) and studies of other featurizations that are not conflated with the tensor design (e.g., radius-1 atom environments).", "strength_and_weaknesses": "Strengths:\n- Simple architecture and featurization\n- Original use of two encoders for 2D and 3D features in the VAE formulation that allows for 3D conformation generation from 2D graphs\n- (Near) state-of-the-art results for conformation generation (and QM9 property prediction)\n\nWeaknesses:\n- It is not clear that the proposed method of producing atom-tokens via 1D convolution on the input tensor is necessary. Any number of possible aggregation steps could have been used, many of which would likely have yielded similar results. One could also imagine performing a Tucker decomposition and using singular values as tokens, etc.\n- Moreover, it's not even clear if the tensor formulation is needed. The 1D convolution aggregates information about the radius-1 atomic environment (including virtual bonds). One could use any radius-1 atomic-environment hash as tokens with the proposed featurization, which would likely yield comparable results.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Na9JUTyDxQ", "reviewer": "Reviewer_8zkT", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a CNN/transformer based variational autoencoder for generation of molecular 3d structure from molecular graphs.\n", "summary_of_the_review": "While I find the subject of the paper very interesting, my primary reason for recommending rejection is the lack of novelty and the limited empirical evaluation.\n", "strength_and_weaknesses": "Strengths:\nThe paper deals with an important problem and presents a working solution that is fairly simple and effective.\n\nWeaknesses:\nThe novelty is limited - the main technical contribution appears to be a model that works well, built using standard CNN, transformer, VAE components.\n\nIt is unclear if the comparison with other methods is fair - results are taken directly from previous papers but is it the same test set that is used? According to the text, the authors randomly select 200 molecules. \n\nConfidence intervals are not provided for the results. With such a small test set, significance is difficult to determine.\n\nBecause a CNN is used, it appears that the method is not invariant to permutation of the atoms. It is not clear if this has negative implications for the performance.\n\nIt is not clear to me what the motivation is behind combining node and edge features in a tensor. This restricts the dimensions of the node and edge features to be identical, but on the other hand allows for standard CNN code to run on the representation. Is that the reason for this choice?\n\nThe relation of the proposed model to common GNN architectures is not fully described. If I understand correctly, a kernel of size 3xN is used? If it had been a 1xN kernel, I believe it would correspond exactly to a fully connected GNN which is commonly used. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-pj2VGf7sg", "reviewer": "Reviewer_iXjq", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper tackles the generation of 3D conformations of a molecule from its 2D graph. They propose to encode a molecular graph using\na fully connected and symmetric tensor. They use the standard VAE framework, where they build two input tensors with one encoding only the 2D molecular graph and the other also encoding 3D coordinates and distance. Both tensors go through the same feature engineering step and the generated feature vectors are fed through two separate transformer encoders. The output of these two encoders is then combined in an intuitive way to form the input for another transformer encoder for generating confirmation directly.", "summary_of_the_review": "\n\nSome detailed comments: \n\n1. I wonder why the model needs two conditional encoders, one conditioned on the graph only and another one conditioned on both the graph and coordinate, why not directly use the latter since it has also graph information? \n\n2. I am a bit confused about section 2.3 where the second main idea is explained. ", "strength_and_weaknesses": "Strength:\nthe idea is to combine Atom and edge features into a single input by adding an additional dimension to\nthe adjacency matrix, making it a tensor where the diagonal section of the tensor holds the atom features sounds interesting.\n\nWeakness: the models seem to be really big (because of the input tensor)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject