TensorVAE: a simple and efficient generative model for conditional molecular conformation generation
Authors: Hongyang Yu
Hongjiang Yu
Venue: Accepted by TMLR ()

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 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 Autoencoder framework for generating conformations 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 outperforming more than a dozen state-of-the-art models employing more sophisticated and specialized generative architecture.

Published in Transactions on Machine Learning Research (01/2024)
TensorVAE: a simple and efficient generative model for conditional molecular conformation generation
Hongyang Yu
hyyu@anniionnerbioo.com
Anticancer Bioscience (ACB), Ltd., Sydney, Australia
Hongjiang Yu
hyyu Shayuuaaanticcccebbo.ooom com
Anticancer Bioscience (ACB), Ltd., Sydney Australia
Reviewed on OpenReview: https: openreview. net/ forum? id=rQgat4gYce
Abstract
Efficient generation of 3D conformations of molecule from its 2D graph is a key challenge in in-silico drug discovery. Deep learning (DL) based generative modelling has recently be- come a potent tool to tackling this challenge. However, many existing DL-based methods are either indirect-leveraging inter-atomic distances or direct-but requiring numerous sam- 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 novel encoding and feature extraction mechanism relying solcly standard convolution operalion to generale token-like feature vector for each atom. These feature vectors are then transformed through standard transformer encoders under a conditional Variational Autocncodor framework for generating conformations directly. We show through experiments on two benchmark datasets that with intuilive feature engineering, a relatively simple and standard model can provide promising generative capability outperforming more than dozen state-of-the-ar models employing more sophisticated and specialized genera- tive architecture. Code is available at https://gitte com/yuh8//ensor//a
pling steps to generate conformations.
1 Introduction
Recent advance in deep learning has enabled signíficant progress computational drug design (Chen el al. 2018). Particularly, capable graph-based based generative models have been proposed generate valid 2D graph reprosentation of novel drug-like molecules (Honda et al., 2019 Mahmood et al. 2021 Yu & Yul 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 et al. 2021 Simm et al. 2021 Gebauer et al. 2022). A stable 3D structure or conformation of a molecule is specified by the 3D Cartesian coordinates of all its atoms. Traditional molecular dynamics or statistical mechanic driven Monte Carlo methods are computationally expensive, making them unviable for generating 3d molecular structures at scale (Hawkins) 2017). In this regard, deep Icarning(DI))-baeed gencrative methods have become an attractive alternative.
DL-based generative methods may be broadly classified into categories: listance-basedd reconstruction- based, sequential and energy- based and direct methods. The main goal of distance-based methods is learning probability distribution over the inter-atomic distances. During inference, distance matrices are sampled from the learned distribution and converted to valid 3D conformations through post-processing algorithms.
Two representative methods of this category include GraphDG (Simm Hernández- Lobato (020) and CGCF Xu al. 2021a). An advantage of modeling distance roto-translation invariance property- an important inductive bias for molecular geometry modeling (Köhler et al.] (2020). Additional virtual edges and their distances between 2nd and 3rd neighbors often introduced to constrain bond angles

---

Published in Transactions on Machine Learning Research (01/2024)
and dihedral angles crucial to generating valid conformation. However, Luo et al. (2021) 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 region. Another weakness of the distanco-based based methods is the error accumulation problem; random noise in the predicted distance can be exacerbated by an Euclidean Distance Geometry algorithm, leading to generation of inaccurate conformations (Xu et al. 2022: 2021b).
To address the above weaknesses, reconstruction-based methods directly model distribution over 3D co- ordinates. Their main idea is to reconstruct valid conformations from distorted coordinates. GeoDiff (Xul ct al. 2022) and Uni-Mol (Zhou et al. 2023) are pioncering studics in this respect. Though sharing sim- ilar idea, they differ in the process of transforming corrupted coordinates to stable conformations. While GeoDiff adapts a reverse diffusion process (Sohl-Dickstein et al.) 2013), Uni-Mol treats conformation recon- struction as an optimization problem. Despite their promising performance, both methods require designing of task-specific and complex coordinate transformation methods. This is to ensure the transformation is rolo-translation or -equivaanar To achieve this, GeoDiff proposed specialized SE(())-quivariint Markov transition kernel. On the other hand, Uni-Mol accomplished the same by combining task-specific adaption of transformer (Vaswani et al. 2017) inspired by the AlphaFold's Evoformer (Jumper et al. 2021 with another specialized equivariant prediction head (Satorras et al. 2021.. Furthermore, GeoDiff requires numerous diffusing steps to attain satisfactory generative performance which can be time consuming.
While promising generative performance has been achieved by directly learning a distribution over the 3D geometries (coordinates or pair-wise distances) of molecules, energy-based learning methods have also recently been shown to yield competitive performance in molecular conformation generation. A unique advantage of using energy minimization as the reward mechanism is that energy-based models can better explore the low-energy regions of the conformational space of a molecule, leading to generating conformations with both high quality and diversity. On the other hand, methods relying directly on minimizing distance metrics to ground-truth conformations may result in generating very similar conformations with high enorgy strain. Two recent. methods adopting the energy minimization paradigm are TorsionNel. (Gogineni et al., 2020) and GFlowNet (Volokhova et al., (023) for conformation generation. Both of which are sequential and energy-based methods that sequentially move an original molecular conformation by modifying the torsion angle of all rotatable bonds towards a lower energy state. Despite their promising potential and advantages, sequential methods are relatively inefficient compared to the direct methods. For instance, GFlowNet requires iterating through 40,000 training steps per molecule to achieve satisfactory performance.
CVGAE Mansimov et al. 2019 and DMCG [Zhu et al. 2022 have attempted to resolve the generative effi- ciency issue by developing models that can produce a valid conformation directly from a 2D molecular graph in single sampling step. Regrettably, the performance of CVGAE is significantly worse than its distance- based counterparts mainly due to the use of inferior graph neural network for information aggregation (Zhu) 2022). DMCG aimed to improve the performance of its predecessor by using more sophisticated graph
et al. neural network and a loss function invariant to symmetric permutation of molecular substructures. Although DMCG achieved superior performance, acquiring such loss function requires enumerating all permutations of a 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 developing 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 instead focus on intuitive input feature engineering. We propose to encode molecular graph using a fully-connected and symmetric tensor. For preliminary information aggregation, we run a rectangle kernel filter through the tensor in a ID convolution manner. This operation has a profound implication; with a filter width of 3, the information from two immediate neighbors as well as 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 a standard transformer encoder for further information aggregation.
The generative framework follows the standard conditional variational autoencoder (CVAE) sctup. 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 engineering step and the generated
2

---

Published in Transactions on Machine Learning Research (01/2024)
feature vectors are fed through two separate transformer encoders. The output of these two encoders are then combined in an intuitive way to form the input for another transformer encoder for generating conformation directly. The complete generative model is abbreviated as Tensor VAE.
In summary, the proposed method has 4 main advantages. (1) Direct and Efficient, generating confor- mation direcity from a 2D molecular graph in a single step. (2) Simple, not requiring task-sepecific design of neural network architecture, relying only on simple convolution and off-the-shelf transformer architecture; (3) Easy to implement, no custom module required as both PyTorch and TensorFlow offer ready-to-use convolution and transformer implementation. These advantages translate directly to excellent practicality of the TensorVAE method. (4) Achieving competitive performance through simplicity, we demon- strate through extensive experiments on two benchmark datasets that the proposed Tensor VAE,, despite its simplicity, can perform competitively against 22 recent state-of-the-art methods for conformation generation and molecular property prediction.
2 Method
2.1 Preliminaries
Problem Definition. We formulate molecular conformation generation as conditional generation task. Given set of molecular graphs G and their corresponding iid conformations R, the goal is to train a generative model that approximates the Boltzman distribution, and from which a valid conformation conditioned on a molecular graph can be easily sampled in a single step.
Story Line. In the ensuing sections, we breakdown the formulation of the proposod method in three novel ideas. We first introduce how molecular graph can be encoded using 3D tensor. Then, we demonstrate how token-like feature vector can be generated from the input tensor by using ID convolution operation. The generated feature tokens resemble those used in the language modelling, thereby allowing the use of standard transformer encoders for effective information aggregation. Finally, we propose novel mechanism to combine the outputs of the transformer encoders under a conditional-VAB framework to arrive at the final generative model.
2.2 Input tensor graph
Message passing graph neural network (GNN) is popular feature extraction backbone for DL-based moloc- ular conformation generation. The input for this backbone is often composed of three components, including atom features, edge features and an adjacency matrix. Atom and edge features normally pass through sepa- rated embedding steps before being fod to the GNN. Adjacency matrix is then used to determine neighboring atoms for layer-wise information aggregation. Although bond features are aggregated onto alorn features and vice versa, these two features are maintained separately throughout the message passing layers (Gilmer) let al.] 2017] Satorras et al.] 2002).. Instead of having separated inputs, our first simple idea to combine them into single input. Specifically, 3D tensor. Each cell on-diagonal of the tensor holds the focal atom feature vector to which information from nearby connected atoms are aggregated.
add additional dimension to the adjacency matrix, making it
We consider three types of atom features comprising atom type, charge and chirality. Each feature one-hot encoded and they are stacked together to form a single atom feature vector. There are two variants of the atom feature vector corresponding to two input tensors for the two encoders of the CVAE framework: an encoder conditioned only on graph (referred to as the G tensor) and the other conditioned on both graph and coordinates (referred to as the GDR tensor). For the GDR tenosr, every focal atom feature vector has three additional channels incorporating the 3D coordinate of the respective atom, and distance channel filled with zeros.
Each off-diagonal cell holds the stacked neighbour atom and bond features. The considered bond features are bond type, bond sterco-chemistry type, ring size and normalized bond length. A virtual bond is also included in the bond type. It is worth noting that all virtual bonds share the same virtual bond type; they only differ in their normalized bond length. The normalized bond length is caleulated edge length (1 for
3

---

Published in Transactions on Machine Learning Research (01/2024)
type encoding
Carbon
1
C
D
charge encoding
Single bond
d
C
Atom chirality
C
encoding
0,1
0.7
coord
Bond stereachem
0
0
L/5
Atom feature vectors summed Coordinate channels zeroed Stacked with bond feature vector
Bond type
In ring size
Normelized Bond length
Euclidean distance
Figure 1: Benzene ring tensor graph example. Note that the values in the feature vector and its dimension are for demonstration purpose only. We explain how they are determined in See
direct neighbor, 2 for 2nd neighbor, etc.) divided by the longest chain length. To construct off-diagonal feature vector, we first sum the atom feature vectors of the connected atoms. This vector is then stacked with one-hot encoded bond type vector, normalized bond length, and one-hot encoded ring size vector to become the off-diagonal feature vector. Since there are no bond features for a focal atom, the bond feature vector channels on-digonal are also filled with Os. Therefore, both on and off-diagonal feature vectors have the same dimension.
There also two variants of the off-diagonal feature vector. For the G tensor, coordinate and distance channles are excluded. For the GDR lensor, to match the size of the on-diagonal feature vector, every off- diagonal feature vector has three more coordinate channels filled with Os, and an additional distance channel holding the Euclidean distance between two connected atoms. This off-diagonal fcature vector obtained for all atom pairs, making the proposed tensor fully-connected and symmetrie.
A tensor encoding 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 the tensor, and train it to predict distribution over the inter-atomic distances. We utilize a standard UNet (Ronneberger et al. 2015) structure to map the input tensor to probability distribution over a 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 2020
We refer to this
model as the NaiveUNet. More details of the NaiveUNet can be found Sec A.4
This naive model achieves unsatisfactory performance as shown in Tab and Tab 9
merely outperforming
GraghDG and is far from that of the state-of-the-art. There are two major issues to this approach. First, with small kernel size (3 3 used in the UNet), it takes many convolution layers to achieve information aggregation between atoms that are far apart; it does not take full advantage of high-order bonds (chemical or val) already made available in the input lensor. Secondly, the output size grows quadratically with the number of atoms, as compared to only linear growth in the reconstruction based or direct generation mothods. The solution to the first issue is rather simple, obtained by increasing the kernel size to expand its "field of view". On the other hand, solving the second issue requires elevating the naive two-step generative model to direct one.
4

