Model Inversion Networks for Model-Based Optimization

1 INTRODUCTION . Data-driven optimization problems arise in a range of domains : from protein design ( Brookes et al. , 2019 ) to automated aircraft design ( Hoburg & Abbeel , 2012 ) , from the design of robots ( Liao et al. , 2019 ) to the design of neural net architectures ( Zoph & Le , 2017 ) and learning from logged feedback , such as optimizing user preferences in recommender systems . Such problems require optimizing unknown reward or score functions using previously collected data consisting of pairs of inputs and corresponding score values , without direct access to the score function being optimized . This can be especially challenging when valid inputs lie on a low-dimensional manifold in the space of all inputs , e.g. , the space of valid aircraft designs or valid images . Existing methods to solve such problems often use derivative-free optimization ( Snoek et al. ) . Most of these techniques require active data collection where the unknown function is queried at new inputs . However , when function evaluation involves a complex real-world process , such as testing a new aircraft design or evaluating a new protein , such active methods can be very expensive . On the other hand , in many cases there is considerable prior data – existing aircraft and protein designs , and advertisements and user click rates , etc . – that could be leveraged to solve the optimization problem . In this work , our goal is to develop an optimization approach to solve such optimization problems that can ( 1 ) readily operate on high-dimensional inputs comprising a narrow , low-dimensional manifold , such as natural images , ( 2 ) readily utilize offline static data , and ( 3 ) learn with minimal active data collection if needed . We can define this problem setting formally as the optimization problem x ? = arg max x f ( x ) , ( 1 ) where the function f ( x ) is unknown , and we have access to a dataset D = { ( x1 , y1 ) , . . . , ( xN , yN ) } , where yi denotes the value f ( xi ) . If no further data collection is possible , we call this the data-driven model-based optimization setting . This can also be extended to the contextual setting , where the aim is to optimize the expected score function value across a context distribution . That is , π ? = arg max π Ec∼p0 ( · ) [ f ( c , π ( c ) ) ] , ( 2 ) where π ? maps contexts c to inputs x , such that the expected score under the context distribution p0 ( c ) is optimized . As before , f ( c , x ) is unknown and we have access to a dataset D = { ( ci , xi , yi ) } Ni=1 , where yi is the value of f ( ci , xi ) . Such contextual problems with logged datasets have been studied in the context of contextual bandits ( Swaminathan & Joachims , a ; Joachims et al. , 2018 ) . A simple way to approach these model-based optimization problems is to train a proxy function fθ ( x ) or fθ ( c , x ) , with parameters θ , to approximate the true score , using the dataset D. However , directly using fθ ( x ) in place of the true function f ( x ) in Equation ( 1 ) generally works poorly , because the optimizer will quickly find an input x for which fθ ( x ) outputs an erroneously large value . This issue is especially severe when the inputs x lie on a narrow manifold in a high-dimensional space , such as the set of natural images ( Zhu et al. , 2016 ) . The function fθ ( x ) is only valid near the training distribution , and can output erroneously large values when queried at points chosen by the optimizer . Prior work has sought to addresses this issue by using uncertainty estimation and Bayesian models ( Snoek et al. , 2015 ) for fθ ( x ) , as well as active data collection ( Snoek et al. ) . However , explicit uncertainty estimation is difficult when the function fθ ( x ) is very complex or when x is high-dimensional . Instead of learning fθ ( x ) , we propose to learn the inverse function , mapping from values y to corresponding inputs x . This inverse mapping is one-to-many , and therefore requires a stochastic mapping , which we can express as f−1θ ( y , z ) → x , where z is a random variable . We term such models model inversion networks ( MINs ) . MINs provide us with a number of desirable properties : they can utilize static datasets , handle high-dimensional input spaces such as images , can handle contextual problems , and can accommodate both static datasets and active data collection . We discuss how to design simple active data collection methods for MINs , leverage advances in deep generative modeling ( Goodfellow et al . ; Brock et al. , 2019 ) , and scale to very high-dimensional input spaces . We experimentally demonstrate MINs in a range of settings , showing that they outperform prior methods on high-dimensional input spaces , perform competitively to Bayesian optimization methods on tasks with active data collection and lower-dimensional inputs , and substantially outperform prior methods on contextual optimization from logged data ( Swaminathan & Joachims , a ) . 2 RELATED WORK . Bayesian optimization . In this paper , we aim to solve data-driven optimization problems . Most prior work aimed at solving such optimization problems has focused on the active setting . This includes algorithms such as the cross entropy method ( CEM ) and related derivative-free methods Rubinstein ( 1996 ) ; Rubinstein & Kroese ( 2004 ) , reward weighted regression Peters & Schaal , Bayesian optimization methods based on Gaussian processes Shahriari et al . ( 2016 ) ; Snoek et al . ; 2015 ) , and variants that replace GPs with parametric acquisition function approximators such as Bayesian neural networks ( Snoek et al. , 2015 ) and latent variable models ( Kim et al. , 2019 ; Garnelo et al. , 2018b ; a ) , as well as more recent methods such as CbAS ( Brookes et al. , 2019 ) . These methods require the ability to query the true function f ( x ) at each iteration to iteratively arrive at a near-optimal solution . We show in Section 3.3 that MINs can be applied to such an active setting as well , and in our experiments we show that MINs can perform competitively with these prior methods . Additionally , we show that MINs can be applied to the static setting , where these prior methods are not applicable . Furthermore , most conventional BO methods do not scale favourably to high-dimensional input spaces , such as images , while MINs can handle image inputs effectively . Contextual bandits . Equation 2 captures the class of contextual bandit problems . Prior work on batch contextual bandits has focused on batch learning from bandit feedback ( BLBF ) , where the learner needs to produce the best possible policy that optimizes the score function from logged experience . Existing approaches build on the counterfactual risk minimization ( CRM ) principle ( Swaminathan & Joachims , a ; b ) , and have been extended to work with deep nets ( Joachims et al. , 2018 ) . In our comparisons , we find that MINs substantially outperform these prior methods in the batch contextual bandit setting . Deep generative modeling . Recently , deep generative modeling approaches have been very successful at modelling high-dimensional manifolds such as natural images ( Goodfellow et al . ; Van Den Oord et al . ; Dinh et al. , 2016 ) , speech ( van den Oord et al. , 2018 ) , text ( Yu et al . ) , alloy composition prediction ( Nguyen et al . ) , etc . MINs combine the strength of such generative models with important algorithmic decisions to solve model-based optimization problems . In our experimental evaluation , we show that these design decisions are important for adapting deep generative models to model-based optimization , and it is difficult to perform effective optimization without them . 3 MODEL INVERSION NETWORKS . In this section , we describe our model inversion networks ( MINs ) method , which can perform both active and passive model-based optimization over high-dimensional input spaces . Problem statement . Our goal is to solve optimization problems of the form x ? = arg maxx f ( x ) , where the function f ( x ) is not known , but we must instead use a dataset of input-output tuples D = { ( xi , yi ) } . In the contextual setting described in Equation ( 2 ) , each datapoint is also associated with a context ci . For clarity , we present our method in the non-contextual setting , but the contextual setting can be derived analogously by conditioning all functions on the context . In the active setting , which is most often studied in prior work , the algorithm is allowed to actively query f ( x ) one or more times on each iteration to augment the dataset , while in the static setting , only an initial static dataset is available . The goal is to obtain the best possible x ? ( i.e. , the one with highest possible value of f ( x ? ) ) . One naïve way of solving MBO problems is to learn a proxy score function fθ ( x ) , via standard empirical risk minimization . We could then maximize this learned function with respect to x via standard optimization methods . However , naïve applications of such a method would fail for two reasons . First , the proxy function fθ ( x ) may not be accurate outside the samples on which it is trained , and optimization with respect to it may simply lead to values of x for which fθ ( x ) makes the largest mistake in the negative direction . The second problem is more subtle . When x lies on a narrow manifold in very high-dimensional space ( such as the space of natural images ) , the optimizer can produce invalid values of x , which result in arbitrary outputs when fed into fθ ( x ) . Since the shape of this manifold is unknown , it is difficult to constrain the optimizer to prevent this . This second problem is rarely addressed or discussed in prior work , which typically focuses on optimization over low-dimensional and compact domains with known bounds . 3.1 OPTIMIZATION VIA INVERSE MAPS . Part of the reason for the brittleness of the naïve approach above is that fθ ( x ) has a high-dimensional input space , making it easy for the optimizer to find inputs x for which the proxy function produces an unreasonable output . Can we instead learn a function with a small input space , which implicitly understands the space of valid , in-distribution values for x ? The main idea behind our approach is to model an inverse map that produces a value of x given a score value y , given by f−1θ : Y → X . The input to the inverse map is a scalar , making it comparatively easy to constrain to valid values , and by directly generating the inputs x , an approximation to the inverse function must implicitly understand which input values are valid . As multiple x values can correspond to the same y , we design f−1θ as a stochastic map that maps a score value along with a dz-dimensional random vector to a x , f−1θ : Y × Z → X , where z is distributed according to a prior distribution p0 ( z ) . To define the inverse map objective , let the data distribution be denoted pD ( x , y ) , let pD ( y ) be the marginal over y , and let p ( y ) be an any distribution defined on Y ( which could be equal to pD ( y ) ) . We can train the proxy inverse map f−1θ under distribution p ( y ) by minimizing the following objective : Lp ( D ) = Ey∼p ( y ) [ D ( pD ( x|y ) , pf−1θ ( x|y ) ) ] , ( 3 ) where pf−1θ ( x|y ) is obtained by marginalizing over z , pf−1θ ( x|y ) = ∫ z p0 ( z ) · 1 [ [ x = f−1θ ( z , y ) ] ] dz , and D is a measure of divergence between the two distributions . Using the Kullback-Leibler divergence leads to maximum likelihood learning , while Jensen-Shannon divergence motivates a GAN-style training objective . MINs can be adapted to the contextual setting by passing in the context as an input and learning f−1θ ( yi , z , ci ) . In standard empirical risk minimization , we would choose p ( y ) to be the data distribution pD ( y ) , such that the expectation Algorithm 1 Generic Algorithm for MINs 1 : Input : pD ( y ) : distribution of y in D 2 : Train inverse map f−1θ : Y × Z → X using objective ( Equation 3 ) with reweighting , active data collection if needed 3 : x ? ← APPROX-INFER ( f−1θ , pD ( y ) ) 4 : return x ? can be approximated simply by sampling training tuples ( xi , yi ) from the training set . However , as we will discuss in Section 3.3 , a more careful choice for p ( y ) can lead to better performance . The MIN algorithm is based on training an inverse map , and then using it via the inference procedure in Section 3.2 to infer the x that approximately optimizes f ( x ) . The structure of the MIN algorithm is shown in Algorithm 1 .

Summary: This paper tackles the problem of solving a black-box optimization problem where only some samples have been observed. This task requires a good model that can be both expressive and generalizable. Instead of learning only a single forward model of x -> y, this paper proposes to additionally use a mapping from y -> x. Optimizing in the space of z instead of x can be much simpler, and this should also act as a strong regularizer during training. Specifically, the paper uses a GAN that transforms [y,z] -> x, where z is stochastically sampled. This paper further proposes a reweighting scheme that interpolates between a uniform weighting and weighting the best sample so far, as well as a sampling procedure that iteratively samples points and refits a second model, which was inspired by Thompson sampling.