MALIZED MAXIMUM LIKELIHOOD ESTIMATION
Year: 2021

Abstract:
In this work we consider data-driven optimization problems where one must maximize a function given only queries at a fixed set of points. This problem setting emerges in many domains where function evaluation is a complex and expensive process, such as in the design of materials, vehicles, or neural network architectures. Because the available data typically only covers a small manifold of the possible space of inputs, a principal challenge is to be able to construct algorithms that can reason about uncertainty and out-of-distribution values, since a naive optimizer can easily exploit an estimated model to return adversarial inputs. We propose to tackle this problem by leveraging the normalized maximum-likelihood (NML) estimator, which provides a principled approach to handling uncertainty and out-of-distribution inputs. While in the standard formulation NML is intractable, we propose a tractable approximation that allows us to scale our method to high-capacity neural network models. We demonstrate that our method can effectively optimize high-dimensional design problems in a variety of disciplines such as chemistry, biology, and materials engineering.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 PRELIMINARIES
text: 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.

heading: 4 NEMO: NORMALIZED MAXIMUM LIKELIHOOD ESTIMATION FOR MODEL-BASED OPTIMIZATION
text: We now present NEMO, our proposed algorithm for high-dimensional offline MBO. NEMO is a tractable scheme for estimating and optimizing the estimated expected value of the target function under the CNML distribution. As mentioned above, the CNML estimator (Eqn. 2) is difficult to compute directly, because it requires a) obtaining the MLE for each value of y, and b) integrating these estimates over y. In this section, we describe how to address these two issues, using amortization and quantization. We outline the high-level pseudocode in Algorithm 1, and presented a more detailed implementation in Appendix A.2.1.

heading: 4.1 AN ITERATIVE ALGORITHM FOR MODEL-BASED OPTIMIZATION
text: We first describe the overall structure of our algorithm, which addresses issue a), the intractability of computing an MLE estimate for every point we wish to query. In this section we assume that the domain of y is discrete, and describe in the following section how we utilize a quantization scheme to approximate a continuous y with a discrete one.
Recall from Section 3 that we wish to construct a proxy for the ground truth, which we will then optimize with gradient ascent. The most straightforward way to integrate NML and MBO would be to fully compute the NML distribution described by Eqn. 2 at each optimization step, conditioned on the current optimization iterate xt. This would produce a conditional distribution pNML(y|x) over output values, and we can optimize xt with respect to some function of this distribution, such as the mean. While this method is tractable to implement for small problems, it will still be significantly slower than standard optimization methods, because it requires finding the MLE estimate for every y value per iteration of the algorithm. This can easily become prohibitively expensive when using large neural networks on high-dimensional problems.
To remedy this problem, we propose to amortize the learning process by incrementally learning the NML distribution while optimizing the iterate xt. In order to do this, we maintain one model per value of y, θ̂k, each corresponding to one element in the normalizing constant of the NML distribution. During each step of the algorithm, we sample a batch of datapoints, and train each model by appending the current iterate xt as well as a label yt,k to the batch with a weight w (which is typically set to w = 1/N ). We then perform a number of gradient step on each model, and use
the resulting models to form an estimate of the NML distribution pNML(yt|xt). We then compute a score from the NML distribution, such as the mean, and perform one step of gradient ascent on xt.
While the incremental algorithm produces only an approximation to the true NML distribution, it brings the computational complexity of the resulting algorithm to just O(K) gradient steps per iteration, rather than solving entire inner-loop optimization problems. This brings the computational cost to be comparable to other baseline methods we evaluated for MBO.

