Title: CONDITIONAL GENERATION OF MOLECULES FROM DISENTANGLED REPRESENTATIONS

Abstract: Though machine learning approaches have shown great success in estimating properties of small molecules, the inverse problem of generating molecules with desired properties remains challenging. This difficulty is in part because the set of molecules which have a given property is structurally very diverse. Treating this inverse problem as a conditional distribution estimation task, we draw upon work in learning disentangled representations to learn a conditional distribution over molecules given a desired property, where the molecular structure is encoded in a continuous latent random variable. By including property information as an input factor independent from the structure representation, one can perform conditional molecule generation via a "style transfer" process, in which we explicitly set the property to a desired value at generation time. In contrast to existing approaches, we disentangle the latent factors from the property factors using a regularization term which constrains the generated molecules to have the property provided to the generation network, no matter how the latent factor changes.

1. INTRODUCTION:

Conditional molecule generation is far from being solved. The main challenge is the enormous and discrete nature of the molecules space and the fact that molecule properties are highly sensitive to molecular structure (Kirkpatrick & Ellis, 2004) . Approaches to conditional generation are typically two-step, either using a model or genetic algorithm to generate candidates which are later filtered, or learning a continuous embedding of the discrete molecules and optimizing in a real-valued representation space. The former is computationally expensive, the latter performs conditional generation only very obliquely.

We propose a conditional generative model that produces candidate molecules which targeting a desired property in a single step. This approach builds on work in structured deep generative models (Kingma et al., 2014; Siddharth et al., 2017) , which aim to learn a disentangled representation that factors into observed properties we want to control for, and latent factors that account for the remaining features which are either hard to annotate or irrelevant to the properties we wish to optimize.

We derive a regularizer for supervised variational autoencoders which exploits property information that we provide as supervision, ensuring that produced molecules adhere to target properties they are conditioned on. We demonstrate the ability of our model to perform accurate conditional molecule generation and a sort of "style transfer" on molecules, where a latent representation for a single molecule can have its target properties perturbed independently of its learnt structural characteristics, allowing direct and efficient generation of candidates for local optimization of molecules.

2. BACKGROUND:

Molecule discovery tasks come in two flavors. Global optimization seeks to find molecules that have a particular target property. Local optimization starts from some initial molecule and searches for molecules which have a desired property while not straying too far from the prototype. There is some overlap in methods used in the two approaches.

2.1. DEEP GENERATIVE MODELS FOR MOLECULES:

Virtual screening methods start from a large database of possible molecules and retain the promising ones (Eckert & Bajorath, 2007) , as measured by some quality function f (•). Machine learning approaches expand on this by dynamically generating additional candidate molecules; Segler et al. (2017) uses a stacked LSTM to produce large numbers of novel molecules which have similar characteristics to an existing database.

For properties which are expensive to evaluate, generating large sets of candidate molecules is not particularly useful. More sample-efficient global search can be achieved using Bayesian optimization methods, which use a generative model with a latent space that functions as a continuous representation of molecules (Gómez-Bombarelli et al., 2016; Kusner et al., 2017) . Optimization is then carried out over this continuous representation space to find candidates which are expected to have the desired property. Local gradient-based search can also be applied on continuous latent spaces to optimize the latent representation with respect to a target property (Jin et al., 2018; Liu et al., 2018) .

A challenge for these latent variable models is to reliably produce valid molecules. Character variational autoencoders (CVAEs) (Gómez-Bombarelli et al., 2016) generate molecules one character at a time, and are prone to syntactic and semantic errors; the grammar-based variational autoencoder (GVAE) (Kusner et al., 2017) and syntax-directed variational autoencoder (SD-VAE) (Dai et al., 2018) instead operate in the space of context-free and attribute grammars, respectively, to ensure syntactic validity. Other work generative models that operates on graph representations (Simonovsky & Komodakis, 2018; De Cao & Kipf, 2018; Jin et al., 2018; You et al., 2018; Liu et al., 2018) , largely improving the ability to generate valid molecules.

Suppose we are given a training set of pairs D = {(x i , y i )}, i = 1, . . . , N , where x corresponds to molecules and y represents a value of some properties of the molecule x. Assume the molecules represent an i.i.d. sample from some unknown distribution p(x), which assigns high probability to molecules believed to be useful for a given task. Aside from Segler et al. (2017) , which has no latent space and thus directly trains via maximum likelihood, these latent variable models are trained by optimizing a standard ELBO objective for variational autoencoders (?). This entails learning a stochastic encoder q φ (z|x) which maps molecules into a latent space, and a stochastic decoder p θ (x|z) for reconstructing molecules, by maximizing

L(θ, φ) = N i=1 E q φ (zi|xi) [log p θ (x i |z i )] -D KL (q φ (z i |x i )||p(z i )) .

(1)