---

Published in Transactions on Machine Learning Research (01/2024)
2.3 Extended kernel and Attention Mechanism
We observe that every row column of the proposed tensor contains global information encompassing 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 a profound implication; global information from the immediate neighbors, 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 1D convolution is permitted. With multiple kernels being applied simultaneously, each stride of these kernels generates feature vector for a single atom. An illustration of the 1D convolution operation is shown in Fig
Token
Token
Tokon N
V
V
V
C
C
C
V
Extend
ID COIV
C
C
Figure 2: Extending kernel and 1D convolution.
We further observe that the generated fcature vectors resemble the token-like feature vectors used in language modeling. This observation combined with the proven success of attention mechanism in other related work leads to the selection of transformer architecture as the backbone of our generative model. A significant advantage of using transformer's self-attention mechanism is, similar to the extended kernel, enables a 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 transformation. We present further insight and a more detailed analysis of the adavantage of this input feature engineering in See A.1 There is also intoresting equivalence between the information aggregation achieved by fully-comnected MPNN (Gilmer et al. 2017 and running a 1 x 1 convolution operation over the proposed input tensor, as detailed in Sec A.2
2.4 Putting everything together
Conditional variational autoencoder framework. We aim at obtaining a generative model pa (R(GG) that approximates the Boltzmann distribution through Maximum Likelihood Estimation. Particularly, given
set of molecular graphs G and their respective ground-truth conformations R, we wish to maximize the following objective.
log (R()) = log / p (2)) po G)
(1)
A molccular graph can have many random conformations. We assume this randomness is driven by a latent random variable z ~ (=), where (z) a known distribution e.g. a standard normal distribution. As pa (R]z, G) is often modeled by complex function e.g. a deep neural network, evaluation of the integral
in Eq 1 is intractable. Instead, resort to the same techniques proposed in the original VAE (Kingma &
Welling 2014) to establish tractable lower bound for Eq
log Po (R|G)= 2 (.. [log Pa (R) 2, G)] DRL. (q
(2)
where DKL is the Kullback-Leibn divergence and qu (2 RR G) is variational approximation of the true posterior p (2|R, G). We assume p (z) N (0, I) and gu (z) R, G) is diagonal Gaussian distribution whose
5

---

Published in Transactions on Machine Learning Research (01/2024)
means and standard deviations are modeled by transformer encoder. The input of this transformer en- coder is the proposed tensor containing both the coordinate and distance information. We denote this tensor the GDR tensor. On the other hand, pe (Alz, G) is further decomposed into two parts: a decoder Po, [R]2, 0a, (G)) for predicting conformation directly and another encoder σg, (G) for encoding the 2D molecular graph. The input tensor for σœ, (G) is absent of coordinate and distance informat ion, and is there- fore denoted the G tensor. Both encoders share the same standard transformer encoder structure. However, there minor modification to the transformer structure for the decoder. Specifically, the Query, Key matrices for the first multi-head attention layer are computed based the output vectors of σg, (G), and the Value matrices come directly from the reparameterization of the output of qu (z) R, G),, as 2 = Mhe + Eme, where Mw and E are the predicted mean and standard deviation respectively. e is sampled from N (0, I). We present the complete picture of how the two encoders and the decoder are arranged in a CVAE framework in Fig 3
Pa. (R|z. oe (G)
KaRtch
tranofornet Block
Reconstruction loss
Add horm
Feed Famward
KL Regularization loss
Linear
Concat
Scaled Dot-praducct
Attention
Feed Fonword
A
Split
taxclrree enooder
Norm
R
Fg
Az
Lineer
Lineor
(G)
qw (z| G)
RC
Figure 3: Variational AutoEncoder framework (left) and modified multi-head attention (right)
Intuition behind the modified attention. There are multiple ways to join together the output of the two encoders to form the imput to the final decoder. Popular methods include stacking or addition. We tried both these mothods with unsatisfactory performance. We notice that, due direct stacking or addition of the sampled output of the onto the output of on, atlention 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 { hly, } of oe, and attribute the variation in conformation to the Value matrices which are directly sampled from {=1,-.,,N) EN) te- 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.

---