heading: 4.2 QUANTIZATION AND ARCHITECTURE
text: The next challenge towards developing a practical NML method is addressing issue b), the intractability of integrating over a continuous y. We propose to tackle this issue with quantization and a specialized architecture for modeling the ground truth function.
Quantization. One situation in which the denominator is tractable is when the domain of y is discrete and finite. In such a scenario, we could train K models, where K is the size of the domain, and directly sum over the likelihood estimates to compute the normalizing factor.
In order to turn the NML distribution into a tractable, discrete problem, we quantize all outputs in the dataset by flooring each y value to the nearest bin bykc, with the size of each interval defined as B = (ymax − ymin)/K. While quantization has potential to induce additional rounding errors to the optimization process, we find in our experiments in Section 5 that using moderate value such as K = 20 or K = 40 provides both a reasonably accurate solution while not being excessively demanding on computation.
This scheme of quantization can be interpreted as a rectangular quadrature method, where the integral over y is approximated as:∫
y p(y|x, θ̂D∪(x,y))dy ≈ B K∑ k=1 p(bykc|x, θ̂D∪(x,bykc))
Discretized logistic architecture. Quantization introduces unique difficulties into the optimization process for MBO. In particular, quantization results in flat regions in the optimization landscape, making using gradient-based algorithms to optimize both inputs x and models p(y|x, θ) challenging. In order to alleviate these issues, we propose to model the output using a discretized logistic archi-
tecture, depicted in Fig. 2. The discretized logistic architecture transforms and input x into the mean parameter of a logistic distribution µ(x), and outputs one minus the CDF of a logistic distribution queried at regular intervals of 1/K (recall that the CDF of a logistic distribution is itself the logistic or sigmoid function). Therefore, the final output is a vector o of length K, where element k is equal to σ(µ(x) + k/K). We note that similar architectures have been used elsewhere, such as for modeling the output over pixel intensity values in images (Salimans et al., 2017).
We train this model by first encoding a label y as a vector ydisc, where ydisc[k ≤ bin(y)] = 1 and elements ydisc[k > bin(y)] = 0. bin(y) denotes the index of the quantization bin that y falls under. The model is then trained using a standard binary cross entropy loss, applied per-element across the entire output vector. Because the output represents the one minus the CDF, the expected value of the discretized logistic architecture can be computed as ymean(x) = Ey∼p(y|x)[g(y)] =∑ k[g(k) − g(k − 1)]o[k]. If we assume that g normalizes all output values to [0, 1] in uniform bins after quantization, the mean can easily be computed as a sum over the entire output vector, ymean = 1 K ∑ k o[k]
Optimization The benefit of using such an architecture is that when optimizing for x, rather than optimizing the predicted output directly, we can compute gradients with respect to the logistic parameter µ. Because µ is a single scalar output of a feedforward network, it is less susceptible to flat gradients introduced by the quantization procedure. Optimizing with respect to µ is sensible as it shares the same global optimum as ymean, and gradients with respect to µ and ymean share a positive angle, as shown by the following theorem:
Proposition 4.1 (Discretized Logistic Gradients). Let µ(x) denote the mean of the discretized logistic architecture for input x, and ymean(x) denote the predicted mean. Then,
1. If x ∈ arg maxx µ(x), then x ∈ arg maxx ymean(x).
2. For any x, 〈∇xµ(x),∇xymean(x)〉 ≥ 0.
Proof. See Appendix. A.1.2.

heading: 4.3 THEORETICAL RESULTS
text: We now highlight some theoretical motivation for using CNML in the MBO setting, and show that estimating the true function with the CNML distribution is close to an expert even if the test label is chosen adversarially, which makes it difficult for an optimizer to exploit the model. As discussed earlier, the CNML distribution minimizes a notion of regret based on the log-loss with respect to an adversarial test distribution. This construction leads the CNML distribution to be very conservative for out-of-distribution inputs. However, the notion of regret in conventional CNML does not easily translate into a statement about the outputs of the function we are optimizing. In order to reconcile these differences, we introduce a new notion of regret, the functional regret, which measures the difference between the output estimated under some model against an expert within some function class Θ.
Definition 4.1 (Functional Regret). Let q(y|x) be an estimated conditional distribution, x represent a query input, and y∗ represent a label for x. We define the functional regret of a distribution q as:
Regretf (q,D,x, y∗) = |Ey∼q(y|x)[g(y)]− Ey∼p(y|x,θ̂)[g(y)]|
Where θ̂ is MLE estimator for the augmented dataset D ∪ (x, y∗) formed by appending (x, y∗) to D.
A straightforward choice for the evaluation function g is the identity function g(y) = y, in which case the functional regret controls the difference in expected values between q and the MLE estimate p(y|x, θ̂). We now show that the functional regret is bounded for the CNML distribution: Theorem 4.1. Let pNML be the conditional NML distribution defined in Eqn. 2. Then,
∀x max y∗
Regretf (pNML,D, x, y∗) ≤ 2gmax √ Γ(D, x)/2.
Γ(D, x) = log{ ∑ y p(y|x, θ̂D∪(x,y))} is the minimax individual regret, and gmax = maxy∈Y g(y).
Proof. See Appendix A.1.2.
This statement states that, for any test input x, the CNML estimator is close to the best possible expert if the test label y is chosen adversarially. Importantly, the expert is allowed to see the label of the test point, but the CNML estimator is not, which means that if the true function lies within the model class, this statement effectively controls the discrepancy in performance between the true function and pNML. The amount of slack is controlled by the minimax individual regret Γ (Fogel & Feder, 2018), which can be interpreted as a measure of uncertainty in the model. For large model classes Θ and data points x far away from the data, the individual regret is naturally larger as the NML estimator becomes more uncertain, but for data points x close to Θ the regret becomes very small. This behavior can be easily seen in Fig. 3, where the CNML distribution is very focused in regions close to the data but outputs large uncertainty estimates in out-of-distribution regions.