Notably, the objective is not a function of y: most existing generative models with latent variables do not perform direct conditional generation, and approaches for targeted molecule discovery are bolted on to the learnt model. Some, e.g. Kusner et al. (2017) , are trained in an "unsupervised" manner, agnostic to any property which later may need to be optimized. Others, e.g. Gómez-Bombarelli et al. (2016) ; Liu et al. (2018) , train the autoencoder jointly alongside a function to predict y from z, hoping to guide the latent space to be also good for predicting the desired property. A recent exception is Assouel et al. (2018) , which learns a deterministic autoencoder where the decoder takes the latent code and the desired property as input, using a mutual information term in training to steer the model towards generating molecules whose target properties match the input. 2018) instead learn generation models optimized towards specific metrics, such as drug-likeliness and solubility; the major downside is that these models must be retrained each time for a new property. In contrast, the autoencoder-based methods can be re-used to optimize towards any particular value of the property.

2.2. STYLE TRANSFER WITH SUPERVISED VAES:

While the latent representations learned through standard VAE models perform well on the task of molecule reconstruction they do not necessarily provide interpretable factorised representations. A disentangled representation gives us additional control on the molecule generation process, allowing us to modify a single property leaving the remaining unaffected (Bengio et al., 2013a) . In many cases important variation in the data is easy to annotate. For example in the case of molecule datasets we have access to different functional descriptors of the molecules obtained by chemoinformatics software such as RDKit (Landrum) . Particularly useful to us here are supervised methods for learning disentangled representations (Kingma et al., 2014; Siddharth et al., 2017) . These are distinct from unsupervised disentangling approaches such as InfoGAN (Chen et al., 2016) or β-VAE (Higgins et al., 2017) , which encourages the latent factor to learn a disentangled representation by modifying the objective to promote component independence.

We will learn representations that specifically disentangle molecular properties of interest which we may later want to modify. Kingma et al. (2014) demonstrates how disentangling can be used to take two MNIST images of different digits, written in different styles, and independently change the digit while holding the style constant. An analogous operation on molecules would involve holding the physical structure of a molecule (its "style") relatively fixed while modifying a salient property. Unlike (say) the style transfer example for the MNIST digits, the conditional distribution of molecules with a particular value of properties might be very diverse; for example, the QED score attempts to measure the drug-likeness of a molecule, and the set of molecules generated at high values of this score would hopefully have high probability on a large, varied set of molecules. An essential challenge here is that the property only provides a very weak signal as to the overall structure of the molecule. To account for this diversity, we model the conditional distribution with a latent variable z, such that p θ (x|y) = p θ (x|z, y)p(z)dz. Concretely, this involves fitting a joint generative model of the form p θ (x, y, z) = p θ (x|y, z)p(y)p(z), in which y and z are independent under the prior, and we assume a unit multivariate normal prior p(z). To infer the latent variable z we will use a variational distribution q φ (z|x), which takes the form of a multivariate normal distribution with parameters a nonlinear function of x, to approximate the true posterior p θ (z|x, y). This objective function

L ELBO (θ, φ) = N i=1 E q φ (zi|xi) [log p θ (x i |y i , z i )] -D KL (q φ (z i |x i )||p(z i )) (2)

corresponds to learning a supervised VAE (Kingma et al., 2014) , and represents a fairly naïve approach to modeling a conditional distribution.

3. CONDITIONAL GENERATION BY DISENTANGLING:

Maximizing this conditional ELBO in Eq (2) will likely yield good reconstructions of molecules from an embedding z (alongside the true property y), but for properties which only weakly inform the generative model there is nothing to enforce that the variable y actually directly has an effect on the generative process. Since the value y is something we know is a derived property of the molecule x, it is completely possible for all information about y to also be encoded in the representation z, in which case there is no guarantee that the learnt likelihood p θ (x|y, z) actually takes into account the value of y -in fact, we know it is possible to fit variational autoencoders where the decoder simply has the form p θ (x|z) -and we are relying on the utility of y in reconstructions to see any sort of disentangling effect.

3.1. CONSTRAINED ELBO:

In the case of conditional generation of molecules, we often have access to some oracle function f (possibly non-differentiable) which for any given x outputs a property estimate y, for instance, the chemoinformatics software RDKit (Landrum) . Since for conditional generation our ultimate goal is to generate a molecule x for any given target property y 0 , which then actually has f (x) = y 0 , we can reframe the problem by introducing hard constraints on the generated values, i.e. if restricting to values of y in the training set, max

θ,φ L ELBO (θ, φ) subject to E x∼p(x|yi) [I[f (x) = y i ]] = 1 decoder network latent space target property encoder network discrete sampling procedure h : generated molecule f (x) f ! (h) ⇡ E[f (x)] ⇡ y ? g ✓ (z, y ? ) y ? f ! (h) p ✓ (x|z, y ? )

Approximate target property:

Non-differentiable property estimation:

Figure 2 : Setting and modeling pipeline for conditional generation of molecules, with supervision provided via an external property prediction oracle. Red lines correspond to non-differentiable components, including both a potentially complex sampling process and the property prediction itself. The blue dashed line corresponds to the approximate property predictor, which aims to predict the expected value of the property from a continuous relaxation, marginalized over the sampling process.