Published in Transactions on Machine Learning Research (01/2024)
Learning to achieve approximate Roto-translation invariant loss. Following ConfVAE [u et al, 2021b), we formulate the reconstruction loss as.
Rij A (f.. (R,
log GR2z,
==1
where A (-) is a function aligning the predicted conformation R onto the reference conformation R. We choose Kabsch algorithm (Arun et al. 1987) the alignment method which translates and rotates the predicted conformation onto its corresponding ground-truth before loss computation. This makes the reconstruction loss roto-translation invariant.
Simultancously, minimizing the KT-loss component DKL m (z|R, G) [=)] compels the output of the poste- rior encoder to adhere to a standard normal distribution. Despite this minimization promoting convergence, achieving exact equality 9n (z)R, G) (2) in practice is challenging, especially in the presence of SE(3) transformations of the input R. Consequently, upon convergence, the objective function defined in Eq only achieves approximate roto-translation invariance.
Direct conformation generation at inference time. To generate single conformation, we first con- struct the G tensor of molecular graph and obtain single latent sample {ıI....Zy from a standard diagonal Gaussian distribution. The G tensor is passed through 0ø, encoder to produce { (......%) 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 L standard attention layers be transformed to the final conformation. The entire generation process depends only on 2D molecular graph, and requires single sampling step and single pass of the TensorVAB model.
3 Experiment
In this section, we first elaborate on the implementation details of the TensorVAB model including determin- ing the size the input tensors, network architecture and how the entire framework is trained end- to-end. We then present conformation generation experiment results of the proposod Tensor on throe benchmark data-sels, including GEOM-QM9, GEOM-Drugs and Platinum data-sets. While the GEOM dalasets contain unbound conformations of molecules, the Planinum dataset contains molecular conformations bound to their respective protein targets. The generative performance of the proposed model is compared to those of 15 the-r-rtttt--ttttt baselines. In addition to conformation generation, in Sec.A.8] we further demonstrate the effectiveness of information aggregation the proposed TensorVAF. architecture by briefly comparing the molecular property prediction performance of the proposed method against more state-of-the-an baselines on the MolecularNet (Wu 2018 benchmark.
3.1 Experiment setup
Dataset. Following existing work (Luo et al. 2021 Shi et al. 2021 Xu et al. 2021 a 2022
Zhou et al.
2023), we utilize the GEOM data-sel for eval evaluating uating the performance the proposed Tensor AE. GEOM contains 37 million energy and statistical weight annotated molecular conformations corresponding to 450,000 molecules (Axelrod
KGmeez-Bombare 2022). This dataset is further divided into two constituent datasets, Drugs and QM,. The Drugs dataset covers 317. 7,0000 median-sizee molecules averaging 44.4 number of atoms. The QM9 dataset contains 133,000 smaller molecules averaging only 18 atoms.
We follow Xu et al. (0022) to randomly select 40,000 molecules from each dataset to form the training set. For each molecule, we choose the top 5 most likely conformations. This results in 200,000 training conformations for each train set. For validation set, we randomly sample 2,500 conformations for both Drugs and QM9 experiments. Finally, for testing, following (Shi et al. 2021 Xu et al. 2022), we randomly select
Ranked by their Boltzmann weight.

---

Published in Transactions on Machine Learning Research (01/2024)
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 Drug.
The GEOM dataset contains conformations of molecules that are not bound to any specific target. To assess the proposed model's ability to generate
bound conformations, we additionally evaluate its performance using the Platinum dataset (Friedrich et al. 2017]. The Platinum dataset derived from the Pretein Data Bank (Berman et al., 2000) and consists of two high-quality ligand-protein bound conforma- tion dataset: comprehensive dataset and diversified subset of 4,626 and 2,912 structures, respectively. Following the setup in (Friedrich et al.) 2017), we test the performance of the proposed Tensor on the diversified subset.
Determining input tensor size and atom ordering. We conduct a basic data analysis on the entire Drugs dataset to determine the 98.5th percentile of the number of atoms to 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 lensor to 69 69 for Drugs experiment. On the other hand, we use the maximum number of atoms 30 for QM9 experiment. The channel features for the imput tensor include atom types, atom charge, atom chirality, bond type, bond stereo-chemistry and bond in-ring size. For the GDR tensor, we also include 3D coordinate channels and a distance channel. The resulting channel depth is 50 for GDR tensor and 46 for G tensor. The detailed information of these features and their encoding method is listed in Sec [A.5] The ordering of the atoms along the diagonal of the tensor is determined by a random Depth-First Traversal (DFT) of the molecular graph.
Implementation details. We implement the proposed Tensor using Tensorflow 2.5.0. All three trans- former encoders of Tensor follow the standard Tensorflow implementation in htttps://www/ tensorflow All of them have 4 layers, 8 heads and latent dimension of 256. Both QM9 and Drugs experiments share the same network architecture and hyper-paramete configuration. We present the detailed training hyperparameter configuration in Sec A.3
htt///eet//utooral//raaaat
Evaluation metrics. We adopt the widely accepted coverage score (COV) and matching score (MAT) Shi
et al. 2021 to evaluate the performance of the proposed TensorVAE model. These two scores are computed
as:
{ FREEC,RMMD RR. R) S co
COV (Cg, |,,1
(4)
MAT (C,: C.)
min RMSD (R, ft)
REC,
where Cg is the set of generated conformations and C, is the corresponding reference set. The size of Cg is twice of that of C,,, as for every molecule, we follow (Xu et al., 2022) to generate twice the number of conformations as that of reference conformations. 8 is predefined threshold and is set to 0.5À for QM9 and 1.25A for Drugs respectively (Shi el al. 2021) RMSD stands for the root-mean-square deviation wet R and R, and is computed using the GetBestRMS method in the RDKit (Riniker & Landrum 2015) package, While COV score measures the ability of model in generating diverse conformations to 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.
To evaluate the accuracy of the proposed model on the Platinum dataset, we employ two metrics: the sutte deviation (RMSD) for four ensemble sizes (10, 50, 250, and 500) and the percentage of molecules with RMSD within specified thresholds (0.5, 1.0, 1.5, and 2) for two ensemble sizes (50 and 250). In terms of generative speed evaluation, we calculate and compare the mean and median generation times for the four ensemble sizes across all 2,912 molecules.
Basclines. We first compare the generative performance of the proposed TensorVAE model to those of classical RDKit method; 5 distance- based methods including GraphDG, CGCF, ConfVAE, ConfGF and DGSM; 2 reconstruction--oesee methods including GeoDiff and Uni-Mol; 3 direct methods including CVGAE,
3This limit on the number of conformations for testing molecules is taken directly from https://github. com/ Deep@rashllarrminn/Coorttt which also followed by other compared methods the GEOM experiment.
8

---

Published in Transactions on Machine Learning Research (01/2024)
GeoMol, and DMCG. For the Platinum dataset, we also incorporate 4 classical methods, namely Ballon DG and Ballon GA (Vainio & Johnson 2007), MultiConf-Dock (Sauton et al. 2008) and ETKDG (Riniker & Landrum 2015
We then compare the molecular property prediction performance of the proposed model (specifically the GDR encoder) to 7 more strong baselines comprising D-MPNN (Vang et al.] 2019), AttentiveFP (Xiong
et al. 2019), N-Gram (Liu et al. 2019), PretrainingGNN (Hu et al. 2020), GROVER (Rong al. 2020
GEM Fang et al. 2022 and finally again Uni-Mol.
3.2 Results and Discussion
Unbound conformation generation. The COV and MAT scores for all compared methods on both QM9 and Drugs datasets are presented in Tab The proposed Tensor achieves the statc-of-the-art generative performance. Additionally, we have conducted 4 ablation studies on the input feature engineering met in 3.3 to demonstrate why ID convolution with a N x 3 kernel is crucial to achieving a good generative performance. While none of the cited baselines quantify confidence of their results, we have included standard deviations in all our results.
In Tab TensorVAEREF results are obtained by running test on the same set of test data that adopted by all other baselines. While Tensor employs the same test set for evaluation, it's essential to note that the training dataset differs from that in ConfGF. Nevertheless, to confirm that none of the molecules in the ConfGF test set are included in our training dataset.
have conducted a thorough examination
Additionally, we observed that the ConfGF DRUGS test set has a maximum of 71 heavy atoms per molecule, exceeding our predetermined maximum of 69 atoms by 2. While there are only molecules with more than 69 heavy atoms, we do not anticipate a significant performance change by allowing Tensor to handle an additional atoms. Therefore, we opt not retrain TensorVAB for this test set. To ensure a fair comparison, for these 3 molecules, we assume worst-case scenario where the trained TensorVARe can only achieve MAT score of 2Å and COV score of 0%. The TensorVAERER Mean/Mediag MAT and COV scores for the DRUGS datasct are computed under this worst-case scenario. For the QM9 dataset, as we have already used the maximum number of heavy atoms, TensorVAE results are obtained as usual.
On the other hand, Tensor results have been obtained on a set of random testset, selected based on the same filtering condition proposed in ConIGF and having a maximum number of heavy atoms of 69 per molecule. This set of 200 molecules contains 23,079 and 14,396 testing conformation for QM9 and Drugs, rospectively. TensorVAE! results and standard deviations obtained by running 10 experiements cach with a different random seed on the same 200 testing molecules.
Tensor results are obtained by running 10 experiements each with different random seed as well as a different set of 200 testing molccules. In this setting, both testsets contain 2,000 testing molccules, amounting to more than 280k and 140k testing conformations for QM9 and Drugs, respectively. The number of testing conformations is more than 70% of that of training conformations. Attaining consistent performance on this much larger testset consolidates the generalization capability of the proposed TensorVAB. and verifies its robustness under random permutation of atom ordering. Additionally, as noted by Xu et. al.
2022 Eqs 4 and 5 are only the recull scores. We also present the precision scores results
in Tab 11 of Sec A.6 where Tensor VAE again achieves the state-of-the-art performance with a considerable margin.
Xu et al. 2021bb discovered that the quality of conformations generated by deep generative models can be further refined by an additional empirical force field (FF) (Halgren) 1996) optimization procedure. Uni- Mol also leverages FF optimization to improve its generative performance. Different from GeoDiff which reconstructs valid conformation directly from random noisy coordinates, Uni-Mol simply refines an initial conformation optimized by RDKit (using ETKGD with FF RRiniker & Landrum 2015
"This datasct is available for download bttps://github www//Deerraphleearinnggg///eo
Originally generated in ConfGI,
it serves the common datasct across all compared baselincs. For the 200 testing molccules, the total mumbers of annotated conformations are 22,408 and 14,324 for QM9 and Drugs respectively. As all compared baselines utilize the same test set, including standard deviation as an uncertainty measure unnecessary.

---

Published in Transactions on Machine Learning Research (01/2024)
Table 1: Performance comparison between Tensor VAE and 10 baselines on GEOM dataset.
QM9
Drugs
Models
COV
(%) 1
MAT
(Å) +
COv
(%)
MAT (A) 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
ConfGI
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
Tensor VAE
97.79
100
0.1985
0.1951
93.05
98.98
0.8087
0.7866
98.11
100
0.1970
0.1926
94.91
100
0.7789
0.7585
VAF!
+0.25
+0
#0.0016
#0.0027
0.335
+0
#0.0027
+0.076
97.11
100
0.2041
0.1920
93.34
99.90
0.8074
0.7927
Tensor
#0.31
+
+
#0.007
#1.17
+0.31
+0.0135
+0.018
*Bold font indicates best result. Results for RdKit, CVGAE, GraphDG, CGCF, ConfGF are taken from
Shi et al. 2021); all other results are taken from (Zhou et al, 2023
Values following tre standard
deviations.
For a fair comparison, we exclude deep generative models relying on FF optimization from Tab compare their performances separately in Tab Again, the proposed TensorVAF. with FF optimization outperforms all of them with significant margin.
Table 2: Performance comparison between methods with FF optimization on GEOM Drugs datasel
COV(%) f
MAT(Á) 4
Method
Mean
Median
Mean
Median
CVGAE
83.08
95.21
0.9829
0.9177
GraphDG
84.68
93.94
0.9129
0.9090
Uni-Mol
91.91
100
0.7863
0.7794
CGCF
92.28
98.15
0.7740
0.7338
ConfVAE
91.88
100
0.7634
0.7312
GeoDiff
92.27
100
0.7618
0.7340
Tensor TensorVAERES RRF
93.36
98.18
0.7267
0.7032
94.74
100
0.6985
0.6845
Tensor V
#0.66
+00
+0.012
+..019
"Results for CVGAE, GraghDG, CGCF, and ConfVAE are taken from (Xu et al. 2021b GeoDiff and Uni-Mol results are from their source paper.
In terms of simplicity, the proposed TensorVAB uses standard transformer encoder and a simple Kabsch alignment loss. On the other hand, due to the lack of effective input feature engineering, both DMCG and Uni-Mol require design of sophisticated network architectures and complex loss functions achieve a good generative performance. A direct consequence these complicated designs is large number of model parameters, shown in Tab
10

---

Published in Transactions on Machine Learning Research (01/2024)
Table 3: Comparison of number parameters among VensorVAB, DMCG and Uni-Mol.
Method
Number of parameters
DMCG
128M
Uni-Mol
47.81M
TensorVAL training
11.5M
TensorVAR: inference
6.65M
*Tensor VAE has less number of parameters dur- ing inference time as the GDR encoder is not needed.
In terms of efficiency, TensorVAE is direct generative model capable of producing conformation from 2D molecular graph in single step. It takes only 62 seconds using single Xeon 8163 CPU to decode 200 QM9 molecules, and 128 seconds for 200 Drug molecules. In comparison, GeoDiff requires 5, 000 diffusion steps per conformation, and takes around 8, 500 seconds for decoding 200 QM9 molecules and 500 seconds for decoding 200 Drugs molecules on single Tesla V100 GPU. The proposed Tensor VAE achieves more than 100x speed up. Finally, some samples of the TensorVAE. generated conformations are shown in Fig
Ground Truth
Generated
Ground Truth
Generated
Ground Truth
Generated
Generated
Generated
Figure 4: Generated samples by the Tensor VAE
11

---

Published in Transactions on Machine Learning Research (01/2024)
Protein-ligand bound conformation generation. The performance evaluation of the proposed model generating ligand-bound conformation is vital to establish its potential application in high throughput virtual screening of drug candidates. Following the setup in Friedrich et al. (2017), we further compare the performance of the proposed TensorVAE model with 5 popular baselines on the Platinum dataset. We took the Tensor model trained on the GEOM-drugs conformations, and applied it directly to the Plat- inum diverse dataset for conformer ensemble generation. Before presenting the evaluation results on the Platinum dataset, we would like to first emphasize the difference between the Platinum dataset which was proposed in Friedrich et al. (017)) and the GEOM dataset which we have used to train TensorVAE. While the GFOM-drugs dataset mainly contains vacuum conformer-rotamee ensembles that are generated using semi-empirical density functional theory, the Platinum dataset only includes protein-bound ligand confor- mations. The energy states of conformers bound to a protein target are different from those of the stable unbound conformers. The underlying distributions governing the generation of these two datasets also differ significantly. Testing on the Platinum dataset without any retraining or finctuning creates a distribution shift from that of the GEOM training data. Inevitably, this will lead to performance degradation of the proposed VensorVAE: However, evaluating the proposed TensorVAB on the Platinum dataset remains valu- able for assessing its ability to generalize and accurately generate valid despite being trained solely on unbound conformations.
bound conformations,
We have repeated the experiments from Table 3 to Table in Friedrich et al. (017)) The results of these experiments are presented below.
Table 4: Arithmetic Mean and Median RMSD in Á Obtained for the Platinum Diverse Dataset.
50
250
500
Maximum ensemble size
Mean
Median
Mean
Median
Mean
Median
Mean
Median
Balloon DG
1.10
0.97
1.00
0.86
0.92
0.77
0.89
0.74
Balloon GA
1.22
1.10
0.90
0.80
0.72
0.63
0.67
0.58
RDKit
1.00
0.89
0.77
0.64
0.63
0.52
0.59
0.48
ETKDG
0.98
0.87
0.77
0.66
0.63
0.54
0.59
0.51
Multiconf DOCK
0.99
0.89
0.84
0.72
0.80
0.69
0.80
0.69
Tensor
1.02
0.95
0.85
0.77
0.73
0.67
0.69
0.63
Table 5: Fraction of Structures of the Platinum Diverse Dataset Successfully Reproduced within a Specified RMSD Threshold.
Maximum ensemble size
50
250
Minimum accuracy [A]
0.5
1.0
1.5
2.0
0.5
1.0
1.5
2.0
Balloon DG
0.29
0.57
0.77
0.92
0.33
0.62
0.81
0.92
Balloon GA
0.30
0.72
0.90
0.97
0.43
0.84
0.96
0.99
RDKit
0.39
0.71
0.89
0.96
0.48
0.82
0.95
0.98
ETKDG
0.36
0.72
0.91
0.97
0.45
0.83
0.95
0.99
Mlulicconf-DOCK
0.32
0.68
0.87
0.96
0.34
0.71
0.89
0.97
Tensor VAE
0.27
0.65
0.89
0.97
0.34
0.76
0.95
0.99
Although the proposed Tensor is trained solely on unbound conformations, it demonstrates comparable performance to 5 popular baselines in terms of accurately generating ligand-protein bound conformations (Tab and Tab which serves to validate its generalization capability. More specifically, it demonstrates a slight performance advantage over Balloon DG/GA and multiconf-Docke however, it falls short of matching the performance achieved by RDkit. and ETKDG. This result appears to contradict the findings obtained from the GEOM dataset, where the proposed TensorVAB outperformed RDkit.
The main reason of this contradiction could be attributed the distribution shift or dataset shift be- tween training and testing. Additionally, for constructing the training dataset, we sampled 40,000 molecules
12

---

Published in Transactions on Machine Learning Research (01/2024)
from GEOM drugs dataset and only retained the top-5 conformations with the highest Boltzmann weight for each molecule. These conditions further restrict the energy search space for conformation generation. Consequently, the Boltzmann distribution approximated (and learned) by the proposed YensorVAH might not be directly suited to prediction of ligand-bound conformations without further fine-tuning.
Table 6: Arithmetic Mean and Median Ensemble Sizes Measured for the Platinum Diverse Dataset.
10
50
250
500
Maximum ensemble size
Mean
Median
Mean
Median
Mean
Median
Mean
Median
Balloon DG
10
10
50
50
249
250
498
500
Balloon GA
9
10
49
50
244
250
487
500
RDKit
10
10
50
50
250
250
500
500
ETKDGG
10
10
50
50
250
250
500
500
Muullicon--OCCh
9
10
36
50
78
57
80
57
TensorVAR
10
10
50
50
250
250
500
500
In terms of generative capability () the proposed TensorVAE able to generate the complete 10-, 50-, 250-, and 500-conformers ensemble sizes for all molecules which puts head-to-head against RDKit and ETKDG. In terms of generative speed (Tab as the proposed TensorVAE only needs single pass of the neural network to generate conformations for each ensemble size, its mean and median runtimes (measured a single core of Xeon 8163 CPU) are significantly faster than the other compared methods.
Table 7: Arithmetic Mean and Median Runtimes in Seconds Measured for the Platinum Diverse Dataset.
10
50
250
500
Maximum ensemble size
Mean
Median
Mean
Median
Mean
Median
Mean
Median
Balloon DG
6
5
27
24
132
117
260
260
Balloon GA
4
3
19
17
105
98
256
234
RDKit
1
1
5
4
22
18
42
34
ETKDD
1
1
4
3
16
12
32
23
Multiconf- -LOC
5
1
8
2
15
3
15
3
Tensor VAE
<1
<1
<1
<1
1
1
2
2
3.3 Ablation studies
In this section, we further demonstrate the effectiveness and necessity of running an ID convolution with N x kernels over the proposed input tensor through 4 ablation studics GEOM drugs datasct. We also show that the transformer attention mechanism is also important contributing factor for competitive generative performance.
Why is 1D convolution necessary. We have shown a model based on a 3 x 3 kernel in See A.4 called NaiveUNet. Here, we provide a more detailed analysis of why NaiveUNet produces unsatisfactory result. The primary reason for this poor performance the "field of view" of a conventional x d (d < N) kernel only sees a partial connection pattern of a focal atom. In comparison,
N x 3 kernel's "field of view" encompasses the complete connection pattern of focal atom. We further observe that when applying 3 3 kernel filter to 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 a 1-hop information aggregation. Secondly when the 3 x 3 kernel moves to an off- diagonal part of the tensor, where most connections are virtual bonds (as atoms of a molecule are often sparscly connected), information aggregation occurs mostly between atoms that are not chemically connected and therefore less meaningful than that on the diagonal part of the tensor. For these two reasons, the Neivs performance the GEOM Drugs dataset is the worst shown in Tab
13

---

Published in Transactions on Machine Learning Research (01/2024)
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 a "2-hop atomic-envirommnn" (because the focal atom can "see" how neighboring atoms are chemically connected to all their direct neigh- bors). Another observation is that after removing all virtual bonds, each column does not correspond to a fully-connected MPNN. Therefore it no longer enables a global information aggregation. The conformation generation results of this variant of TensorVAF. on Drugs dataset shown as as TensorVAE ablal table below. It is observed that due to a less effective local information aggregation as result of removing all virtual bonds (and related atom features), the performance is worse than of the complete Tensor version.
What happens if a N x 1 kernel is used. The third ablation study concerns with using a N x 1 kernel with a smaller "field of view" as compared to that of a N x 3 kernel. Its performance on Drugs dataset is shown as TensorVAR abla2 in Tab[9] 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 global information aggregation for the focal atom. Nevertheless, underperforms the complete Tensor version due to a smaller 'field of view' for information aggregation.
What happens if a 1 x 1 kernel is used. This setup corresponds to connecting a fully-connected MPNN (GNN) with a standard transformer backbone for conformation generation. Since using a 1 x 1 kernel leads to a model with a signicantly less model compacity as compared to the models in previous ablation studies, we experimented with 6 parammerr configurations listed as following to ensure this variant has roughly the same model capacity (number of parameters).
Table 8: Experimental setups for 1 x 1 kernel.
Model name
Embedding size
KL weight schedule
No. of transformer layers
No. of parameters
GNN
base
256
same as Tensor
6.5M
GNN largel
320
same as TensorVAD
4
11M
GNN
large2
256
same as Tensor
6
10M
GNN
large3
320
1c-5 doubling every 16 epochs
6
11M
GNN
large4
320
1e-6 doubling every 16 epochs
6
11M
GNN
large5
320
1e-7 doubling every 16 epochs
6
11M
There are two ways to increase the number of parameters of the MPNN-based variant to match that of the TensorVAF employing N x 3 kernel for fair comparison, including a. larger embedding size and more transformer layers. These two setups correspond to large 1 and large 2. Unfortunately, training for these 3 setups failed to reduce RMSD error after more than 10 epochs of training; we kept facing the KL vanishing problem. To tackle this, we experimented with 3 more configurations (large 3,4 and 5) with much lower KL weights and shorter step period to force training to focus more on reducing the RMSD loss. Unfortunately again, after more than 40 epochs (25+ hours) of training all three efforts have also failed to resolve this issue. We have included the training and validation curve for all experiments Fig
It is observed that for all cases, while KL error quickly decreases to close to zero, the RMSD loss stays almost constant at 4.0, indicating model's inability to learn. It seems that the MPNN-based models struggle to learn any meaningful information that contribute to producing valid conformations. Instead, they always resort to reducing KL loss which is a much easier learning task. This fact combined with previous 3 ablation studics manifest an emerging trend that the TensorVAE modcl's capacity to learn difficult conformation generation task improves with the increase of expressive power of its aggregation mechanism. In other words, the extra flexibility introduced by the increasod kernel size (from x 1 to N 3) is main contributing factor to the promising performance of the Tensor VAE model. Therefore, we conclude that the design choice made to use a N x 3 kernel is sensible and fully justified.
What happens if the transformer architecture is replaced by a MLP. In this sctup, we replace the transformer encoder block with a MLP block as following:
m = W Wa(RELU + b1)) + b2
h+ DropOut(LayerNorrtnn)
14

---

Published in Transactions on Machine Learning Research (01/2024)
train
validation loss
2000
2000
KL (log scale)
80-
logs GRIN, basntrair
KL KLisoss ttop scale)
logs, CNN, bass/vaidaaion
logs ONN large L/wrin
logs GNN large Livalicttion
GNN, 2/trann logs, BwN.ibrre.3/ran
GNN, large 2ovalicttion
1og1_ ONN hrrg__/trsh
logs, CHNL jarge. hetl_ierr__/aalliattte Alvalicatioon
logs_ FRIN Jerge Sirain
0.12
logs, CHN Jarge Sivaliiwtte
0.12
tetsorvwe ben
0.005
0.0
0.005
0.9035
6.0005
10
20
30
40
0
10
20
30
Epoch
Epoch
RMSD train loss
RMSD validation loss
3-
3-
2
CRNN besetrain
GNN hww.://we/eelllnee vidaiio
1.2
logs,
Lisalisation
logs GNN large 2uriin
logs GNN larga 2ovalication
0.8
logs GNN large 3/trann
0.8
logs, GNN, large 3ivalication
logt_ GNNN harge 4/train
15gs_ GNNLaaggg, @tt.brrge_haalloote 4/valcottto
log1_ GeitttnrgeSSttra
RMSD
0.4
logs_ Jersorvee basettan
lngs,
0.-
20 Epoch
20
30
Epoch
Figure 3: Ablation study: performance comparison between using N x 3 kernel and 1 x kernel. The model architectural difference between GNN large models and the GNN base model can be found in Tab
Where h/ E R256x1 is the 1th layer output for the 7th atom, W1 E R2024X256 by E R1024X1 W2 1256 and be E R256x1, Additionally, due to the absence of allention mechanism, the output (hf.... likb) of the graph encoders σe, (G) and the sampled latent output {a1, ... EN} J of the posterior encoder (z G) are simply summed to become the input of the of the decoder po, (R|2, da, (G)) to generate conformation directly. Akin to the TensorVAE, all three components consist of 4 MLP blocks. The total number of parameters corresponding to this setup 11M which is similar to that of the TensorVAF.
The dropout rate was initially set to 0.1. In this configuration, we trained the model for 100 epochs and observed a severe overfitting issue, as illustrated in Fig where the training and validation curves of the ML. variant and TensorVAF, are compared. The MLP variant exhibited not only a significantly higher KI. loss but also a much higher RMSD validation loss compared to both its training loss and that of TensorVAE. Upon observing this behavior, we decided to experiment with higher dropout rates, including 0.3, 0.5, and 0.7, to mitigate overfitting.
After training the model for 100 epochs for each dropout rate, we found that the dropout rate of 0.3 achieved the best validation KL and RMSD losses without encountering any overfitting issues. However, with an increase in the dropout rate, the RMSD training loss decreased (while the validation RMSD error remained the same), and the KL loss incrcased significantly. This behavior mirrored that of the 1 x 1 convolution kernel (fully-connected MPNN) variant mentioned earlier. Essentially, the model increasingly relied on posterior encoder information to reconstruct the conformation and reduce the RMSD error, which is an casier task compared to reconstructing conformation from a 2D molccular graph in the absence of any coordinate information. This trend suggested that a higher dropout rate led to reduction in the model's capacity to learn.
15

---

Published in Transactions on Machine Learning Research (01/2024)
Despite achieving the best performance among all tested dropout rates, after 120 epochs of training, the MLP variant with 0.3 dropout still performed significantly worse than the TensorVAE, with transformer backbone. Although its RMSD validation loss matched that of the TensorVAE, its KL validation loss was more than double that of the TensorVAE, indicating significantly lower learning capacity. Observing this behavior led us to the conclusion that was no longer necessary to complete the training to remonstrate the necessity of TensorVAF. with transformer architecture to obtain competitive performance. This experiment suggests that the attention mechanism is also a crucial contributing factor for effective information aggregation among atoms.
train loss
KL validation loss
2000
logs, www.ornnnceeee.ooooooooooooeeeee basettrain
logt, for sorvoc, baseivaloationn
logs, gepp_tenoonvee_mipe mip.train
logs, ton sorvie mipooalicattr
SoPs_tenorvve_mte dr
logy_ twe
2000
nip d:5.ttr
1,000
logs, mip dt Stealidation
logs Lensorvaw nip d: 0.7kmann
logs sanservae mp dr 7.eeadidation
800
115510 Da1e)
[log Malle)
1.00
500
400
500
400
100
300
250
250
200
200
100
100
20
40
80
20
40
80
Epoch
Epoch
RMSD train less
RMSD validation less
logs tonservae, baseivaliation
2-
3
logr_tensorvvee, mmp_or_0.3vvlldettee
2-
Ings, twe sorooeeee,eee __p__pppp_____ Rrvaiiidation
logo, tersorv mip dr Predlidtioon
0.9
IMSD loss 1kgg scul))
basetrain
0.6
logs, tensorvee mipitrann
0.4
loga, Lemsorvaa nip d: 3train
logs tansorvae mip dr 0. 0.5kran
0.3
logs, www.tencorvee..o mip 0.7kaan
0.15
0.24
0.38
0.09
0.32
20
-00
60
40 Epoch
80
Epoch
Figure 6: Ablation study: performance comparison among MLP backbones with different dropout rates. While the base TensorVAB with a transformer backbone and the base MLP ablation studay have default dropout rate of 0.1, the other ablation studies have dropout rates ranging from dr 0.3 to dr 0.7, respectively.
Table 9:
Performance comparison among models with different input feature engineering setup on GEOM
Drugs dataset
COV
MAT
Method
Mean
Median
Mean
Median
UNe
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 VAE 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.
16