heading: 5 EXPERIMENTS
text: In our experimental evaluation, we aim to 1) evaluate how well the proposed quantized NML estimator estimates uncertainty in an offline setting, and 2) compare the performance of NEMO to a number of recently proposed offline MBO algorithms on high-dimensional offline MBO benchmark problems. Our code is available at https://sites.google.com/view/nemo-anonymous

heading: 5.1 MODELING WITH QUANTIZED NML
text: We begin with an illustrative example of modeling a function with quantized NML. We compared a learned a quantized NML distribution with a bootstrapped ensemble method (Breiman, 1996) on a simple 1-dimensional problem, shown in Fig. 3. The ensemble method is implemented by training 32 neural networks using the same model class as NML, but with resampled datasets and randomized initializations. Both methods are trained on a discretized output, using a softmax cross-entropy loss function. We see that in areas within the support of the data, the NML distribution is both confident and relatively accurate. However, in regions outside of the support, the quantized NML outputs a highly uncertain estimate. In contrast, the ensemble method, even with bootstrapping and random initializations, tends to produce an ensemble of models that all output similar values. Therefore, in regimes outside of the data, the ensemble still outputs highly confident estimates, even though they may be wrong.

heading: 5.2 HIGH-DIMENSIONAL MODEL-BASED OPTIMIZATION
text: We evaluated NEMO on a set of high-dimensional MBO problems. The details for the tasks, baselines, and experimental setup are as follows, and hyperparameter choices with additional implementation details can be found in Appendix A.2.

heading: 5.2.1 TASKS
text: We evaluated on 6 tasks from the Design-bench (Anonymous, 2021), modeled after real-world design problems for problems in materials engineering (Hamidieh, 2018), biology (Sarkisyan et al., 2016), and chemistry (Gaulton et al., 2012), and simulated robotics. Because we do not have access to a real physical process for evaluating the material and molecule design tasks, Design-bench follows experimental protocol used in prior work (Brookes et al., 2019; Fannjiang & Listgarten, 2020) which obtains a ground truth evaluation function by training a separate regressor model to evaluate the performance of designs. For the robotics tasks, designs are evaluated using the MuJoCo physics simulator (Todorov et al., 2012).
Superconductor. The Superconductor task involves designing a superconducting material that has a high critical temperature. The input is space is an 81-dimensional vector, representing properties such as atomic radius and valence of elements which make up the material. This dataset contains a total of 21,263 superconductoring materials proposed by Hamidieh (2018).
GFP. The goal of the green fluorescent protein (GFP) task is to design a protein with high fluorescence, based on work proposed by Sarkisyan et al. (2016). This task requires optimizing a 238-dimensional sequence of discrete variables, with each dimension representing one amino acid in the protein and taking on one of 20 values. We parameterize the input space as logits in order to make this discrete problem amenable to continuous optimization. In total, the dataset consists of 5000 such proteins annotated with fluorescence values.
MoleculeActivity. The MoleculeActivity task involves designing the substructure of a molecule that exhibits high activity when tested against a target assay (Gaulton et al., 2012). The input space
is represented by 1024 binary variables parameterized by logits, which corresponds to the Morgan radius 2 substructure fingerprints. This dataset contains a total of 4216 data points.
The final 3 tasks, HopperController, AntMorphology, and DKittyMorphology, involve designing robotic agents. HopperController involves learning the parameters of a 5126-parameter neural network controller for the Hopper-v2 task from OpenAI Gym (Brockman et al., 2016). The Ant and DKitty morphology tasks involve optimizing robot parameters such as size, orientation, and joint positions. AntMorphology has 60 parameters, and DKittyMorphology has 56 parameters.