for all i = 1, . . . , N . This is an unreasonably hard constraint, unlikely to be satisfied by any distribution other than one which simply places a point mass on the single training x i associated with y i , but we can relax it by considering that (unlike the molecular space x) the property space y is typically smooth, as many properties are continuous-valued and correspond to a human-interpretable scale. Following Ma et al. (2018) and Hu et al. (2017) , we reframe the constraint as a soft penalty on the ELBO,

EQUATION

so that they are consistent with the property prediction, i.e., as we have an oracle function f which enable us to access the property of any generated data, we can explicitly add a soft constraint to our loss function to provide explicit guidance for the generative model such that f (x) = y. This constraint is expected to hold for any pair (x, y) we may happen to come across, not just those in the training data. We also show optimizing the relaxed constraint is equivalent to maximizing mutual information with the target y i and generated molecule x; for details see appendix Section 6.1.

3.2. APPROXIMATING THE PROPERTY PREDICTOR:

Introducing the regularizer as in Eq. ( 3) implicitly guides the reconstruction to take into account the property information, such that the reconstructed data should exhibit properties which match the input properties it is conditioned on. However, existing implementations of f are often non-differentiable or CPU-bound, and x are discrete samples from a categorical distribution, all of which means the gradient of the regularizer can't flow back to the generator. This is outlined in Figure 2 . To enable the gradient based methods on GPUs during training and avoid discrete sampling, one approach would be to first fit a differentiable approximation to f , and then use either a Gumbel-softmax relaxation (Jang et al., 2016) or tricks like a "straight-through" estimator (Bengio et al., 2013b) as a continuous approximation for the discrete samples. Instead, we propose bypassing the discrete sampling step entirely and learning a function f ω that can map from a learned representation of the molecules directly to molecules property (Hu et al., 2017) .

To do this, we take as input the last hidden layer of the decoder network which parameterizes p θ (x|z, y), denoting this deterministic transformation as g θ (z, y). For the grammar VAE and the syntax-directed VAE, this last layer h = g θ (z, y) is the output of a recurrent layer that generates logits corresponding to unmasked and unnormalized log probabilities for each character at each position in the string; see Kusner et al. (2017) and Dai et al. (2018) for details on the implementation of the somewhat complex sampling process in the decoder. Ideally, f ω would estimate the property distribution obtained by marginalizing out the discrete sampling step, with

EQUATION

where we condition on z, and y 0 refers to an arbitrary input target property.

Assuming the approximation in Eq. ( 4), we have

E p θ (x|yi) f (x) -y i 2 = E p(z) E p θ (x|yi,z) f (x) -y i 2 ≈ E p(z) f ω (g θ (z, y i )) -y i 2 ,

an expectation over a real-valued variable which does not depend on any of the parameters we are estimating, meaning we can use a simple path estimate of the gradient with respect to θ, ω by exchanging the gradient with the expectation. We thus define a regularization term

L disent (θ, ω) = λ 1 2 N i=1 E p(z) f ω (g θ (z, y i )) -y i 2 (5)

which can be used as a drop-in replacement for the non-differentiable penalty term in Eq. ( 3), yielding a candidate objective function

EQUATION

3.3. LEARNING THE PROPERTY ESTIMATOR JOINTLY WITH GENERATIVE MODEL:

While one could imagine attempting to learn f ω jointly with φ, θ by direct optimization of Eq. ( 6), in practice this is very unstable, as values of g θ (z, y i ) early in training may correspond to very poor generated molecules xi which may not have properties at all similar to y i . This can be sidestepped by training the property estimator jointly as part of an extended generative model on [x, y].

We note that the property estimator f ω parameterizes a probability distribution p ω (f (x)|z, y 0 ), where x ∼ p θ (x|z, y 0 ) and f is the oracle function that f (x) = y. With a Gaussian distribution over the error, we can consider

EQUATION

for small, fixed λ 2 . Therefore, we propose defining a new ELBO based on a joint autoencoder for {f (x i ), y i )}, albeit with a factorization such that the input y i bypasses the encoder and is passed directly into the decoder, with a joint likelihood

EQUATION

This yields a joint ELBO for the training set of

EQUATION

)

Note that we can rewrite this ELBO as a function of the previous one, with

EQUATION

where we also see that L ELBO (ω, θ, φ) ≤ L ELBO (θ, φ), allowing us to define an objective

EQUATION

which is a lower bound on Eq. ( 6). Notice the two terms we have added to the original ELBO are quite similar, differing only in choice of distribution: for learning f ω , we wish to use values of z simulated form the approximate posterior q φ (z|x), whereas for enforcing a constraint across all possible generations we simulate z from the prior p(z).

3.4. GRADIENT ESTIMATION:

