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):
# Summary The paper considers the problem of generating molecules with desired properties using a variant of supervised variational auto-encoders. The key novelty is the idea to learn a representation with disentangled molecular properties which can be modified during generation of novel molecules. To account for the diversity of molecules in a particular target class the conditional probability of a molecule x given a target property y is modelled with a latent variable z, i.e., p(x | y) = int p(x | z, y) p(z) dz. The *style transfer* is achieved by taking an initial molecule x and computing the posterior of novel molecule x* with modified property y* via p(x* | y*, x) = int p(x* | y*, z) p(z | x) dz. The latent representation is learned using a supervised auto-encoder (Kingma et al., 2014), formally specified in Eq. (2). The optimization problem is then further extended by imposing a soft constraint that the molecules from a particular class exhibit a certain property (Eq. 3). The reason for this constraint is to enforce that the conditional probability p(x | y, z) takes into account the information in label y. This part might not be sufficiently well-motivated and would benefit from strengthening the argument for the property predictor and soft constraint (would an illustration be possible here?). In particular, why should it not be possible to encode (and keep fixed during training) the information present in y into the sufficient statistics of p(x | y, z)? The soft constraint involves an oracle function (see Eq. 3), which evaluates the designed molecules. Section 3.2 deals with the approximation of the oracle via a property predictor function to side-step the fact that the oracle is not necessarily differentiable or CPU-bound. Section 3.3 and 3.4 deal with joint optimization of the supervised variational auto-encoder and gradient estimates. The approach is evaluated on two datasets, relative to state-of-the-art baselines based on variational auto-encoders and deep learning. The main goals of the experiments are to establish that the approach can learn useful disentangled conditional distributions over molecules and to assess the extent of improvement in the data generation process as a result of using the soft constraint term. The focus of the experiment is on generating molecules with certain range of logP values. The data generating process is simulated/evaluated while controlling for the logP value range and molecular structure. The experiments provide an objective assessment of the merits of the approach and indicate clear advantages over prior work. # Recommendation The paper is very well written, easy to follow, and properly structured. The work is novel and the idea to disentangle molecular properties from the target property is quite interesting. The experiments are detailed and compare to baselines based on variational auto-encoders and deep learning. The part that could be improved is the choice of the target property. In particular, it would be great to evaluate the approach on the problem of finding molecules binding to a certain protein site. Such a problem is likely to exhibit scaffold hops and activity cliffs which make drug design problems very difficult (e.g., see [6-8]). # Related work (additional references) - While mapping of the discrete space of molecules to a continuous latent space (in active learning approaches combined with VAE) allows for gradient computation and the use of active learning/optimization techniques, the very difficult problem of finding latent space pre-images persists (in general, the problem should be in the NP class). The latter refers to finding a molecule that maps to a fixed point in the latent space. There have, however, been some efficient approximations in special cases [1-3]. - In [4] and [5], an approach for generation of molecules with desired properties has been pursued with posterior sampler based on a conditional exponential family model. The representation is based on a tuple kernel mapping mapping (x, y) to some reproducing kernel Hilbert space (without disentanglement). The approach provides a consistent data generation process and a mean to deal with the fact that designs are not independent and identically distributed (e.g., Eq. 2 assumes that examples are IID). # Black-box target property - The logP property might not be a good proxy for the effectiveness in generating molecules with desired binding activities (to some protein site) because (to the best of my knowledge) it does not exhibit the *activity cliff property* characteristic to drug design. This refers to the property that a small change in molecular structure results in a completely different activity level (e.g., see [6-8] and [4]). In this sense, the *style transfer* might be beneficial for finding molecules binding to a certain protein site that are structurally similar to some available molecules with bad activity levels (for which a synthesis path might be known or easy to derive). - An actual black-box property (realized via a docking program) that is expensive to evaluate served as a target property in [5]. That docking program (supplied with suitable binding/docking constraints) can provide a nice proxy for the actual generation of molecules. # References [1] G. H. Bakir, J. Weston, and B. Schoelkopf. Learning to find pre-images, NIPS 2004. [2] C. Cortes, M. Mohri, and J. Weston. A general regression technique for learning transductions, ICML 2005. [3] S. Giguere, A. Rolland, F. Laviolette, and M. Marchand. On the string kernel pre-image problem with applications in drug discovery, ICML 2015. [4] D. Oglic, R. Garnett, and T. Gaertner. Active search in intensionally specified structured spaces, AAAI 2017. [5] D. Oglic, S. Oatley, S. Macdonald, T. Mcinally, R. Garnett, J. Hirst, and T. Gaertner. Active search for computer-aided drug design, Molecular Informatics 2018. [6] G. Schneider and U. Fechner. Computer-based de novo design of drug-like molecules, Nature Reviews Drug Discovery, 2005. [7] P. Schneider and G. Schneider. De novo design at the edge of chaos. Journal of Medicinal Chemistry, 2016. [8] J. Scannell, A. Blanckley, H. Boldon, and B. Warrington. Diagnosing the decline in pharmaceutical R&D efficiency. Nature Reviews Drug Discovery, 2012.