---

Published in Transactions on Machine Learning Research (01/2024)
4 Reproducibility statement
We did not introduce any task-specific neural network archiecture. The results presented in this study can be straightforwardly reproduced using publically available datasets and ready-to-use implementation of convo- lution and Transformer from either PyTorch or TensorFlow. We have also provided detail syper-paramett sctup to ensure reproducibility. We have included the complete code for reproducing the conformation gen- eration results lattps://anonymontt woaeooyyouu..ppen..coo / and code for reproducing the property prediction results in https: httpt:://annntttie http::////oommmuussuuu wttpa///aaaaaaeee/////oo/a/aa/..
5 Conclusion
We develop TensorVAB a simple yet powerful model able to generate 3D conformation directly 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 decompose these techniques into three main idcas, and explain how one idea naturally connects to 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 1D convolution manner achieve global information aggregation. Finally, present the complete CVAE-based framework featuring 2 transformer-baser encoders and another transformer- based decoder, and propose novel modification to the first multi-head attention layer of the decoder to enable sensible integration of the output of the other two encoders.
The proposed Tensor VAE demonstrates state-of-the-art generative performance compared to recently pro- posed deep-learning- based generative models on the GEOM dataset, utilizing DFT-generated unbound con- formations. When directly applied to the Platinum dataset, which contains ligand-protein bound confor- mations, the proposed method offers faster generation speed while maintaining competitive accuracy compared to 5 popular and classical methods.
Limitations and Future Directions. Despite achieving promising performance in conformational gen- cration, the current work has three major limitations that pave the way for future improvements. Firstly, the proposed lensor graph representation lacks invariance under random permulations of atom ordering. While experimentally robust, achieving true invariance such transformations would enhance the stability of conformation generation.
Secondly, the training process only achieves approximate SE(3) invariance due to the presence of R in the imput of the posterior encoder. Aiming for exact invariance has the potential to further improve the TensorVAB framcwork's performance. To address this, we plan to replace In (2| G) with the recently proposed equivariant posterior encoder component from the Geometric AutoEncoder framework in GeoLDM [Xu al. 2023).
Thirdly, the current TensorVAB can only predict the local structure of molecules, specifically heavy atom coordinates with respect to an arbitrary origin. It lacks the capability to predict the SE(3) transformation necessary to obtain the bounding pose with respect to protein target-a crucial aspect for structural- based drug discovery tasks. Our next objective involves expanding Tensor VAE's architecture to predict both unbound ligand and protein conformers as input and produce valid bound ligand conformation as output. To achieve this, we are integrating the SE(3) equivariant convolution operation proposed in the Tensor Field Networks Thomas al. 2018) into the TensorVAF framework. This expansion aims to enhance the modcl's ability to generale conformer ensembles suitable for docking to specific protein targels.
References
K Somani Arun, Thomas S Huang, and Steven Blostein. Least-squares fitting of two 3-d point sets. IEEE Transactions on pattern analysis and machine intelligence, (0)=098-70,, 1987.
Simon Axclrod and Rafacl
Gcom, energy-annotated molccular conformations for property prediction and molecular generation. Scientiffc Data, 91):1885 2022. doi: 01/////111997-22-28-----3
URL https::/tpsatt httpppp:/aaa.oorr/....o/////////
17