As the regularizer L disent (θ, ω) encourages disentangling by constraining the molecules generated from y i to have property y i no matter what value z takes, we found that it does not necessarily evaluate at meaningful values of z when sampled randomly from p(z). This roughly corresponds to the notion that not all combinations of "style" and property are physically attainable; ideally for style transfer we would like the generated molecule to stay "close" in structure to the original molecule that we intended to modify. When estimating (gradients of) the soft constraint term L disent (θ, ω), we found it advantageous to use samples of z which correspond to encodings of actual data points, as opposed to random samples from the prior. We approximate expectations with respect to p(x) by looking at the so-called marginal posterior; we note that

p(z) = p θ (z|x)p θ (x)dx ≈ 1 N j p θ (z|x j ) ≈ 1 N j q φ (z|x j ),

where the first approximation uses the empirical data distribution as an approximation to the model marginal p θ (x), and the second uses our variational posterior approximation q φ (z|x). We define this quantity as q(z) = 1 N j q φ (z|x j ), a mixture of Gaussians, which we can sample from by drawing random values from our dataset and then drawing from their encoding distributions.

When we use this in estimating gradients of the soft constraint, we can use samples from the same minibatch, exactly corresponding to a property transfer task. That is, for any particular y i in the dataset, we can estimate

E p(z) ∇ θ,ω f ω (g θ (z, y i )) -y i 2 ≈ E q(zj |xj ) ∇ θ,ω f ω (g θ (z, y i )) -y i 2 .

for any uniformly randomly sampled j = i. By sampling z j from q(z j |x j ) where j = i, we make sure that all the label information decoder is receiving comes from the actual y i that is feed to the decoder and z j does not include any information about label. This can be evaluated easily by simply evaluating the penalty term of Eq. (10) twice per minibatch; once as in Eq. ( 10), and once to approximate L disent (θ, ω) by permuting the properties in the minibatch to be assigned to incorrect molecules. We detail the training algorithm in Section 6.2 of the appendix.

4. EXPERIMENTS:

We experiment with the QM9 dataset (Ramakrishnan et al., 2014) , that contains 134k molecules with up to 9 heavy atoms, and the ZINC dataset (Sterling & Irwin, 2015) containing 250k druglike molecules. Our goal here is two-fold: we would like to understand (1) whether a supervised variational autoencoder is capable of learning suitable conditional distributions over molecules, and

(2) to what extent this task is assisted by the additional regularization term corresponding to the soft constraint.

We represent molecules using the one-hot encoding of their SMILES production rules (Kusner et al., 2017) and add a semantic constraint (Dai et al., 2018) on the decoder network to avoid generating syntactically correct but semantically invalid molecules. We use 80 production rules to describe molecules and set the maximum SMILES sequence length to 100 for the QM9 dataset and 278 for the Zinc dataset. We experiment with the logP property of the molecules (Wildman & Crippen, 1999). We use the same encoder and decoder network structure as Dai et al. (2018) with the only difference that our decoder takes as input the concatenation of y, z. We give the details of the architecture in the appendix section 6.2.

We evaluate the reconstruction accuracy and the quality of the molecules generated by our method, which we denote by CGD-VAE (conditional generation with disentangling) and compare against CVAE (Gómez-Bombarelli et al., 2016) , GVAE (Kusner et al., 2017) , and SD-VAE (Dai et al., 2018) . We explore its conditional generation performance in two settings: controlling only the property value and controlling both the property value and the molecule structure to what can be seen as property transfer. We took the results of CVAE, GVAE from the literature. For SD-VAE we used the authors code with the default values to generate results for QM9 since these were not available for QM9. We also implemented supervised VAE versions of SD-VAE which we denote Sup-VAE-X-GRU (X ∈ {1, 3}, denotes GRU layers) and which can do conditional generation. Before proceeding with the experiments we will give some additional details on how we do conditional generation from p θ (x|y 0 ) given the target property y 0 . Instead of marginalizing over the prior p(z), we mirror the approach taken during training and integrate over an approximation to the marginal inference distribution q φ (z) = 1 N N i=1 q φ (z|x i ) which better characterizes where the mass of the dataset is in the latent space. However, as N is large and we do not wish to keep the entire dataset available at test time, we approximate q φ (z) with an isotropic Gaussian distribution qσ (z) = N (z|0, σ 2 I). We estimate σ for each model by Monte Carlo samples from q φ (z). For the supervised VAE without the soft constraint regularizer this yields 0.053 for QM9 and 0.118 for ZINC. For our model with the soft constraint we get 0.0354 for QM9 and 0.096 for ZINC. We do conditional generation of x given y 0 by sampling from p θ (x|y 0 ) = qσ (z)p θ (x|z, y 0 )dz.

