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: 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).