---

Published in Transactions on Machine Learning Research (01/2024)
Helen M Berman, John Westbrook, Zukang Feng, Gary Gilliland, Talapady N Bhat, Helge Weissig, Ilya N Shindyalov, and Philip E. Bourne. The protein data bank. Nucleic acids research, 28(1 (21())233-220))) 2000.
Hongming Chen, Ola Engkvist, Yinhai Wang, Marcus Olivecrona, and Thomas Blaschke. The rise of deep learning in drug discovery. Drug discovery today, 23(6):1241-12500 2018.
Xiaomin Fang, Lihang Liu, Jieqiong Lei, Donglong He, Shanzhuo Zhang, Jingbo Zhou, Fan Wang, Hua Wu, and Haifong Wang. Geometry-enhanced molecular representation learning for property prediction. Nature Machine Intelligence, 4(2):127 134, 2022.
Nils-Ole Friedrich, Agnes Meyder, Christina de Bruyn Kops, Kai Sommer, Florian Flachsenberg, Matthias Rarey, and Johannes Kirchmair. High-quality dataset of protein-bound ligand conformations and its ap- plication benchmarking conformer ensemble generators. Journal of chemical information and modeling, 57(3):529-539, 2017.
Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, and Lawrence Carin. Cyclical annealing schodule: A simple approach mitigating kl vanishing. In Proceedings of the 2019 Conference of the North American Chapler of the Associalion for Computational Linguislies: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 240-250, 2019.
Niklas WA Gebauer, Michael Gastegger, Stefaan SP Hessmann, Klaus-Robert Müller, and Kristof T Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. Nature communi- cations, 13(::1 11, 2022.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263 1272. PMLR, 2017.
Tarun Gogineni, Ziping Xu, Exequiel Punzalan, Runxuan Jiang, Joshua Kammeraad, Ambuj Tewari, and Paul Zimmerman. Torsionnet: A reinforcement learning approach to sequential conformer search. Advances Neural Information Processing Systems, 14:2012--2013, 2020.
Thomas A Halgren. Merck molecular force field. v. extension of mmff94 using experimental data, additional computational data, and empirical rules. Journal of Computational Chemistry, 17(5-6):616-644 1996.
Paul CD Hawkins. Conformation generation: the state of the art Journal of chemicol information and modeling, 1777711777 2017.
Shion Honda, Hirotaka Akita, Katsuhiko Ishiguro, Toshiki Nakanishi, and Kenta Oono. Graph residual flow for molecular graph generation. arXiv preprint arXiv:1909.1985 2019.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. In International Conference on Learning Representa- tions, 2020. URL http://s///gppeeeeeeeaaeeeeerrea htt..//foorrrrr//d/11/......
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(7873) 589, 2021.
Diederik P. Kingma and Jimmy Ba. Adam: method for stochastic optimization. ICLR (Poster), 2015.
URL http://arxxv.orgg/ wwwpoo/g///bb//o/o/eee hetp://11...60000 6980
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Represeniations, 2014. URL https://openreviess http:////e.....oooooo//////////0 9FyZd.
Jonas Köhler, Leon Klein, and Frank Noó. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pp.
PMLR, 2020.
Yibo Li, Jianfeng Pei, and Luhua Lai. Structuro-basedd de novo drug design using 3d deep generative models. Chemical science, 12(41):13664 -36755 2021.
18

---

Published in Transactions on Machine Learning Research (01/2024)
Shengchao Liu, Mehmet F Demirel, and Yingyu Liang. N-gram graph: Simple unsupervised representation for graphs, with applications to molecules. Advances in neural information processing systems, 32, 2019.
Shitong Luo, Chence Shi, Minkai Xu, and Jian Tang. Predicting molecular conformation via dynamic graph scorc matching. Advances in Neural Information Processing Systems, 34:19784-19795, 2021.
Omar Mahmood, Elman Mansimov, Richard Bonneau, and Kyunghyun Cho. Masked graph modeling for molecule generation. Nature communications, 12(1):1-12, 2021.
Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. Molecular geometry prediction using deep generative graph neural network. Scientific reports, 9(1):1 11, 2019.
Sereina Riniker and Gregory A Landrum. Better informed distance geometry: using what we know to improve conformation generation. Journal of chemical information and modeling, 55(1)):2562 2574, 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 Neural Information Processing Systems, 33:12559 12571, 2020.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference Medical image computing and computer-assisted interven- tion, pp. 241. Springer, 2015.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pp.
PMLR, 2021.
Nicolas Sauton, David Lagorce, Bruno O Villoutreix, and Maria A Miteva. Ms-dock: accurate multiple conformation generator and rigid docking protocol for multi-step virtual ligand screening. BMC bioinfor- matics, 911::1112, 2008.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International Conference on Machine Learning, pp. 9558-9568. PMLR, 2021.
Gregor N. C. Simm, Robert Pinsler, Gábor Csányi, and José Miguel Hernández-Lobato. Symmetry-awarre actor-critic for 3d molecular design. International Conference on Learning Representations, 2021. URL lattps://openrevier .ett/forrun/iirrj88//EEE
Gregor NC Simm and José Miguel Hernández-Looato. A generative model for molecular distance geometry. In Proceedings of the 37th International Conference on Machine Learning, pp. 8949 8958, 2020.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan,, and Surya Ganguli. Deep insupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp. 2256-2265. PMLR, 2015.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and cranslation-equiivriint neural networks for 3d point clouds. arXiv preprinl arXiv:1802.0920 2018.
Mikko J Vainio and Mark S Johnson. rithm. Journal of chemical information and modeling, 47(6):2462 2474, 2007.
Generating conformer ensembles using multiobjective genetic algo-
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkorcit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Alexandra Volokhova, Michal Koziarski, Alex Hernández- Garefa, Cheng-Hao Liu, Santiago Miret, Pablo Lemos, Luca Thiede, Zichao Yan, Alán Aspuru- Guzik, and Yoshua Bengio. Towards equilibrium molecular conformation generation with gflownets. arXiv preprint ar Xrxi::2310. 14782, 2023.
19