We evaluate reconstruction performance in terms of the correctly reconstructed molecules on test sets of size 10k for QM9 and 5k for ZINC, for the latter we used the default test set. We evaluate the generated molecules' quality by the percentage of valid, unique (i.e. percentage of unique molecules among the generated valid molecules) and novel (i.e. percentage of molecules never seen in the training set among the generated molecules) molecules. We estimate these quantities by sampling 10k (5K for ZINC) z from the qσ (z) and coupling each one of them with a logP value, y, randomly selected from the test set, and we subsequently decode the z, y concatenation. We can see that our model has a better reconstruction performance compared to the baselines while in some cases generating slightly less valid molecules table 1. In terms of the three quality measures achieves an excellent performance across all three metrics being always one of the two best performing methods for any metric.

To visualise how the conditional generation operates we randomly sample from the test set some molecule and obtain its property value y 0 . We then draw 50 random samples z i from qσ (z) and decode the [z i , y 0 ] vectors. Among the generated valid molecules we compute the percentage of those that have a property value y i that is within a 15% range from the y 0 property value. In Figure 3 we present the molecules obtained for a test molecule that had a logP of -0.5759. Out of the 50 generated molecules 46 were valid of which we give the five that were within a 15% range from the y 0 value in Figure 3 . As we can see we get molecules that are structurally very different from the original one yet they have similar logP value.

To quantify the quality of the conditional generations we measure the correlation between the property value we obtain by the conditional generation and the property value on which we conditioned the generation. We randomly sample 1000 y values from the test set and 1000 z values from the approximate learned prior qσ (z). We decode each pair, obtain x ∼ p θ (x|y, z), and then measure the correlation of the original y with the ŷ of generated x. In Table 2 , we give the correlation estimates for our method and the Sup-VAE baselines. As we can see our method has a considerably higher correlation score between the input and the obtained property than Sup-VAE. Conditional generation seems considerably harder for the ZINC dataset for all methods.

To visualise the style transfer behavior of our model we randomly sample two molecules x A , x B from the test set. We then sample z A from the learned posterior q φ (z|x A ). We subsequently decode [z A , y B ], y B is the property of x B , and get a new molecule xAB . Ideally, the obtained molecule xAB should have a property value (logP) close to the target y B and be similar to x A . In Figure 4 we give one such example. To put the results into context in Figure 8 in appendix, we give the results of a virtual screening method, where we select from the full dataset five molecules which are structurally To quantify the style transfer performance we proceed in exactly the same manner as we did to quantify the conditional generation performance. However, now instead of sampling z from the approximate learned prior, qσ (z), we first sample some x from the test set and then we sample z from the learned posterior q(z|x). The results are in the second column of Table 2 . As we can see the correlation values are now lower than the ones we obtained in the simple conditional generation case. This can be explained by the fact that now we are forcing a specific combination of structure, z comes from a real molecule, and property, which might simply be physically infeasible since the molecule space is discrete and not all combinations are possible. In addition, as it was the case for the conditional generation, style transfer is considerably more difficult for the ZINC dataset.

We further explore the style transfer and visualize how our model covers the combined space of molecule structure and properties. We sample nine molecules from the QM9 test set, and get their z encodings. For each such encoding we decode the vectors [z, y], y ∈ [-4.9, 4.9], with the y (logP) interval sampled at 11 points. We give in Figure 5 the resulting valid molecules, each column there corresponds to one of the nine original molecules, the ones surrounded by dotted rectangle, and their decodings with different logP values.For each original molecule we give the generated molecules ordered along the y axis according to the y property that they actually exhibit. The x-axis does not provide an ordering of the original molecules according to z, in fact we have ordered the original molecules by their y property. As we can see not all (z, y) combinations produce a result. These holes can be explained either by the physical infeasibility of the combination and/or a limitation of the learned model.

We can use conditional generation to control in a fine manner the value of the desired property, to what can be seen as direct property optimization. We visualise the level of control we have on an experiment with a single molecule (with logP is -1.137), which we randomly sample from the test set. We obtain its z encoding and perform generations with increased logP taking values in 1000 point grid in [-1.137, 4.9] . We then decode [z, y i ] and compute the logP value of the generated molecules. Among the 1000 generated molecules only 19 are unique. We get an increase of logP of a very discrete nature, Figure 6 . As already discussed not all combinations of structure and properties are possible. The generated molecules themselves are shown in the supplemental material. Table 2 : Correlation between the desired input property and the obtained property . z ∼ qσ (z) corresponding to conditional generation), and x, z ∼ q(z|x) to property transfer case. 

4.1. CONDITIONAL LSTM BASELINE:

