Conditional generation of molecules from disentangled representations

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 ( Gó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 ) ( Gó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 = { ( xi , yi ) } , 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 { Eqφ ( zi|xi ) [ log pθ ( xi|zi ) ] −DKL ( qφ ( zi|xi ) ||p ( zi ) ) } . ( 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 . Gó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 . Guimaraes et al . ( 2017 ) ; De Cao & Kipf ( 2018 ) ; You et al . ( 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 . Disentangling the latent code z from the property y enables style transfer . This is done by taking an initial x , computing the posterior over the latent variable z , and then generating a new x′ with the property modified to have a target value y′ , with pθ ( x′|y′ , x ) =∫ pθ ( x ′|y′ , z ) pθ ( z|x ) 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 LELBO ( θ , φ ) = N∑ i=1 { Eqφ ( zi|xi ) [ log pθ ( xi|yi , zi ) ] −DKL ( qφ ( zi|xi ) ||p ( zi ) ) } ( 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 .

Summary: This paper proposes a VAE-based conditional molecular graph generation model. For that purpose, the disentanglement approach is adopted in this paper: learn to separate property information from the structure representation of a molecular graph. The authors use the supervised VAE objective since the KL regularizer in the objective has reportedly disentanglement-promoting effect. The final objective function is a standard VAE ELBO plus a penalty term of the property value prediction error. Non-differentiable property estimation is conducted via stochastic sampling expectation with the help of an external program (RDKit).