heading: 5.2.2 BASELINES
text: In addition to NEMO, we evaluate several baseline methods. A logical alternative to NEMO is a forward ensemble method, since both NEMO and ensemble methods maintain a list of multiple models in order to approximate a distribution over the function value, and ensembles are often used to obtain uncertainty-aware models. We implement an ensemble baseline by trainingK networks on the task dataset with random initializations and bootstrapping, and then optimizing the mean value of the ensemble with gradient ascent. In our results in Table 1, we label the ensemble as “Ensemble” and a single forward model as “Forward”. Additionally, we implement a Bayesian optimization baseline wth Gaussian processes (GP-BO) for the Superconductor, GFP, and MoleculeActivity tasks, where we fit the parameters of a kernel and then optimize the expected improvement according to the posterior. We use an RBF kernel for the Superconductor task, and an inner product kernel for the GFP and MoleculeActivity tasks since they have large, discrete input spaces. Note that the GP baseline has no variance between runs since the resulting method is completely deterministic.
We evaluate 3 state-of-the-art methods for offline MBO: model inversion networks (MINs) (Kumar & Levine, 2019), conditioning by adaptive sampling (CbAS) (Brookes et al., 2019), and autofocused oracles (Fannjiang & Listgarten, 2020). MINs train an inverse mapping from outputs y to inputs x, and generate candidate inputs by searching over high values of y and evaluating these on the inverse model. CbAS uses a generative model of p(x) as a trust region to prevent model exploitation, and autofocused oracles expands upon CbAS by iteratively updating the learned proxy function and iterates within a minimax game based on a quantity known as the oracle gap.

heading: 5.3 RESULTS AND DISCUSSION
text: Our results are shown in Table 1. We follow an evaluation protocol used in prior work for design problems (Brookes et al., 2019; Fannjiang & Listgarten, 2020), where the algorithm proposes a set of candidate designs, and the 100th and 50th percentile of scores are reported. This mimics a realworld scenario in which a batch of designs can be synthesized in parallel, and the highest performing designs are selected for use.
For each experiment, we produced a batch of 128 candidate designs. MINs, CbAS, and autofocused oracles all learn a generative model to produce candidate designs, so we sampled this batch from the corresponding model. Ensemble methods and NEMO do not maintain generative models, so we instead optimized a batch of 128 particles. We report results averaged of 16 random seeds.
NEMO outperforms all methods on the Superconductor task by a very large margin, under both the 100th and 50th percentile metrics, and in the HopperController task under the 100th percentile metric. For the remaining tasks (GFP, MoleculeActivity, AntMorphology, and HopperMorphology), NEMO also produces competitive results in line with the best performing algorithm for each task. These results are promising in that NEMO performs consistently well across all 6 domains evaluated, and indicates a significant number of designs found in the GFP and Superconductor task were better than the best performing design in the dataset. In Appendix A.3, we present learning curves for NEMO, as well as an ablation study demonstrating the the beneficial effect of NML compared to direct optimization on a proxy function. Note that unlike the prior methods (MINs, CbAS, Autofocused), NEMO does not require training a generative model on the data, only a collection of forward models.

heading: 6 CONCLUSION
text: We have presented NEMO (Normalized Maximum Likelihood Estimation for Model-Based Optimization), an algorithm that mitigates model exploitation on MBO problems by constructing a conservative model of the true function. NEMO generates a tractable approximation to the NML distribution to provide conservative objective value estimates for out-of-distribution inputs. Our theoretical analysis also suggests that this approach is an effective way to estimate unknown objective functions outside the training distribution. We evaluated NEMO on a number of design problems in materials science, robotics, biology, and chemistry, where we show that it attains very large improvements on two tasks, while performing competitively with respect to prior methods on the other four.
While we studied offline MBO in this paper, we believe that NML is a promising framework for building uncertainty-aware models which could have numerous applications in a variety of other problem settings, such as in off-policy evaluation, batch learning from bandit feedback, or offline reinforcement learning.

heading: ACKNOWLEDGEMENTS
text: We would like to thank Brandon Trabucco and Aviral Kumar for assisting with the implementation and evaluation of baselines and helping with the benchmark. This research was funded by the Office of Naval Research, C3.ai, and computational resources from Amazon Web Services.