Finally, we consider a variant of the stacked LSTM model of Segler et al. (2017) , with no latent space, where the model is modified to take a target logP value as an additional input at each generation step. This model forms a very strong baseline for many distribution matching tasks (Liu et al., 2018; ?) , though as best we are aware this has never been used directly for conditional generation given a target property. We use a modification of the implementation provided by (?) with three layers and default settings, and fit the model by maximum likelihood training on p θ (x|y) = T t=1 p θ (x t |x 1:t-1 , y). Training on the ZINC dataset, we find the generated molecules from this model have a very high correlation 0.975 with the target logP value, greatly outperforming any of the latent variable models we consider. This suggests that such a model would be very useful for generating candidates globally, but as the model has no latent variable it is not amenable to style transfer. We observe this in Figure 7 , which samples 100 candidate molecules from both the stacked LSTM model and for CGD-VAE-3-GRU, conditioning on the property of one randomly-chosen test set example, while computing the Tanimoto similarity (computed using Morgan fingerprints of radius 2) to a second randomly-chosen test set example, across 200 pairs. The VAE has higher Tanimoto similarities as it can condition on the latent variable of the target molecule, representing a trade-off against the better adherence to the target property value of the unconditioned LSTM.

5. CONCLUSION:

We presented a single step approach for the conditional generation of molecules with desired properties. Our model allows also to condition generation on a prototype molecule with a desired high-level structure. This work thus directly inverts the traditional relationship between molecules and their properties. We found that training the deep generative models conditional on target properties, following a supervised VAE approach, does not appreciably harm the quality of the unconditional generative model as measured by validity, novelty, and uniqueness of samples. Furthermore, we see that the additional act of regularizing the output using an approximate property predictor helps improve both reconstruction accuracy and property correlations in most combinations of tasks and datasets, particularly for the smaller QM9 dataset and for smaller models with fewer RNN layers. We also note that although none of the deep latent variable models are competitive with an LSTM baseline when purely considering generation conditioned on a target property value, the low Tanimoto similarity between randomly sampled candidates and an arbitrary style transfer target makes clear that such a model is not suitable for targeted generation of candidates which are close in structure to a particular prototype.

In future work, we want to explore how to further improve the correlation between the desired input properties to the decoder, and the properties of the generated molecules. Moreover, we want also to condition on multiple properties; while this is in principle possible in our framework, we do not explore it empirically here. Modifying a single property while constraining the remaining to be close to the original can further aggravate the infeasibility problem, as not all combinations of molecular properties may even be feasible, perhaps requiring learning a dependency structure between multiple properties.

6. APPENDIX:



6.1. THE REGULARISER AND ITS RELATION TO THE MUTUAL INFORMATION MAXIMIZATION:

The soft constrain in the loss 3, in fact , is equivalent to a simple maximizing mutual information formulation between generated molecules x and the target property y provided to the generator. Assume the true conditional distribution is p(y|x):

EQUATION

We do not know the true distribution p(y|x), however, the RDKit provides an estimation p(y|x) of the distribution assuming a Gaussian distribution over error:

EQUATION

)

where f is the molecule property estimator, i.e., RDKit. We have: 