---

Published in Transactions on Machine Learning Research (01/2024)
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Ancesh S Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science, 922):513 530, 2018.
Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li, Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, et al. Pushing the boundaries of molecular representation for drug discovery with the graph attention mochanism. Journal of medicinal chemistry, 63(16))8749-8700, 2019.
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. In International Conference on Learning Representations, 2021a.
URL http::///ppparrrrreee..e///rrrrr
Minkai Xu, Wujie Wang, Shitong Luo, Chence Shi, Yoshua Bengio, Rafael Boanez--mmbarllii, and Jian Tang.
An end-to-end framework for molecular conformation generation via bilevel programming. In International Conference Machine Learning, pp. 11537-11547. PMLR, 2021b.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference Learning Repre- sentations, 2022. URL
Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovee. Geometric latent diffusion models for 3d molecule gencration. In International Conference on Machine Learning, pp. 38592-38610. PMLR, 2023.
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 representations for property prediction. Journal of chemical information and modeling, 59(8):3370 33003-3333888 2019.
Hongyang K Yu and Hongjiang C Yu. Powerful molecule generation with simple convnet. Bioinformatics, 38(13):3438 3443, 2022.
Gengmo Zhou, Zhifeng Gao, Qiankum Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. In The Eleventh International Conference on Learning Representations, 2023. URL https::/opparrriiee..effffooo 6K2RM6HVQKu
Jinhua Zhu, Yingce Xia, Chang Liu, Lijun Wu, Shufang Xie, Yusong Wang, Tong Wang, Tao Qin, Wengang Zhou, Houqiang Li, Haiguang Liu, and Tic- Yan Liu. Direct molecular conformation generation. Thansac- lions on Machine Learning Research, 2022. ISSN 2835-8856. URL. httpp:://penrreviex net/forum??d io
ICPOHiztuw
A Appendix
A.1 Global information aggregation beyond the -oop
A geometric interpretation of GNN's message passing layer is it aggregates information between atoms (and their bond) that are 1-hop away. With layers, information from atoms that are I-hop apart can be aggregated. Here, we define a global information aggregation as the -oop aggregation with N being the total number of atoms, where each alom is able to aggregate information from all other atoms.
It is worth noting that for fully-connected GNN, 1-hop message passing can already achieve this global information aggregation. Transformer's 's self-attention can be considered as type of fully-connected GNN. However, a vanilla transformer can only aggregate features from each loken/atom; if edge fealures are not included, they needed to 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 represontation is we want each generated token contain both atom and bond features, such that we can eliminate the pair interaction bond matrix. To achieve this, we fill each column of the fully-connected tensor with;
20

---