heading: A APPENDIX

heading: A.1 PROOFS

heading: A.1.1 PROOF OF PROPOSITION 4.1
text: This statement directly from monotonicity. Because ymean is a sum of monotonic functions in µ, ymean must also be a monotonic function of µ. This implies that the global maxima are the same, and that gradients must point in the same direction since 〈∇xµ,∇xymean〉 = 〈 ∇xµ, dymeandµ ∇xµ 〉 = dymean dµ ||∇xµ|| 2 2 ≥ 0.

heading: A.1.2 PROOF OF THEOREM 4.1
text: In this section we present the proof for Thm. 4.1, restated below:
∀x max y∗
Regretf (pNML,D, x, y∗) ≤ 2gmax √ Γ(D, x)/2
There are two lemmas we will use in our proof. First, the difference in expected value between two distributions p(x) and q(x) can be bounded by the total variation distance TV (p, q) and the maximum function value fmax = maxx f(x):
|Ep(x)[f(x)]− Eq(x)[f(x)]| = | ∑ x [p(x)− q(x)]f(x)|
≤ fmax| ∑ x [p(x)− q(x)]|
= fmax2TV (p(x), q(x))
Second, Fogel & Feder (2018) show that the NML distribution obtains the best possible minimax individual regret of
max y Regretind(pNML,D, x, y) = log{ ∑ y p(y|x, θ̂D∪(x,y))} def = Γ(D, x)
Using these two facts, we can show:
Regretf (pNML,D, x, y∗) = |Ey∼pNML(y|x)[g(y)]− Ey∼p(y|x,θ̂D∪(x,y∗))[g(y)]|
≤ 2gmaxTV (p(y|x, θ̂D∪(x,y∗)), pNML(y|x))
≤ 2gmax
√ 1
2 KL(p(y|x, θ̂D∪(x,y∗)), pNML(y|x)) ≤ 2gmax √ 1
2 max q
Ey∼q(y|x)[log p(y|x, θ̂D∪(x,y))− log pNML(y|x)] = 2gmax √ Γ(D, x)/2
Where we apply the total variation distance lemma from lines 1 to 2. From lines 2 to 3, we used Pinsker’s inequality to bound total variation with KL, and from lines 3 to 4 we used the fact that the maximum regret is always greater than the KL, i.e.
KL(p(y|x, θ̂D∪(x,y∗)), pNML(y|x)) = Ey∼p(y|x,θ̂)[log p(y|x, θ̂D∪(x,y∗))− log pNML(y|x)]
≤ Ey∼p(y|x,θ̂)[log p(y|x, θ̂D∪(x,y))− log pNML(y|x)]
≤ max q Ey∼q(y|x)[log p(y|x, θ̂D∪(x,y))− log pNML(y|x)]
On the final step, we substituted the definition of Γ(D, x) as the individual regret of the NML distribution pNML.

heading: A.2 EXPERIMENTAL DETAILS

heading: A.2.1 DETAILED PSEUDOCODE
text: There are a number of additional details we implemented in order to improve the performance of NEMO for high-dimensional tasks. These include:
• Using the Adam optimizer (Kingma & Ba, 2014) rather than stochastic gradient descent. • Pretraining the models θ in order to initialize the procedure with an accurate initial model. • Optimizing over a batch of M designs, in order to follow previous evaluation protocols. • Optimizing with x with respect to the internal scores µ instead of EpNML [g(y)]. • Using target networks for the NML model, originally proposed in reinforcement learning
algorithms, to improve the stability of the method.
We present pseudocode for the practical implementation of NEMO below:
Algorithm 2 NEMO – Practical Instantiation Input Model class Θ, Dataset D = (x1:N , y1:N ), number of bins K, batch size M , learning rates αθ, αx, target update rate τ Initialize K models θ1:K0 Initialize batch of optimization iterates B0 = x1:M0 from the best performing x in D. Pretrain θ1:K0 using supervised learning on D. Initialize target networks θ̄1:K0 ← θ1:K0 . Quantize y1:N into K bins, denoted as bYc = {by1c, · · · bykc}. for iteration t in 1 . . . T do
for k in 1 . . .K do Sample x′t from batch Bt. Construct Augmented dataset: D′ ← D ∪ (x′t, bykc) Compute gradient gθ ← ∇θyt LogLikelihood(θ y t ,D′)
Update model θyt using Adam and gθ with learning rate αθ. end for Update target networks θ̄yt+1 ← τθ y t+1 + (1− τ)θ̄ y t for m in 1 . . .M do Compute internal values µk(xmt ) from target networks θ̄ y t for all k ∈ 1, · · · ,K
Compute gradient gx for xmt : ∇x 1K ∑ k µ
k(xmt ) Update xmt using Adam and gradient gx with learning rate αx.
end for end for

