Offline Model-Based Optimization via Normalized Maximum Likelihood Estimation

1 INTRODUCTION . Many real-world optimization problems involve function evaluations that are the result of expensive or time-consuming process . Examples occur in the design of materials ( Mansouri Tehrani et al. , 2018 ) , proteins ( Brookes et al. , 2019 ; Kumar & Levine , 2019 ) , neural network architectures ( Zoph & Le , 2016 ) , or vehicles ( Hoburg & Abbeel , 2014 ) . Rather than settling for a slow and expensive optimization process through repeated function evaluations , one may instead adopt a data-driven approach , where a large dataset of previously collected input-output pairs is given in lieu of running expensive function queries . Not only could this approach be more economical , but in some domains , such as in the design of drugs or vehicles , function evaluations pose safety concerns and an online method may simply be impractical . We refer to this setting as the offline model-based optimization ( MBO ) problem , where a static dataset is available but function queries are not allowed . A straightforward method to solving offline MBO problems would be to estimate a proxy of the ground truth function f̂θ using supervised learning , and to optimize the input x with respect to this proxy . However , this approach is brittle and prone to failure , because the model-fitting process often has little control over the values of the proxy function on inputs outside of the training set . An algorithm that directly optimizes f̂θ could easily exploit the proxy to produce adversarial inputs that nevertheless are scored highly under f̂θ ( Kumar & Levine , 2019 ; Fannjiang & Listgarten , 2020 ) . In order to counteract the effects of model exploitation , we propose to use the normalized maximum likelihood framework ( NML ) ( Barron et al. , 1998 ) . The NML estimator produces the distribution closest to the MLE assuming an adversarial output label , and has been shown to be effective for resisting adversarial attacks ( Bibas et al. , 2019 ) . Moreover , NML provides a principled approach to generating uncertainty estimates which allows it to reason about out-of-distribution queries . However , because NML is typically intractable except for a handful of special cases ( Roos et al. , 2008 ) , we show in this work how we can circumvent intractability issues with NML in order to construct a reliable and robust method for MBO . Because of its general formulation , the NML distribution pro- vides a flexible approach to constructing conservative and robust estimators using high-dimensional models such as neural networks . The main contribution of this work is to develop an offline MBO algorithm that utilizes a novel approximation to the NML distribution to obtain an uncertainty-aware forward model for optimization , which we call NEMO ( Normalized maximum likelihood Estimation for Model-based Optimization ) . The basic premise of NEMO is to construct a conditional NML distribution that maps inputs to a distribution over outputs . While constructing the NML distribution is intractable in general , we discuss novel methods to amortize the computational cost of NML , which allows us the scale our method to practical problems with high dimensional inputs using neural networks . A separate optimization algorithm can then be used to optimize over the output to any desired confidence level . Theoretically , we provide insight into why NML is useful for the MBO setting by showing a regret bound for modeling the ground truth function . Empirically , we evaluate our method on a selection of tasks from the Design Benchmark ( Anonymous , 2021 ) , where we show that our method performs competitively with state-of-the-art baselines . Additionally , we provide a qualitative analysis of the uncertainty estimates produced by NEMO , showing that it provides reasonable uncertainty estimates , while commonly used methods such as ensembles can produce erroneous estimates that are both confident and wrong in low-data regimes . 2 RELATED WORK . Derivative-free optimization methods are typically used in settings where only function evaluations are available . This includes methods such as REINFORCE ( Williams , 1992 ) and reward-weighted regression ( Peters & Schaal , 2007 ) in reinforcement learning , the cross-entropy method ( Rubinstein , 1999 ) , latent variable models ( Garnelo et al. , 2018 ; Kim et al. , 2019 ) , and Bayesian optimization ( Snoek et al. , 2012 ; Shahriari et al. , 2015 ) . Of these approaches , Bayesian optimization is the most often used when function evaluations are expensive and limited . However , all of the aforementioned methods focus on the active or online setting , whereas in this work , we are concerned with the offline setting where additional function evaluations are not available . Normalized maximum likelihood is an information-theoretic framework based on the minimum description length principle ( Rissanen , 1978 ) . While the standard NML formulation is purely generative , the conditional or predictive NML setting can be used Rissanen & Roos ( 2007 ) ; Fogel & Feder ( 2018 ) for supervised learning and prediction problems . Bibas et al . ( 2019 ) apply this framework for prediction using deep neural networks , but require an expensive finetuning process for every input . The goal of our work is to provide a scalable and tractable method to approximate the CNML distribution , and we apply this framework to offline optimization problems . Like CNML , conformal prediction ( Shafer & Vovk , 2008 ) is concerned with predicting the value of a query point ŷt+1 given a prior dataset , and provides per-instance confidence intervals , based on how consistent the new input is with the rest of the dataset . Our work instead relies on the NML framework , where the NML regret serves a similar purpose for measuring how close a new query point is to existing , known data . The offline model-based optimization problem has been applied to problems such as designing DNA ( Killoran et al. , 2017 ) , drugs ( Popova et al. , 2018 ) , or materials ( Hautier et al. , 2010 ) . The estimation of distribution algorithm ( Bengoetxea et al. , 2001 ) alternates between searching in the input space and model space using a maximum likelihood objective . Kumar & Levine ( 2019 ) propose to learn an inverse mapping from output values to input values , and optimize over the output values which produce consistent input values . Brookes et al . ( 2019 ) propose CbAS , which uses a trust-region to limit exploitation of the model . Fannjiang & Listgarten ( 2020 ) casts the MBO problem as a minimax game based on the oracle gap , or the value between the ground truth function and the estimated function . In contrast to these works , we develop an approach to MBO which explicitly reasons about uncertainty . Approaches which utilize uncertainty , such as Bayesian optimization , are commonly used in online settings , and we expect these to work in offline settings as well . There are several related areas that could arguably be viewed as special cases of MBO . One is in contextual bandits under the batch learning from bandit feedback setting , where learning is often done on logged experience ( Swaminathan & Joachims , 2015 ; Joachims et al. , 2018 ) , or offline reinforcement learning ( Levine et al. , 2020 ) , where model-based methods construct estimates of the MDP parameters ( Kidambi et al. , 2020 ; Yu et al. , 2020 ) . Our work focuses on a more generic function optimization setting , but could be applied in these domains as well . 3 PRELIMINARIES . We begin by reviewing the problem formulation for offline model-based optimization , as well as necessary background on the normalized maximum likelihood estimator . Problem statement . We define the offline model-based optimization ( MBO ) problem as follows . Assume the existence of a stochastic ground truth function f ( y|x ) . The MBO algorithm is given a dataset D of inputs x along with outputs y sampled from f ( y|x ) . Like in standard optimization problems , the goal of MBO is to find the input value that maximizes the true function : x∗ = argmaxxEy∼f ( y|x ) [ y ] . ( 1 ) However , in offline MBO , the algorithm is not allowed to query the true function f ( y|x ) , and must find the best possible point x∗ using only the guidance of a fixed dataset D = { x1 : N , y1 : N } . One approach to solving this problem is to introduce a separate proxy function f̂θ ( y|x ) ≈ f ( y|x ) , which is learned from D as an estimate of the true function . From here , standard optimization algorithms such as gradient descent can be used to find the optimum of the proxy function , x̂∗ = argmaxxEy∼f̂θ ( y|x ) [ y ] . Alternatively , a trivial algorithm could be to select the highestperforming point in the dataset . While adversarial ground truth functions can easily be constructed where this is the best one can do ( e.g. , if f ( x ) = −∞ on any x /∈ D ) , in many reasonable domains it should be possible to perform better than the best point in the dataset . Conditional normalized maximum likelihood . In order to produce a conditional distribution pNML ( y|x ) we can use for estimating the ground truth function , we leverage the conditional or predictive NML ( CNML ) framework ( Rissanen & Roos , 2007 ; Fogel & Feder , 2018 ; Bibas et al. , 2019 ) . Intuitively , the CNML distribution is the distribution closest to the MLE assuming the test label y is chosen adversarially . This is useful for the MBO setting since we do not know the ground truth value y at points we are querying during optimization , and the CNML distribution gives us conservative estimates that help mitigate model exploitation ( see Fig . 1 ) . Formally , the CNML estimator is the minimax solution to a notion of regret , called the individual regret defined as Regretind ( h , y ) = log p ( y|x , θ̂D∪ ( x , y ) ) − log h ( y|x ) , and pNML ( y|x ) = arg minh maxy′ Regretind ( h , y′ ) ( Fogel & Feder , 2018 ) . The notation D ∪ ( x , y ) refers to an augmented dataset by appending a query point and label ( x , y ) , to a fixed offline dataset D , and θ̂D∪ ( x , y ) denotes the MLE estimate for this augmented dataset . The query point ( x , y ) serves to represent the test point we are interested in modeling . The solution to the minimax problem can be expressed as ( Fogel & Feder , 2018 ) : pNML ( y|x ) = p ( y|x , θ̂D∪ ( x , y ) ) ∫ y′ p ( y′|x , θ̂D∪ ( x , y′ ) ) dy′ , ( 2 ) where θ̂D∪ ( x , y ) = arg maxθ 1 N+1 ∑ ( x , y ) ∈D∪ ( x , y ) log p ( y|x , θ ) is the maximum likelihood estimate for p using the dataset D augmented with ( x , y ) . Algorithm 1 NEMO : Normalized Maximum Likelihood for Model-Based Optimization Input Model class { fθ : θ ∈ Θ } , Dataset D = ( x1 : N , y1 : N ) , number of bins K , evaluation function g ( y ) , learning rates αθ , αx . Initialize K models θ1 : K0 , optimization iterate x0 Quantize y1 : N into K bins , denoted as bYc = { by1c , · · · bykc } . for iteration t in 1 . . . T do for k in 1 . . .K do construct augmented dataset : D′ ← D ∪ ( xt , bykc ) . update model : θkt+1 ← θkt + αθ∇θkt LogLikelihood ( θ k t , D′ ) end for estimate CNML distribution : p̂NML ( y|xt ) ∝ p ( y|xt , θyt ) / ∑ k p ( bykc|xt , θkt ) Update x : xt+1 ← xt + αx∇xEy∼p̂NML ( y|x ) [ g ( y ) ] end for The NML family of estimators has connections to Bayesian methods , and has shown to be asymptotically equivalent to Bayesian inference under the uninformative Jeffreys prior ( Rissanen , 1996 ) . NML and Bayesian modeling both suffer from intractability , albeit for different reasons . Bayesian modeling is generally intractable outside of special choices of the prior and model class Θ where conjugacy can be exploited . On the other hand , NML is intractable because the denominator requires integrating and training a MLE estimator for every possible y . One of the primary contributions of this paper is to discuss how to approximate this intractable computation with a tractable one that is sufficient for optimization on challenging problems , which we discuss in Section 4 .

Summary: The paper proposes an approximation method, called NEMO (Normalized maximum likelihood Estimation for model-based optimization)  to compute the conditional normalized maximum log-likelihood of a query data point as a way to quantify the uncertainty in a forward prediction model in offline model-based optimization problems. The main idea is to construct a conditional NML (CNML) distribution that maps the high-dimensional inputs to a distribution over output variables. In addition, the paper provides a theoretical motivation that estimating the true function with the CNML is close to the best possible expert even if the test label is chosen adversarially, which is a great challenge for an optimizer to exploit the model. By using this CNML on three offline optimization benchmark datasets (Superconductor, GFP, MoleculeActivity) with gradient ascent-based optimization, the NEMO outputs all the other four baselines on the Superconductor dataset by almost 1.4x to 1.7x, the generate comparable results as the other four baselines method on the GFP and MoleculeActivity datasets.