Published in Transactions on Machine Learning Research (01/2024)
focal atom features;
chemical and virtual bond features indicating how the focal atom is connected to all other atoms;
atom features of all connected atoms, since for off-diagonal cell, we sum atom features of both the neighbour atom and the focal atom.
Running a N x 1 kernel filter on the proposed tensor also achieves global information aggregation. By increasing kernel width to 3, the aggregation window includes global information from two immediate neigh- bours. This type of information aggregation extends far beyond just Noh-hop.
More interestingly, when multiple kernels are applied simultaneously kernel is free to choose whichever group of atom /bond features to attend to depending on its kernel weights. This resembles the multi-head atlention mechanism of a transformer, where each kernel(head) contributes to portion of the generated feature token. We believe the effective global information aggregation driven by these two (tensor representation + 1D Conv) simple yet intuitive ideas is the main reason why the proposed TensorVAS achieves superior performance with a much less number of parameters.
the same N 3 x region, each
A.2 Connection to fully-connected Message Passing GNN
We show that information aggregation achieved by running a 1 x 1 convolution over the proposed tensor representation is similar to that achieved by a fully-connected MPNN (Gilmer et al. 2017
When running a x 1 convolutional operation over the proposed tensor, a W E RixIxExC kernel matrix is shared among all x cells of the tensor, where F is the number of kernels and C the channel depth. Since each cell, regardless it is on-diagonal bond feature vector, the weight matrix can be decomposed into two parts, W, E RFxC, and We RFxx. where C, is the atom feature vector size and C. is the bond feature vector size. The bond feature vector for on-diagonal cells is filled with zeros, since there is no self connection for focal atoms. Subsequently, for each column n of the tensor, a 1 x convolution operation followed by a sum-aggregation over the rows can be decomposed into 3 steps;
off-diagonal, is stacked with
atom feature vector and
Off-diagonal cell aggregation. For each off-diagonal cell, we first sum the atom feature vectors of the focal atom and its cell-specific neighbour atom, as described in Fig[]] Due to convolution operation, the dot product of the summed vector and W, is then computed. Simultaneously, the dot product between the bond feature vector and W. is also computed. The resulting two feature vectors are added together. This aggregation process can be expressed as;
Woh W W.hh Ween,m E Rx11
where No E RC.x1, 4o E RCox1 and 6n,m E RCox1 are the focal atom feature of the n1th col- umn, neighbour atom feature the mh cell in column n, and bond feature between the th focal atom and its th neighbour atom, respectively. If we concat ho and en.m into single vector (he 1o m this operation also be represented as;
M (he, No Cn,m) = (W,, We, W.) en (n (h
nr Ag Cm.) ) (C(0))))
Row-wise aggregation. The above aggregation operation generates feature vector (of size RFx1) for each off-diagonal row of column n. The sum- num-aggregation over these rows generates feature vector which contains the aggrogated information from all neighbour atoms.
m
Σ MEN,
M (hg Ag Cn,m)
Complete aggregation. Finally, aggregate this feature vector my onto the focal atom feature to complete the sum-aggregation operation over all the rows of column n.
we define (+) concatenation operator
MPNN Gilmer al. 2017
21

---