heading: A.2.2 HYPERPARAMETERS
text: The following table lists the hyperparameter settings we used for each task. We obtained our hyperparameter settings by performing a grid search across different settings of αθ, αx, and τ . We used 2-layer neural networks with softplus activations for all experiments. We used a smaller networks for GFP, Hopper, Ant, and DKitty (64-dimensional layers) and a lower discretization K for computational performance reasons, but we did not tune over these parameters.
Superconductor MoleculeActivity GFP Hopper Ant DKitty Learning rate αθ 0.05 0.005 Learning rate αx 0.1 0.01 0.001
Network Size 256,256 64,64 Discretization K 40 20
Batch size M 128 Target update rate τ 0.05
For baseline methods, please refer to Anonymous (2021) for hyperparameter settings.

heading: A.3 ABLATION STUDIES
text: In this section, we present 3 ablation studies. The first is on the effect of NML training, by comparing NEMO to optimizing a pretrained baseline neural network. The second ablation study investigates the architecture choice, comparing the discretized logistic architecture to a standard feedforward neural network. The final ablation study investigates the ratio of model optimization steps to input optimization steps. Each logging iteration in these figures corresponds to 50 loop iterations as depicted in Algorithm A.2.1.

heading: A.3.1 EFFECT OF NML
text: In this section, we present learning curves for NML, as well as an ablation study comparing NML (orange curve) against a forward optimization algorithm without NML (blue curve), labeled “No NML”. The “No NML” algorithm is identical to the NML algorithm detailed in Alg. A.2.1, except the NML learning rate αθ is set to 0.0. This means that the only model training that happens is done during the pretraining step. For illustrative purposes, we initialize the iterates from the worstperforming scores in the dataset to better visualize improvement, rather than initializing from the best scores which we used in our final reported numbers.
The scores on the Superconductor task are shown in the following figure. Removing NML training makes it very difficult for training on most designs, as shown by the poor performance on the 50th percentile metric.
The scores on the MoleculeActivity task follow a similar trend.
And finally, the scores on the GFP task also display the same trend.

heading: A.3.2 ARCHITECTURE CHOICE
text: In this ablation study, we investigate the efficacy of the discretized logistic architecture. As a baseline, we compared against a standard feedforward network, trained with a softmax cross-entropy loss to predict the discretized output y. We label this network as ”Categorical”, because the output of the network is a categorical distribution. All other hyperparameters, including network layer sizes, remain unchanged from those reported in Appendix. A.2.2.
On the Superconductor task, both the discretized logistic and categorical networks score well on the 100th percentile metric, but the Categorical architecture displays less consistency in optimizing designs, as given by poor performance on the 50th percentile metric.
On the MoleculeActivity task, the Categorical network performs comparatively better, but still underperforms the discretized logistic architecture.

heading: A.3.3 RATIO OF OPTIMIZATION STEPS
text: In this ablation study, we investigate the effect of the ratio of model optimization steps to input optimization step. For this experiment, we fix the learning rate αx to the hyperparameter values in Appendix. A.2.2, fix the input optimization steps to 1, and vary the number of model optimization steps we take.
We first investigate the Superconductor task, using a small model learning rate αθ = 0.0005. In this setting, we see a clear trend that additional model optimization steps are helpful and increase convergence speed.
Using a higher learning rate of αθ = 0.05, a smaller amount of steps works better, which suggests that it is easy to destabilize the learning process using a higher learning rate.
A similar trend holds true in the MoleculeActivity task, albeit less pronounced. The following figure uses a learning rate of αθ = 0.0005, and we once again see that more model optimization steps leads to increased performance.
And using a learning rate of αθ = 0.05, the advantage becomes less clear.
Overall, while we performed a grid search over the learning rates to achieve the highest performance, using a large number of model optimization steps with a small learning rate αθ appears to be a consistent strategy which performs well.

Review summaries:
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.