I(y; x) = H(y) + E x∼p θ (

As H(y) is constant, minimizing E x∼p θ (x|yi) f (x)-y i 2 is equivalent to maximizing I(y; x) under the assumption that p(y|x) is close to p(y|x)

6.2. ARCHITECTURE AND TRAINING PROCEDURE DESCRIPTION:

We use the same encoder and decoder network structure as Dai et al. (2018) with the only difference that our decoder takes as input the concatenation of y, z. As GRU layers become computationally expensive when the sequences length increase, we also examined the model using less layer GRU. To be precise, the decoder in Dai et al. (2018) takes the from of a dense hidden layer with ReLU activation followed by three layers GRU Chung et al. (2014) . We tried two different settings of decoder: in the first setting, we feed the concatenation of y, z to dense layer then apply one layer GRU, in the second setting, to enhance the effect of y in the decoder, we feed y not only to the dense layer but also to each layer of GRU. Furthermore, we set the dimension of the latent representation to 56. For the oracle function estimator f w , we use the same network architecture as the encoder (there is no parameter sharing) and we add one more fully connected layer followed by a Tanh transformation.

To speed up convergence, we initialize the f w from a pre-training, where we train f w on the welltrained (maximum 500 epochs with early stopping) supervised VAE's decoder output to predict the molecules property value. We also initialize the parameters of the encoder/decoder networks with the partially trained supervised VAE model (after 40 epochs for QM9, 100 epochs for ZINC). We do not update ω and φ, θ simultaneously, instead we do an alternate optimization. We update ω continuously for five epochs while holding φ, θ and do the same for updating φ, θ. We set the hyper-parameter value λ 1 to 50 and λ 2 to 1. The mini-batch size is set to 300 for QM9 and 100 for ZINC. We use ADAM optimizer with learning rate 0.0001 and pytorch lr-schedular on the validation loss. The general training algorithm is described in below algorithm block1. In our experiment, to train f ω , we skipped the second term in step 7, which means we only train f ω on the training data but not the newly generated molecules obtained by permuting the property. The reason for this is that, during the training, we found that it is easy for the model to learn to reconstruct but hard to conditionally generate the molecules with given properties while we have no guidance of what the molecules should look like. Further more, some combination of z and y are physically not feasible. In this case, when the conditional generation is not good enough yet during the training, we end up fitting f ω on the miss represented molecules representations and it makes the optimization harder. θ j = θ j-1 -γ(-∇ θ L ELBO (θ, φ, ω) + λ1 2 (xi,yi)∈D * E q(z|xi) ∇ θ f ω (g θ (z, y i )) -y i 2 )

6:

φ j = φ j-1 -γ(-∇ φ L ELBO (θ, φ, ω))

7:

w j = w j-1 -γ( λ2 2 (xi,yi)∈D E q(z|xi) ∇ ω f ω (g θ (z, y i )) -y i 2 + λ1 2 (xi,yi)∈D * E q(z|xi) ∇ ω f ω (g θ (z, y i )) -y i 2 )

6.3. USE A FIXED PRE-TRAINED PROPERTY PREDICTION FUNCTION:

We also investigate the case where we train the property prediction function f ω on the well trained supervised VAE output and keep it fixed during the training of the main model. We give the performance in tables 3, 4. Using a fixed f ω , in terms of reconstruction and generation performance, delivers mixed results 3. However, in terms of conditional generation 4, it does perform better than the baselines but worse compared to the case where we also update f ω during the learning. Table 4 : Correlation between the desired input property and the obtained property when z is sampled from the approximate learned prior, z ∼ qσ (z), (conditional generation), and when z is sampled from the learned posterior given some x, z ∼ q(z|x), (property transfer).

investigate if this findings holds for other baseline models or not. We explored the behavior of baselines when the z is sampled from the approximate marginal posterior, we observe that for CVAE and GVAE the validity did not change (as the q(z) and p(z) are essentially identical); for the SD-VAE the validity increases (Table 5 ). 

:



Figures and Tables:

Figure 1: A demonstration of style transfer

Figure 3: Conditional generation given the desired logP=-0.5759, row molecules have a logP within a 15% range of the desired one.Figure 4: Property transfer

Figure 3: Conditional generation given the desired logP=-0.5759, row molecules have a logP within a 15% range of the desired one.Figure 4: Property transfer

Figure 5: Style transfer. The z of the nine real molecules placed in the x-axis is combined with 11 y property values, sampled in [-4.9, 4.9], the resulting pair is decoded to a molecule.

Figure 6: Property optimization. Given a start molecule (red), we combine its z with 1000 logP values and decode (blue)

Training algorithm 1: Initialize p θ (x|z, y), q φ (z|x), f w 2: for i=1,2, . . . , N (maximum epoch number) do 3:for j = 1,2, . . . , L, sample a minibatch D = (X, Y) = {x i , y i )} M i=1 of M samples do 4:randomly permute the property set Y to obtain Y * and define a label permuted mini batch D * = (X, Y * ) = {x i , y * i

Figure11is the visualization of the generated molecules from property optimization task, given in figure6section 4. The molecules are generated by increasing the logP of a given molecule, i.e. we hold z fixed and increase the y value. Among the all generated molecules, 19 of them are unique.

Figure 11: Molecules generated with an increasing logP

Reconstruction performance and generation quality (Valid, Unique, Novel).

Table content: <table><tr><td>QM9</td><td>ZINC</td></tr></table>

Sup-VAE-1-GRU /Sup-VAE-3-GRU: supervised version of SD-VAE model where y is been feed to only the first layer/all layer decoder. CGD-VAE: our model, conditional generation with disentangling. The result for CVAE and GVAE are taken from the literature, "-" refers that those measures are not reported. The rest of baseline result is obtained by rerunning (SD-VAE) and editing the original code (Sup-VAE).

Table content: <table><tr><td/><td/><td>QM9</td><td/><td/><td/><td>ZINC</td><td/><td/></tr><tr><td>Model</td><td colspan="8">Reconstruction % Valid% Unique % Novel % Reconstruction % Valid% Unique % Novel %</td></tr><tr><td>CVAE</td><td>3.61</td><td>10.30</td><td>-</td><td>90.00</td><td>44.60</td><td>0.70</td><td>-</td><td>100</td></tr><tr><td>GVAE</td><td>96.00</td><td>60.20</td><td>-</td><td>80.90</td><td>53.70</td><td>7.20</td><td>-</td><td>100</td></tr><tr><td>SD-VAE</td><td>97.84</td><td>98.40</td><td>99.28</td><td>91.97</td><td>76.20</td><td>43.50</td><td>-</td><td>-</td></tr><tr><td>Sup-VAE-1-GRU</td><td>97.53</td><td>93.66</td><td>91.30</td><td>92.05</td><td>74.12</td><td>32.84</td><td>95.61</td><td>100</td></tr><tr><td>CGD-VAE-1-GRU</td><td>98.96</td><td>95.03</td><td>92.77</td><td>89.74</td><td>67.46</td><td>17.38</td><td>82.44</td><td>100</td></tr><tr><td>Sup-VAE-3-GRU</td><td>97.81</td><td>97.9</td><td>95.09</td><td>89.47</td><td>82.40</td><td>36.16</td><td>86.26</td><td>100</td></tr><tr><td>CGD-VAE-3-GRU</td><td>96.9</td><td>93.8</td><td>98.29</td><td>89.55</td><td>81.80</td><td>37.78</td><td>98.70</td><td>100</td></tr><tr><td/><td/><td/><td colspan="3">Model z ∼ qσ(z) z ∼ q ( z|x)</td><td/><td/><td/></tr><tr><td/><td/><td colspan="2">Sup-VAE-1-GRU</td><td>0.5420</td><td>0.2526</td><td/><td/><td/></tr><tr><td/><td>QM9</td><td colspan="2">CGD-VAE-1-GRU Sup-VAE-3-GRU</td><td>0.6331 0.6958</td><td>0.3835 0.4204</td><td/><td/><td/></tr><tr><td/><td/><td colspan="2">CGD-VAE-3-GRU</td><td>0.6665</td><td>0.4507</td><td/><td/><td/></tr><tr><td/><td/><td colspan="2">Sup-VAE-1-GRU</td><td>0.2301</td><td>0.0481</td><td/><td/><td/></tr><tr><td/><td>ZINC</td><td colspan="2">CGD-VAE-1-GRU Sup-VAE-3-GRU</td><td>0.2981 0.3638</td><td>0.0866 0.1818</td><td/><td/><td/></tr><tr><td/><td/><td colspan="2">CGD-VAE-3-GRU</td><td>0.3765</td><td>0.1310</td><td/><td/><td/></tr></table>

Baseline model performance on ZINC 6.7 MOLECULE PROPERTY OPTIMIZATION

Table content: <table><tr><td/><td>SDVAE</td><td/><td/></tr><tr><td/><td colspan="3">Valid% Unique % Novel %</td></tr><tr><td>z ∼ qσ(z)</td><td>56.05</td><td>80.69</td><td>100</td></tr><tr><td>z ∼ p(z)</td><td>43.50</td><td>82.26</td><td>100</td></tr></table>

Peer review (iclr2020):
The authors introduce a variational autoencoder for conditional generation of molecules. The model is borrowed from text-based style transfer, applied here on sequence (SMILES) representation of molecules rather than viewing molecules as graphs (as more recent approaches). From a modeling point of view, the main new part is an additional regularizer whose role is to 1) ensure that the property used as input during generation matches the property derived from the generated molecule, and 2) to dissociate the latent molecule representation in the autoencoder (loosely speaking, its overall structure) from the property being controlled. This regularizer is just a squared difference between predicted and actual properties, averaged over independent samples from latent states and properties which are parametrically mapped to predicted properties via the decoder state (so as to be able to backprop). The resulting ELBO criterion for the VAE, together with the regularizer, could be in principle directly used to train all the model parameters (encoder, decoder, property mapping) but apparently this criterion does not result in a *stable* estimation procedure. The model is then amended in various ways. The first change consists of introducing a property prediction part directly into ELBO so that the posterior over latent states will also be affected by property prediction (previously property prediction would only appear in the additional regularizer). Beyond stability that the authors cite as the reason, this is likely also needed because in eq (2) the posterior over latent states given the molecule does not depend on the property. The posterior is derived from a v-structure where latent states and properties appear as independent parents of the molecule. Since for the training cases the property is fixed, this would be fine with an arbitrarily complex posterior encoding. However, when the encoder is defined parametrically, this dependence may be needed, and may relate to the stability issues. It seems a bit roundabout way of putting the dependence back in, if so. The authors should comment on this further. The other simplification that is introduced is that the prior over the latent states in the regularizer is modified to be the average posterior of training encodings, and later further collapsed (for computational reasons) to an adaptively estimated simple Gaussian with variance that matches the average posterior variance. It*s a bit unsatisfying that the authors only use and experiment with a simple logP as the property to control in the model. This is not really a particularly relevant property to control. The approach should in principle generalize to other properties but the many adjustments needed to make it work makes this a little questionable. Also, the results are a bit limited (only logP) and not encouraging (see comments below). - how is validity defined? SMILES syntax or validity as a molecule? - is table 1 reconstruction really autoencoder reproduction percentage? Is the correct property for the molecule being reconstructed offered as an additional input to the authors* model? isn*t this a bit unfair? - the correlation between property given as input and the property derived from the generated molecule does not seem like a good metric. Wouldn*t a model that overfits (essentially spitting out the nearest property neighbor from the training set) do particularly well according to this metric? - it seems Figure 6 can be interpreted as also indicating some degree of overfitting? - the stacked LSTM molecular generator baseline that takes the property as input seems to do much better in terms of maintaining the property in its generated molecules and it also realizes a more diverse collection of molecules. It*s true that it doesn*t support incremental molecular manipulation directly, eg, if one is keen on keeping the modified molecule close to a starting point. But it is so much better at maintaining the input property in its diverse realizations that perhaps one should instead invest in a tailored sampling procedure to obtain realizations close to a chosen reference. The literature references seem to omit all molecular generation work since 2018