Published in Transactions on Machine Learning Research (01/2024)
h U (h2 m.)) ReLu (W.hh + mm,)
Noticeably, the M and U operators correspond exactly to the message passing phase of a single forward pass of a fully-connocted MPNN, as described by Eqs.1 and 2 of the MPNN paper (Gimeer et al. 2017)
Similarly, the feature aggregation operation of N kernel can be expressed as;
h ReLu W +
Σ MEN
(W"Thng Wh + W en.m)
This type of aggregation is more flexible and has more expressive power as different node and edge features are weighted differently. This llexibility further increased with N x kernel whose corresponding aggregation can be expressed as;
(aj.k)
(R...)
Noc
h/ = ReLu
Σ
We +
Σ
Σ
We hg + Worn + Weo Ec.m
where i,j. are the indices of three adjacent columns. In this respect, the information aggregation achieved by a fully-connected MPNN is a special case (the simplest form) of more general framework embodied by single convolution operation over the proposed tensor representa- tion.
22

---

Published in Transactions on Machine Learning Research (01/2024)
A.3 Training hyperparameters
Conformation generation. Training is conducted single Tesla V100 GPU. We follow a similar learning rate schedule, shown by Eq.3 of the original Transformer paper (Vaswani et al. 2017) but with dmodel 9612. This results a maximum learning rate of 1.6e To tackle the notorious issuc of KL vanishing (Fu ct al. 2019), we set a minimum KL weight of le- and double every 62.5e iterations until a maximum weight of 0.0256 reached. We select Adam optimizer (Kingma & Ba 2015) default hyperparameters for training. validatiin
We present some interesting observations of the training/validatiion curve corresponding to this setup in A.7 For both experiments, the TensorVAE is trained for 1e iterations with batch size of 128. The implementation details of NaiveUNet. is explained in A.4
Molecular property prediction. For the molecular property prediction task, we use the same GDR transformer encoder structure (4 attention layers and approximately 5M parameters) and add an additional mean pooling layer, which is then followed by a linear layer for property prediction. We follow the same data train-val-test split in Uni-Mol and GEM and standardize the output property data. We train the GDR model for 300 epochs with batch size of 128. The learning rate schedule the same that of the TensorVAL
23

---

Published in Transactions on Machine Learning Research (01/2024)
A.4 NaiveUnet model architecture
Encoder
Decoder
3 x Conv Rernel
Masited output
Skip connection
due to symimetry
Fenture depth
9
88
5
9
V
5
6
61
0.30
0.66
C
6
64
123
128
9
6
8
0.1
0.25
C
256
184
N120
NP16
NAOSE
NP12
Encoder Block
Decorer Block
Figure 7: Naive UNet model. N = 69
We train the above NaiveUNet on the Drugs dataset for 30 epochs with constant learning rate of 1e 1e-4 and batch size of 32. We follow the same method presented in GraphDG (Simin & Hernández-Lobail 2020 to convert the predicted distance matrix to conformation.
A.5 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 10
Table 10: Atom and bond features used construct input tensor.
Feature name
Feature value
Encoding method
Atom type
H, C, N, O, F, S. CL, Br, P. I, Na, B, Si, Se, K, Bi
one-hot
Atom charge
-1, 1, 2
one-hot
Atom chirality
Unspecified, Tetrahedral CW Tetrahedral CCoo Other Single, Double, Triple, Aromatic, Virtual
one-hot
Bond type Normalized bond length Bond stereochem
onc-hot
real-value
StereoNone, StereoAny, StereoZ StereoE,, StereoCIS, StereoTTrans 3 10
one-hot
Bond in-ring size Coordinate (3 channels) Pair wise atom distance
one-hot
real-value
real-value
A.6 COV and MAT precision results
The precision COV and MAT scores are defined as;
{REE C,IRD (R. &.
COV (C,,(C,,C)) = C
MAT (C,, C2)
min RMSD (R. R
C Rec.
24

---

Published in Transactions on Machine Learning Research (01/2024)
Table 11: Precision performance comparison on GEOM Drugs dataset
COV (%)
pATP((A +
Method
Mean
Median
Mean
Median
GraphDG
2.08
0.00
2.4340
2,4100
CGCF
21.68
13.72
1.8571
1.8066
ConfVAE
22.96
14.05
1.8287
1.8159
ConfGF
23.42
15.52
1.7219
1.6863
GeoDiff
61.47
64.55
1.1712
1.1232
72.12
79.02
1.0655
1.0355
VAn5
+1.5
+1.
#0.0145
#
*Results for GraphDG, CGCF, Con(VAE, ConfGF and
GeoDiff are taken from (Xu et al.
2022
25

---

Published in Transactions on Machine Learning Research (01/2024)
A.7
Training and validation curve
We present the train and validation plots for KL and reconstruction loss based on Drugs datasel in Fig 8a and Fig 8b respectively. Both plots are based on an initial KL weight of Ie- doubling every 62.5k iterations (40 epochs). While KL validation loss reached 18.29 after 1e iterations (640 epochs), the reconstruction/RMM)) loss reached 0.64À at the end of training. During the first epochs of training, model learning focused on reducing the KT. loss due it orders of magnitude larger than the RMSD loss. We were expecting this trend continue for while until both losses converge roughly in the same range. However, much to surprise, the model seemed to find way to drastically reduce RMSD loss much carlier 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 low as 0.08A. On the other hand, the KL loss climbed to 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 cheat since training is forcing the output of GDR encoder to conform to standard uninformative Gaussian distribution. The KL loss started drop while the RMSD loss remained steady, indicating increasing reliance 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 to initiate the training with a much larger initial KL weight (Le (le=3) to prevent 'cheating" from begining. However, this quickly led to the notorious KI vanishing issuc (Fu al., 2019). We figure that "cheating" is actually beneficial in that it reduces learning difficulty particularly for the decoder; its weights are tuned on easy training task, simply reversing what GDR encoder has done. In other words, the tuned weights of 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-propagation 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.
800
300
70
0.
20
0.06
100
200
300
400
500
600
0.
100
200
300
400
500
600
(a) KL loss.
(b) RMSD loss.
Figure Training and validation plots for Drugs dataset. Orange line: Train; Blue line: Validation
26

---

Published in Transactions on Machine Learning Research (01/2024)
A.8 Molecular property prediction
Following Uni-Mol (Zhou el al., 2023 and GEM Fang et al. 2022), report properly prediction result on the MolecularNet (Wu et al. 2018 QM9 regression task. The goal of this task is to estimate homo, lumo, and homo-lumo gap propertics of molecules in the QM9 dataset based on their molecular structure. We adapt the proposed GDR encoder to this regression task by changing its prediction head. We defer the details of this adaption and training procedure to Sed A.3. We report the mean average error(MAF. over all the test samples.
The result of the adapted model is compared to those of 7 other models including:
D-MPNN Yang et al. 2019), AtlentivePP Xiong al. 2019 and GEM which are GNN based models without pretraining;
N-Gram [Lu et al. 2019). pretraining;
PretrainingGNN
Hu et al. 2020 and GROVER (Rong et al. 2020 with
variant of Uni Mol without pretraining.
The MAE for compared methods are summaried in Tab 12 The proposed GDR encoder produces a
SOTA performance with less than 5M parameters. This experiment demonstrates that the proposed feature engincering method is very effective at information aggregation.
Table 12: Property prediction result comparison based on MolecularNet QM9 benchmark.
Method
MAE
D-MPNN
0.00814 (0.00001)
AttentiveFPP
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 pretraining
0.00653 (0.00040)
GDR encoder (ours)
0.00553 (0.00012)
All results are taken from Zhou et al. 2023
Values in heses are standard deviation ob-
tained by repeating experiments times.
27


Reviews:
[{"review_id": "S7NOq9iTjn", "reviewer": "Reviewer_1HGp", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The paper introduces a new representation for 2D molecular graphs that appears useful for conformer generation and additional tasks, such as property prediction. \n* The paper provides an extensive set of experiments across relevant datasets (GEOM, Platinum) and compares with relevant methods. TensorVAE generally shows performance improvements and distinct advantages, such as better compute efficiency. \n* The paper provides significant detail on the method and experimental settings.\n\nWeaknesses:\n* I think the claims of single step generation of conformers and simple architectures are somewhat overstated. Part of the text and the ablations in the appendix seem to indicate that neural network architecture does matter in performance, so it would be good to be clearer about the paper claims as important.\n* It is unclear why the paper includes zero-shot experiments on Platinum only and why molecular property prediction experiments are also included in the main text. Given the focus on conformer generation, I think the results of the ablation would be more interesting to understand the important components of the proposed method.\n* Related work is discussed primarily in the introduction and is missing a discussion on sequence-based methods for conformer generation, such as reinforcement learning [1] and GFlowNets [2].\n\n[1] Gogineni, Tarun, et al. \"Torsionnet: A reinforcement learning approach to sequential conformer search.\" Advances in Neural Information Processing Systems 33 (2020): 20142-20153.\n\n[2] Volokhova, Alexandra, et al. \"Towards equilibrium molecular conformation generation with GFlowNets.\" arXiv preprint arXiv:2310.14782 (2023).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new method for generating 3D molecular conformers from 2D graphs. The primary contribution of the method is encoding the input features of the 2D molecular graph as a tensor graph which includes atomic encodings (atom type, atom charge, atom chirality, atom coordinate) as well as bond encoding (bond type, bond stereochemistry, normalized bond length, distances) for all of the atoms in the system. In the tensor graph, the diagonal elements include the primary while the off-diagonal atoms represent information about the neighbors of the atom. Based on this tensor graph, the paper then proposes using a 1D-convolution as the primary neural network operation to process the tensor graph followed by a self-attention transformer. Taken together, the above make up the primary components of TensorVAE which is the paper's proposed method.\n\nNext, the paper describes the formulation of the problem, including the VAE training loss and how it applies to the conformer generation case study and outline their primary experiments. The experiments focus on three main datasets: GEOM-QM9, GEOM-Drugs, and Platinum. The paper outlines the experimental conditions and metrics in detail before presenting the results on GEOM, where TensorVAE generally out performs the other methods the author compare against. The results include methods that apply force-fields for conformer generation in Table 2 and methods with and without force-fields in Table 1. In Table 3, the paper claims that TensorVAE outperforms other method while having less parameters and outline details about the compute efficiency during of TensorVAE compared to other methods.  Next, Table 4, Table 5 and Table 6 the paper provides an analysis of zero-shot performance of TensorVAE on the Platinum dataset, which includes molecular conformation in a protein environment. The results generally show reasonable performance on the Platinum dataset. Lastly, the authors provide a brief study on molecular property prediction along with a reproducibility statement and a conclusion.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Important Requests that would sway my opinion:\n* Please adjust the claims related to single-step generation and simple architectures. Are there other methods that perform conformer generation in a single step? How do they compare to TensorVAE.\n    * You mention that your method does not require sophisticated neural network design, so it would be good to see performance with even simpler parts, such as MLPs.\n* Add related work of sequence-based methods for conformer generation.\n* Could you clarify why you only ran zero-shot experiments for Platinum? It seems like fine-tuning or training experiments for TensorVAE on Platinum would be appropriate.\n\nAdditional requests (nice to have):\n* I would prefer seeing results of the ablation in the main text and would be OK moving the property prediction experiments to the appendix in exchange. Feel free to explain why you included property prediction in the main text.\n* Could you clarify if the tensor graph representation is symmetric. If so, is there a way to take advantage of the symmetry? This could be interesting to discuss in the conclusion and future work section.\n* It would be nice to expand potential future work directions.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "bwgYw7Sq8v", "reviewer": "Reviewer_cqxm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The paper is well-written and easy to follow. The key ideas and contributions are claimed clearly.\n\n2. Many baselines are considered, and experimental results seem promising.\n\n3. The method itself is simple and easy to reproduce.\n\n\nWeaknesses:\n\n1. If my understanding is correct, the 1D convolution applied to the graph tensor is not permutation-invariant, i.e., permuting the order of the rows and columns of the graph tensor leads to different token sequences. Accordingly, although the method is simple, I think it will lead to different 3D conformation coordinates when changing the order of atoms in the 2D graphs, which does harm to the rationality of the proposed method.\n\n2. The proposed method seems to focus on generating conformations conditioned on 2D graphs. It might be more suitable to change the title to \"conditional generation\".", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Summary:\n\nIn this submission, the authors propose a simple but effective method for conditional molecular conformation generation, which generates 3D molecular conformations from the corresponding 2D graphs. Technically, this work 1) proposes a tensor-based graph representation, encoding the atom and bond information jointly in a tensor format, and 2) applies 1D convolution to tokenize the tensor, leading to the following Transformer-based encoder. Therefore, the main technical contributions of this work are feature engineering and encoder architecture, in my opinion. Experimental results demonstrate the usefulness of the proposed method in several datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Each molecule may have various 3D conformations, is the proposed method able to generate various 3D conformations from a single 2D graph?\n\n2. What is the full name of \"GDR\"?\n\n3. It seems that the outputs of the decoder are 3D atom coordinates. Did the authors use other tools like OpenBabel to generate conformations from the coordinates?\n\n4. When implementing the 1D convolution, a kernel with size N x 3 is applied. If my understanding is correct, N is the number of atoms in a molecule, which varies for different molecules. How to learn such a size-varying kernel?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "hxiY6ZOMl2", "reviewer": "Reviewer_SZqZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "An overarching weakness of this paper is its clarity. Although I have certainly read less clear papers, many things were not described very clearly in this paper: particularly the model (e.g. how $p_\\theta(R|z,G)$ decomposes into several sequential decoders and how the attention masking works). I think an explicit algorithmic statement of the model would help for clarity. This issue makes it difficult to properly assess the correctness and impact of the claims made in the paper. The following review content is therefore based on my best guess about what I think the authors are doing.\n\nThe next thing I would comment on is the model and featurization (the efficacy of which is a key claimed contribution from the authors). The proposed featurization seems to just be an adjacency matrix with features, differing from a standard MPNN only in that the atom and edge features are explicitly mixed rather than being maintained separately. This seemed like a fairly small implementation detail to me. Next, the authors propose a kind of masked convolution, which seemed very similar to the attention mechanism of a transformer. In the majority of the paper the authors seem to refer to their model simply as a transformer, making me believe that it is just a transformer variant. The authors go on to use this transformer as the encoder/decoder of a VAE, which is not very novel (e.g. https://arxiv.org/abs/2207.13529).\n\nFinally, I will focus on the experiments. The reported results of the experiments are quite strong, and if correct would be one of the strongest parts of the paper. However, it is not totally clear to me that these results are correct/comparable with previous work. The authors state that their method is \"simple\" and mainly relies on a good choice of features. It is therefore surprising that it outperforms a lot of other methods by such a large margin. My intuition is that the authors have made some kind of mistake here (I would bet ~3:1 odds). Possible sources of this are:\n- Different random selection of train/test molecules (it was not clear if the exact same set of molecules from Xu et al 2022 was used or whether it was just generated with a similar same procedure)\n- Data discarding: the authors seem to have discarded a fraction of the data which had too many atoms. Even though the fraction is small, presumably the other methods did not discard it. Maybe it has an outsized influence on the overall error?\n- Subtle data leakage: for example, the authors use a depth-first traversal to order the nodes in the tensor. Maybe there is actually some spurious correlation here which the model picks up on?\n- Some other small inconsistency (possible because the authors are copying results from previous papers rather than re-running things). E.g. a different normalization or different value of $\\delta$ in equation 4.\n\nI should clarify that I have not *found* a bug, just that I am highly suspicious of the results.\n\nSumming up:\n\nStrengths of the paper: a generally sensible method, good experimental results (if correct)\n\nWeaknesses of the paper: lack of clarity, little novelty", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes TensorVAE: a VAE based on a specific tensor representation which can be used to both generate 3D conformations of molecules and predict their properties. The contributions are:\n\n- The TensorVAE architecture and data format. However, I am uncertain to what degree this can be considered a contribution: TensorVAE seems to fall into the category of \"specialized GNN\" or \"special case of transformer\", but it is described in a somewhat unclear way that makes me uncertain about this equivalence.\n-  A transformer-VAE hybrid with a modified attention (although again this is not very clearly described so its contribution is unclear)\n- Experiments showing strong results for conformer generation", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "I am not really sure what to write for this: TMLR asks for claims which are well-supported and of interest to the audience. I can't really decide what the key claims of the paper are. I would guess either:\n\n- \"TensorVAE is a great method for predicting conformers\". This claim seems fairly well-supported, but given that TensorVAE seems to just be a transformer variant, I'm not sure how interesting such a claim is: have the authors essentially just tuned a transformer?\n- \"The featurization of TensorVAE somehow really improves upon previous work\": this would be an interesting claim but it is not well-supported, since when comparing to other methods the authors are changing not just the featurization but also the model, so it is not clear what is actually causing the performance.\n\nUltimately my request to the authors is to clarify what the key claims are for the paper and explain how they are supported. For this reason I have written \"claims not supported\" in my review below.\n\nAlso, it is not clear that the method is actually equivariant as the authors suggest: when describing equation 3, the authors say that the encoder \"does not involve any coordinate\", implying it is equivariant. However, it appears that it *does* actually involve coordinates, as these are included in $R$, no?", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work proposes a method for conditional molecular conformation generation. After extensive engagement from the authors, all reviewers agree that the work meets TMLR's acceptance criteria. I thus recommend accepting the paper, and ask that the authors please do revise figures 5 and 6 as requested by reviewer 1HGp.
Decision: Accept (poster)