Improving black-box optimization in VAE latent space using decoder uncertainty

1 Introduction . We consider the task of optimizing an expensive black-box objective function taking inputs in a high-dimensional discrete space . This could be for example finding new molecules for drug design , or automatically generating a computer program that matches a desired output . Solving this task directly in the original space ( e.g. , with discrete local search methods such as genetic algorithms ) may be challenging given the complex structure and high dimensionality of the data . Recently , Variational autoencoders ( VAEs ) [ 1 , 2 ] have been successfully leveraged to model a wide range of discrete data modalities — from natural language [ 3 ] , to arithmetic expressions [ 4 ] , computer programs [ 5 ] or molecules [ 6 ] . By learning a lower-dimensional continuous representation of objects in their latent space , VAEs allow to transform the original discrete optimization problem into a simpler continuous optimization one in latent space . For example , this can be achieved via Bayesian Optimization in the latent space , or via gradient ascent with a jointly-trained neural network predicting the black box 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . property from the latent space representation [ 6 , 7 ] . Initial methods in this area have suffered from the fact that the search in latent space may explore areas for which no data was available at train time , and therefore where the decoder network of the VAE will be unreliable [ 8 ] : seemingly good candidate points in latent space may be decoded into objects that are invalid , unrealistic or low quality . While several methods have been introduced to promote validity of decoded objects ( §2.1 ) , they either focus on modifying the generative model learning procedure or adapting the decoder architecture to satisfy the syntactic requirements of the data modality of interest . We propose instead to quantify and leverage the uncertainty of the decoder network to guide the optimization process in latent space ( Fig . 1 ) . This approach does not require any change to the model training nor architecture , and can easily be integrated within several optimization frameworks . It results in a better trade-off between the values of the black-box objective and the validity of the newly generated objects , sometimes improving both simultaneously . To be effective , this method requires robust estimates of model uncertainty for high dimensional structured data . Existing methods for uncertainty estimation in this domain often rely on heuristics or make independence assumptions to make computations tractable ( §2.2 ) . We demonstrate that such assumptions are not appropriate in our setting , and propose new methods for uncertainty estimation in high dimensional structured data instead . Our contributions are as follows : • We introduce an algorithm to quantify the uncertainty of high-dimensional discrete data , and use it to estimate the uncertainty of the decoder ( §3 ) ; • We show how the uncertainty of the decoder can be incorporated across several optimization frameworks , including gradient ascent and Bayesian optimization ( §4 ) ; • We illustrate our approach in a digit generation setting — a simple setup to provide intuition for the method — then quantify its benefits in the more complex tasks of arithmetic expressions approximation and molecule generation , covering a diverse set of decoder architectures across experiments ( Convolutional , Recurrent and Graph Neural Networks ) ( §5 ) . 2 Background . 2.1 Optimization of high-dimensional discrete objects with generative models . Focusing on molecular generation , Gómez-Bombarelli et al . [ 6 ] propose to train a VAE to learn a distribution over the so-called SMILES representation of molecules ( ie . linear sequences of characters ) [ 9 ] , and subsequently perform the optimization in the latent space . Since the SMILES representation follows strict syntactic requirements that are not explicitly enforced by the generative model , promising points in latent space may be decoded into invalid molecules . To improve the validity of decoded sequences , Kusner et al . [ 4 ] and Dai et al . [ 5 ] develop task-specific grammar rules into the VAE decoder , focusing on use cases in molecular and computer program generation . However , crafting the corresponding rules requires domain-specific knowledge , needs to be designed from scratch for each new task , and may not be straightforward to elicit in the first place ( e.g. , digit generation example in 5.1 ) . Griffiths and Hernández-Lobato [ 10 ] and Liu et al . [ 11 ] propose instead to formulate the problem as a constrained Bayesian Optimization and chance-constrained optimization task respectively to simultaneously optimize the target property as well as the probability to generate valid sequences . These two approaches require nonetheless access to a function that quantifies the validity or realism of objects in the training data , which is not readily available in many practical applications . A different line of research has focused on representing high-dimensional structured objects as graphs instead [ 12 , 13 ] . The Junction Tree VAE ( JT-VAE ) [ 14 ] generates systematically valid molecular graphs , by first generating a tree-structured scaffold over a finite set of molecular clusters , and then assembling these clusters back into molecules with a message passing network . The MolDQN [ 15 ] casts the optimization problem as a reinforcement learning task ( double Q-learning ) , which allows in turn to more naturally extend to simultaneous optimization of different objectives . GraphAF [ 16 ] combines the strengths of autoregressive and flow-based approaches to efficiently generate realistic and valid molecular graphs . Lastly , Tripp et al . [ 17 ] show that the black-box optimization performance can be further enhanced by iteratively retraining the generative model on the points selected during optimization , with weights that depend on their objective function value . Our approach deviates from all the above in that it is is representation-agnostic ( works with sequences or graphs ) , does not require domain-knowledge to craft custom rules or constraints , does not change the model architecture nor the learning procedure and can be combined with several of these approaches to reach even stronger optimization performance ( see § 5 ) . 2.2 Quantifying model uncertainty . Adopting a Bayesian viewpoint , the overall uncertainty of a model in a given region of the input space can be broken down into two types of uncertainty [ 18 ] : • Epistemic uncertainty : Uncertainty due to lack of knowledge about that particular region of the input space — the posterior predictive distribution is broad in that region due to lack of information that can be reduced by collecting more data ; • Aleatoric uncertainty : Uncertainty due to inherent stochasticity/noise in the observations in that region — collecting additional data would not further reduce that uncertainty . We denote input points as x , outputs as y and the training data as D. The total uncertainty U of a model at an input point x is typically measured by the predictive entropy , ie . the entropy of the predictive posterior distribution P py|x , Dq : Upxq “ HpP py|x , Dqq “ ÿ y ´P py|x , Dq logP py|x , Dqdy . ( 1 ) Denoting P pθ|Dq the posterior distribution over model parameters θ , we can further decompose the predictive entropy U as the sum of two terms : Upxq “ pHpP py|x , Dqq ´ EP pθ|DqpHpP py|x , θqqq looooooooooooooooooooooooomooooooooooooooooooooooooon Mutual Information M ` EP pθ|DqpHpP py|x , θqqq looooooooooooomooooooooooooon Expected entropy E ( 2 ) The first term — the Mutual Information M between model parameters θ and the prediction y — is a measure of epistemic uncertainty , as it quantifies the magnitude of the change in model parameters that would result from observing y . If the model is uncertain about its prediction for y , the change in model coefficients from observing y should be high . Conversely , if the model is confident about its prediction for y , model parameters will not vary from observing y : Mpxq “ HpP py|x , Dqq ´ EP pθ|DqpHpP py|x , θqqq . ( 3 ) The second term — the Expected Entropy E — is a measure of the residual uncertainty , ie . the aleatoric uncertainty : Epxq “ EP pθ|DqpHpP py|x , θqqq . ( 4 ) In high dimensional settings , an exact estimation of these different quantities is not tractable , therefore , several approximations and heuristics have been introduced . The softmax variance , i.e . the variance of predictions across model parameters , has been shown to approximate epistemic uncertainty well in certain settings [ 19–21 ] . In the context of sequential data , the inherent structure in the data generating process often introduces strong dependencies between the output dimensions , e.g. , the tokens in a generated sentence . In cases where there exist weak correlations between tokens , quantifying the different types of uncertainties above can be made tractable by ignoring these dependencies [ 22 ] , in which case the predictive entropy for a sequence y “ py1 , y2 , ... , yLq may be approximated as the sum of token-level predictive entropies over the L tokens : Upxq “ L ÿ l “ 1 EP py|x , DqrlogP pyl|x , ykăl , Dqs « L ÿ l “ 1 EP pyl|x , ykăl , DqrlogP pyl|x , ykăl , Dqs . ( 5 ) Unlike the standard expectation definition which integrates over all y , here only yl is integrated over , and we condition on ykăl , which are obtained from a sample from P py|x , Dq . The process is repeated for several of these samples and an average is finally computed to reduce variance . While the above has been shown to work well in certain experiments [ 22 ] , valuable information is being discarded when we ignore dependencies across tokens . These dependencies are likely to be informative in applications such as the ones considered in §5.3.1 . Naive Monte Carlo-based approximations are expected to perform poorly in high dimensions since the majority of samples will be in regions with negligible contribution to the sum ( Appendix A ) . In an active learning setting , Kirsch et al . [ 23 ] derive an importance sampling-based estimator of total uncertainty to mitigate these issues . In Natural Language Processing , alternative approaches using domain-specific metrics such as the BLEU score have been suggested [ 24 ] , but these are difficult to extend to other domains . 3 Importance sampling estimator . We consider discrete output points y that belong to a high-dimensional structured object space S ( e.g. , long sequences , large graphs ) , of cardinality |S| . An exact estimation of the Mutual Information between outcomes y and model parameters θ ( equation 3 ) is impractical because the expectation involves a sum over exponentially many possible outcomes for y P S. In lieu of the heuristics previously discussed , we obtain a principled approximation to the Mutual Information via Monte Carlo estimation using importance sampling , with an adequately chosen importance distribution . We denote qpθq « P pθ|Dq the learnt approximation to the posterior , and assume we can approximate expectations over model parameters by sampling M independent samples from qpθq . We can then re-write the Mutual Information M in equation 3 as follows : Mpxq « ´ |S| ÿ s “ 1 ps log ps ` 1 M M ÿ m “ 1 |S| ÿ s “ 1 ps , m log ps , m “ |S| ÿ s “ 1 « 1 M M ÿ m “ 1 ps , m log ps , m ´ ps log ps ff loooooooooooooooooooooomoooooooooooooooooooooon hpysq ( 6 ) where ps and ps , m are shorthands , respectively , for P py “ ys|x , Dq — the posterior predictive distribution — and P py “ ys|x , θ “ θmq— the probability of a given output ys P S given x and a sample θm from the approximate posterior distribution over model parameters qpθq . We can then obtain a tractable approximation to equation 6 via importance sampling : Mpxq “ |S| ÿ s “ 1 hpysq ¨ 1 p̄pysq ¨ p̄pysq “ Ep̄ „ hpyq ¨ 1 p̄pyq « 1 N N ÿ s “ 1 „ hpỹsq ¨ 1 p̄pỹsq ( 7 ) with ỹs „ p̄p.q , where p̄ is the importance distribution . We choose the importance distribution to be the approximate posterior predictive defined over the outputs . We generate an outcome ỹs by first sampling a set of parameters θ̃0 from the approximate posterior , and then generating ỹs from a model defined by that set of parameters θ̃0 . This distribution will sample mostly from regions in S with high probability under the true posterior predictive for input x . This is in contrast to a naive sum over all possible outcomes y , many of which will have a negligible contribution to the sum . This gives rise to an estimator of Mutual information ( obtained with Algorithm 1 ) with lower variance than its naive Monte Carlo counterpart ( see Appendix A ) . Algorithm 1 : Importance sampling estimator of MI for s “ 1 to N do Sample θ̃0 „ qpθq ; ỹs „ P py|x , θ “ θ̃0q ; for m “ 1 to M do Sample θ̃m „ qpθq ; Compute ps , m = P py “ ỹs|z , θ “ θ̃mq ; end for Compute ps “ 1M řM m “ 1 ps , m ; hs “ 1 M řM m “ 1pps , m log ps , mq ´ ps log ps ; end for Return Mpxq = 1N řN s “ 1 ” hs ¨ 1ps ı

Summary: The paper considers the problem of black-box optimization of high-dimensional discrete objects where the inputs are embedded into a continuous latent space (typically using a deep generative model) and the optimization process takes place in this latent space. The main hypothesis in the paper is that existing methods for this problem might not be robust because they explore uncertain parts of the latent space with no training data in that region. The paper proposes to leverage the epistemic uncertainty of the decoder (measured in terms of the mutual information between model parameters and outputs) in order to navigate the latent space for black-box optimization. The mutual information is computed via an importance sampling estimate and employed in Bayesian optimization (BO) and gradient based search on the latent space. Experiments are performed on digit generation, arithmetic expression and molecular generation benchmarks.