Offline Model-Based Optimization via Normalized Maximum Likelihood Estimation
Authors: Justin Fu
Sergey Levine
Venue: ICLR 2021 Poster (2021)
Keywords: model-based optimization
normalized maximum likelihood

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 the MBO 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.

Published as a conference paper at ICLR 2021
OFFLINE MODEL-BASED OPTIMIZATION VIA NOR- MALIZED MAXIMUM LIKELIHOOD ESTIMATION
Justin Fu & Sergey Levine
Department of Electrical Engineering and Computer Science University of California, Berkeley (justinj fu, svlevine) deecs www.eeer..eeeeee................
ABSTRACT
In this work we consider data-driven optimization problems where one must max- imize a function given only queries at a fixed set of points. This problem setting emerges in many domains where function evaluation a complex and expensive process, such as in the design of materials, vehicles, or neural network architec- tures. Because the available data typically only covers a small manifold the possible space of imputs. principal challenge is to be able construct algorithms that can reason about uncertainty and titto--distribttion values, since a naive op- timizer easily exploit estimated model to return adversarial inputs. We propose to tackle this problem by leveraging the normalized maximum-likeliton (NML) estimator, which provides a principled approach to handling uncertainty and dittr--distribution inputs. While in the standard formulation NML is in- tractable, we propose a tractable approximation that allows to scale our method to high-capacin 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.
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. [019 Kumar & Levine 2019), neural network architectures (Zoph) & Le 2016), or vehicles (Hoburg & Abbeel] 2014). Rather than settling for 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 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 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 fo using supervised learning, and to optimize the input æ with respect to this proxy. However, this approach is brittle and prone 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 fo could easily exploit the proxy to produce adversarial inputs that nevertheless are scored highly under So (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 principled approach to generating uncertainty estimates which allows it to reason about out-of-distributiin queries. How- ever, because NML is typically intractable except for handful of special cases (Roos 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-

---

Published as a conference paper at ICLR 2021
vides flexible approach to constructing conservative and robust estimators using high nigh-dimensional models such as neural networks.
The main contribution of this work is to develop an offline MBO algorithm that utilizes novel approximation to the NML distribution to obtain an uncertaint--waae forward model for opti- mization, 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 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 sep- arate 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 MBÔ setting by showing regret bound for modeling the ground truth function. Empirically, we evaluate our method on selection of tasks from the Design Benchmark (Anonymous) 202T), where we show that our method performs competitively with state-of-the -rr baselines. Additionally, we provide qualitative analy- sis 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.
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 (Rubin- stein 1999), latent variable models (Garnelo et al. 2018 Kim et al., 2019), and Bayesian optimiza- tion (Snoek al., 2012 Shahriari al. 2015). Of these approaches, Bayesian optimization is the most often used when function evaluations are expensive and limited. However, all of the aforemen- tioned 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 -heroati framework based on the minimum description length principle (Rissanen) 1978). While the standard NML formulation is purely gen- erative, 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 frame- work for prediction using deep neural networks, but require an expensive finetuning process for every input. The goal of our work is to provide 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 tt+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 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 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. [ummar & Levine (2019) pro- pose 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 prob- lem 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 approach to MBO which explicitly reasons about uncertainty. Approaches which utilize uncertainty, such as Bayesian optimization, are commonly used in online settings, and 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 con- textual 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 reinforce- ment learning (Levine et al. 2020), where model-based methods construct estimates of the MDP
2

---

Published as a conference paper at ICLR 2021
An illustrative example of the construction of the CNML distribution. Left We wish to estimate the p(y(a) at some query point æ, marked by the vertical red line. Middle We compute the MLE assuming we knew the true label y, for every possible value of y. Right Finally, we normalize the predictions across all MLE models to produce PNML- The final prediction will likely exhibit large amounts of uncertainty queries far from the dataset, because it is easier for the individual MLE estimates to overfit to these outliers.
Figure 1:
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 well.
3 PRELIMINARIES
We begin by reviewing the problem formulation for offline model-based optimization, as well as necessary background the normalized maximum likelihood estimator.
Problem statement. We define the offline model-based optimization (MBO) problem as follows. Assume the existence of stochastic ground truth function fyy(x) The MBO algorithm is given a dataset 'D of inputs æ along with outputs y sampled from [ (z). Like in standard optimization problems, the goal of MBO is to find the input value that maximizes the true function:
x² = D
However, in offline MBO, the algorithm is not allowed to query the true function fy((z), and must find the best possible point x:" using only the guidance of fixed dataset = ( V1.N). One approach to solving this problem is to introduce a separate proxy function (v]x)) = f(ylær,, which is learned from D
estimate of the true function. From here, standard optimiza- tion algorithms such as gradient descent can be used to find the optimum of the proxy function, x argmax, y (w)la) [v]- Alternatively, trivial algorithm could be select the highest- performing point the dataset. While adversarial ground truth functions can easily be constructed where this is the best one can do (e.g., if f(x = -o 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 PNNL. (v) we can use for estimating the ground truth function, we leverage the conditional or pre- dictive NMI (CNML) framework (Rissanen & Roos, 2007 Fogel & Feder 2018 8 Bibas et al., 2019). Intuitively, the CNML distribution is the distribution closest the MLE assuming the test label 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. [D. Formally, the CNML estimator is the minimax solution to a notion of regret, called the individual regret defined as Regret (h, y)
log p(yla, ( - log (ylzz, and (NMM(U| (z) = arg min, max,'' Regret, v') (Fogel & Feder 2018). The notation D ( y) refers to an augmented dataset by appending a query point and label (x,y). y). to a fixed offline dataset D. and @Du(EE.) denotes the MLE estimate for this aug- mented dataset. The query point (z, y) serves to represent the test point we are interested in model- ing. The solution to the minimax problem can be expressed as (Fogel & Feder 2018):
ply|æ,, (Du(E.)) Syy p(u0 (
PNML (y|x)
where (p.())))) arg max, N+1 (ma))E bog((yla, xx, 8) is the maximum likelihood estimate for p using the dataset D augmented with (x. y).
1 Σ
3

---

Published as a conference paper at ICLR 2021
Algorithm 1 NEMO: Normalized Maximum Likelihood for Model-Based Optimization
Input Model class {fo : 0 E (), Dataset D (æ1:NN Vl:N), number of bins K, evaluation function g(y), learning rates Ctg, 0zz.
Initialize K models 0 O optimization iterate Zo
Quantize V1:N into K bins, denoted as [}'] [[al,.... [3k]}.
for iteration in 1... do for k in K do construct augmented dataset: D' + DU (z),, lvne)). update model: of t of + ag V LozLogiikchhhood( D end for
estimate CNML distribution: PNMN (v) x p(y| 02)/ Σn 02)
Update æ: 2t+1 < 2t + az V. y V (ulm) (vlm) [gy)]
end for
The NML family of estimators has connections to Bayesian methods, and has shown to be asymp- totically 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 MLE estimator for every possible y. One of the primary contributions of this paper is to discuss how to approximate this intractable computation with tractable one that sufficient for optimization challenging problems, which we discuss in Section 4
NEMO: NORMALIZED MAXIMUM LIKELIHOOD ESTIMATION FOR MODEL-BASED OPTIMIZATION
We now present NEMO, our proposed algorithm for high-dimensional offline MBO. NEMO is 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 MIE for each value y, and b) integrating these estimates over y. In this section, we describe how to address these two issues, using amortiza- tion and quantization. We outline the high-level pseudocode in Algorithm [ and presented a more detailed implementation in Appendix A.2.1
4.1 AN ITERATIVE ALGORITHM FOR MODEL-BASED OPTIMIZATION
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 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 each optimization step, conditioned on the current optimization iterate K;. This would produce a conditional distribution PNML (y/ 2) over output values, and we can optimize 2, with respect to some function 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 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 24, order to this, we maintain one model per value of y. Ôx, each corresponding to one element in the normalizing constant of the NML distribution. During each step of the algorithm, model by appending the current iterate 2, as well as label yr, to the batch with a weight w (which is typically set w 1/N). We then perform a number of gradient step each model, and use
sample batch of datapoints, and train each
4

---

Published as a conference paper at ICLR 2021
the resulting models to form estimate of the NML distribution PNML (Me (3)(2)). We then compute score from the NML distribution, such as the mean, and perform one step of gradient ascent on 2;.
While the incremental algorithm produces only 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.
4.2 QUANTIZATION AND ARCHITECTURE
The next challenge towards developing a practical NML method addressing issue b), the in- tractability of integrating over continuous y. We propose 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 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 prob- lem, we quantize all outputs in the dataset by flooring each y value to the nearest bin |. with the size of each interval defined as B = (Umax Ymin ) K. While quantization has potential to in- duce additional rounding errors to the optimization process, we find in our experiments in Section [S] that using moderate value such
1.0
0.1
0.0
0.0
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:
J, plyla, (D dy xo > p(luxJla (bl
A diagram of the discretized logistic architecture. The mean μ is denoted by an arrow, which is then passed through offset sigmoid functions to produce o.
Figure 2:
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 cradient-basedd algorithms to optimize both inputs and models ply|2, 8) challenging. In order to alleviate these issues, 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 logistic distribution p(æ), and outputs one minus the CDF of logistic distribution queried at regular intervals of 1/K (recall that the CDF of a logistic distribution is itself the logis- tic or sigmoid function). Therefore, the final output vector o of length K, where element is equal to o (p(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 as vector Ydise, where Jdise(k < bin(y)) = 1 and elements ydise[ [k > bin(y)) O. bin(y) denotes the index of the quantization bin that y falls under. The model is then trained using 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 Mmean (x) [y><p(ll=][[[)]
Ek((k) g(k- [)]][K]. If we assume that y 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, 3mcan * Σ. oik]
Optimization The benefit of using such an architecture is that when optimizing for æ, rather than optimizing the predicted output directly, we can compute gradients with respect to the logistic pa- rameter μ. Because j is a single scalar output of feedforward network, 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, shown by the following theorem:
5

---

Published as a conference paper at ICLR 2021
Proposition 4.1 (Discretized Logistic Gradients). Let p(z) denote the mean of the discretized lo- gistic architecture for input 2, and Ymeon (2) (x) denote the predicted mean. Then,
1. If E arg max (()) then a E arg max Umeaa (z))
2. For any x, (V=H(<)) V Vemmman (x)) 2 0.
Proof. See Appendix. A.I.2
4.3 THEORETICAL RESULTS
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 expert even if the test label chosen adversarially, which makes it difficult for an optimizer to exploit the model. As discussed earlier, the CNML distribution minimizes notion of regret based the log-loss with respect to adversarial test distribution. This construction leads the CNML distribution to be very conservative for cittof--iistibution 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 O
Definition 4.1 (Functional Regret). Let q(y(az) be estimated conditional distribution, represent a query input, and y² represent a label for x. We define the functional regret of a distribution q as:
Regret, q D, x. y²) = Ey--auul)) [g())l)llllll y-rlullm [g(1)]
Where is MLE estimator for the augmented dataset D U (x. v²) formed by appending (a.y") to D.
A straightforward choice for the evaluation function 9 is the identity function g(v) = y, in which case the functional regret controls the difference in expected values between and the MLE estimate p(y|æ, 8). We now show that the functional regret is bounded for the CNML distribution:
Theorem 4.1. Let pomr. be the conditional NML distribution defined in Equ. 2 Then,
Ve max Regret (PNML, D, 20max V I(D, //2
y¹
T(D,*) x)) = -og p(yle, ()) is the minimax individual regret, and Gmax = max (VV).
Proof. See Appendix A.1.2
This statement states that, for any test input æ, the CNML estimator is close to the best possible expert if the test label y is chosen adversarially. Importantly, the expert is allowed to 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 I' (Fogel & Feder 2018)), which can be interpreted as measure of uncertainty in the model. For large model classes & and data points æ far away from the data, the individual regret is naturally larger as the NML estimator becomes more uncertain, but for data points æ close ② 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.
5 EXPERIMENTS
In our experimental evaluation, we aim to 1) evaluate how well the proposed quantized NML estima- tor estimates uncertainty in an offline setting, and 2) compare the performance of NEMO number of recently proposed offline MBO algorithms on high dimensional offline MBO benchmark prob-
lems. Our code is available httpttss//sssse goog le com/ eew nemo anonymous

---

Published as a conference paper at ICLR 2021
MNN
NML
Ensemble
X
Figure 3: A comparison uncertainty estimates between quantized NML and bootstrapped ensem- bles. Left: A small training dataset collected using the ground truth function sin(x) and quantized into 32 bins. Each dot represents a data point. Middle: Predictions from quantized CNML, where darker regions indicate outputs with high estimated probability. The ground truth is marked by dotted blue line. Note that in regions where there is little data, the NML distribution tends to cor- rectly output very diffuse distribution with uncertain outputs. Right: Predictions from bootstrap ensemble of 32 models. In out-of-support regions far from the data, the bootstrap ensemble tends to underestimate uncertainty and produce overconfident predictions.
5.1 MODELING WITH QUANTIZED NML
We begin with an illustrative example of modeling a function with quantized NML. We compared learned quantized NML distribution with bootstrapped ensemble method (Breiman 1996) on simple I-dimensional problem, shown in Fig. 3 The ensemble method implemented by training 32 neural networks using the same model class 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.
5.2 DICHH-DIENSIOALLL MODEL-BASER OPTIMIZATION
We evaluated NEMO on set of high-dimensiona MBO problems. The details for the tasks, base- lines, and experimental setup are follows, and hyperparameter choices with additional implemen- tation details can be found in Appendix A.2
5.2.1 TASKS
We evaluated 6 tasks from the Design-bench (Anonymous) 2021). modeled after real-world de- sign problems for problems in materials engineering (Hamidieh] 2018). biology (Sarkisyan et al.] 2016), and chemistry (Gaulton et al. 2012), and simulated roboties. Because we do have access to real physical process for evaluating the material and molecule design tasks, Design-bench fol- lows experimental protocol used in prior work (Brookes 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).
The Superconductor task involves designing superconducting material that has high critical temperature. The input is space is an 81 -dimensionnl vector, representing properties such as atomic radius and valence of elements which make up the material. This dataset contains total of 21,263 superconductoring materials proposed by Hamidieh (2018).
Superconductor.
GFP. The goal of the green fluorescent protein (GFP) task to design a protein with high flu- orescence, based work proposed by [Sarkisyan al. (2016). This task requires optimizing 238-dimensional sequence of discrete variables, with each dimension representing one amino acid in the protein and taking to make this discrete problem amenable to continuous optimization. In total, the dataset consists of 5000 such proteins annotated with fluorescence values.
of 20 values. We parameterize the input space logits in order
MoleculeActivity. The MoleculeActivity task involves designing the substructure of a molecule that exhibits high activity when tested against target assay (Gaulton et al. 2012). The input space

---

Published as a conference paper at ICLR 2021
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, AntMorpholegy, and DKittyMorphologe ing robotic agents. HopperController involves learning the parameters of a 5126- 5126-parameter neural network controller for the Hopper-v2 task from OpenAl Gym (Brockman al. 2016). The Ant and DKitty morphology tasks involve optimizing robot parameters such as size, orientation, and joint positions. AntMorphology has 60 parameters, and MKittyMorpholog has 56 parameters.
involve design-
5.2.2 BASELINES
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 distribution over the function value, and ensembles are often used to obtain uncertainty-awaan -wwre models. We implement an ensemble baseline by training 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 we label the ensemble as "Ensemble" and a single forward model "Forward". Additionally, baseline wth Gaussian processes (GP-BO) for the Superconductor, GFP, and Molecule 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.
implement Bayesian optimization
We evaluate 3 state-of-the-ar methods for offline MBO: model inversion networks (MINs) (Kumar & Levine |019. conditioning by adaptive sampling (CbAS) (Brookes et al. (219)) and autofocused oracles (Fannjiang Listgarten 2020). MINs train an inverse mapping from outputs y inputs æ, and generate candidate inputs by searching over high values of 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.
5.3 RESULTS AND DISCUSSION
Our results are shown in Table We follow 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 reported. This mimics real- world scenario in which 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 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 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, MoleculeActivitt. 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 evalu- ated, 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 cused), NEMO does not require training a generative model on the data, only a collection of forward models.
proxy function. Note that unlike the prior methods (MINs, CbAS, Autofo-
8

---

Published as a conference paper at ICLR 2021
100 Perc.
Supereonductor
GFP
MolcculeActivity
NEMO (ours)
127.0 7.292
3.359 = 0.036
6,682 = 0.209
Ensemble
88.01 + 10.43
2.924 + 0.039
6.525 + 0.1159
Forward
89.64 9.201
2.894 + 0.001
6.636 + 0.066
MINs
80.23 + 10.67
3.315 0.033
6.508 + 0.236
CbAS
72.17 + 8.652
3,408 = 0.029
6.301 + 0.131
Autofoc.
77.07 + 11.11
3.365 = 0.023
6.345 = 0.141
GP-BO
89.72 + 0.000
2.894 0.000
6.745 + 0.000
Dataset Max
73.90
3.152
6.558
HopperController
AntMorpholog 393.7 = 6.135
DKityyMophoolog
NEMO (ours)
2130.1 + 506.9
431.6 = 47.79
Ensemble
1877.0 + 704.2
Forward
1050.8 284.5
399.9 4.941
390.7 + 49.24
MINs
746.1 + 636.8
388.5 9.085
352.9 + 38.65
CbAS
547.1 423.9
393.0 3.750
369.1 + 60.65
Autofoc.
443.8 + 142.9
386.9 + 10.58
376.3 1 47.47
Dataset Max
1361.6
108.5
215.9
Table 1: 1h0th percentile ground truth scores and standard deviations over a batch of 128 designs for each task, averaged across 16 trials.
thth Perc.
Superconductor
GFP
MoleculeActivity 5.814 1 0.092
NEMO (ours)
66.41 1: 4.618
3,219 1 0.039
Ensemble
48.72 -
2.910 = 0.020
6.412 = 0.123
Forward
54.06 5.060
2,894 0.000
6.401 0.186
MINs
37.32 + 10,50
3.135 = 0.019
5.806 = 0.078
ChAS
32.21 + 7.255
3.269 0.018
5.742 + 0.123
Autofoc.
31.57 7.457
3,216 0.029
5.759 + 0.158
GP-BO
72.42 0.000
2.894 0.000
6.373 + 0.000
HopperController
AntMorphology 326.9 5.229
DKityyMophollogy 180,8
NEMO (ours)
390.2
Ensemble
362.5 + 80.09
Forward
185.0 72.88
318.0 + 12.05
255.3 + 6.379
MINs
520.4 + 301.5
184.8 29.52
211.6 + 13.67
CbAS
132.5 23.88
267.3 16.55
203.2 + 3.580 199.3 + 8.909
Autofoc.
116.4 + 18.66
176.7 + 59.94
Table 2: 50th percentile ground truth scores and standard deviations batch of 128 designs for each task, averaged across 16 trials.
CONCLUSION
We have presented NEMO (Normalized Maximum Likelihood Estimation for Model-Based Op- timization), an algorithm that mitigates model exploitation on MBO problems by constructing conservative model of the true function. NEMO generates a tractable approximation to the NML distribution to provide conservative objective value estimates for disttribbbtt inputs. Our the- oretical analysis also suggests that this approach is 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 show that it attains very large im- provements 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 variety of other problem settings, such as in off-po evaluation, batch learning from bandit feedback, or offline reinforcement learning.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGEMENTS
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, ai, and computational resources from Amazon Web Services.
REFERENCES
Anonymous. Design-bench: Benchmarks for data-driven offline model-based optimization. In Submitted to International Conference on Learning Representations, 2021. URL https:// Under review.
openreview net/ wow.:rodddd0222222me
Andrew Barron, Jorma Rissanen, and Bin Yu. The minimum description length principle in coding and modeling. IEEE Transactions on Information Theory, 44(6):2743-72660, 1998.
Endika Bengoetxea, Pedro Larrañaga, Isabelle Bloch, and Aymeric Perchant. Estimation of distri- bution algorithms: A new evolutionary computation approach for graph matching problems. In International Workshop Energy Minimization Methods in Computer Vision and Pattern Recog- nition, pp. 454-469. Springer, 2001.
Koby Bibas, Yaniv Fogel, and Meir Feder. Deep puml: Predictive normalized maximum likelihood for deep neural networks. arXiv preprint 19X4v 2019.
Lco Breiman. Bagging predictors. Machine learning, 24(2):123-140, 1996.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint 1XXx::66060015000 2016.
David H Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. arXiv preprint /901iv:/991//0000 2019.
Clara Fannjiang and Jennifer Listgarten. Autofocused oracles for model-based design. arXiv preprint ar arXiv:2006.0810 08052, 2020.
Yaniv Fogel and Meir Feder. Universal supervised learning for individual data. arXiv preprint arXiv: 1rXXvv/8/2.0999 18/2.1122...2.... 2018.
Marta Garnelo, Dan Rosenbaum, Chris Maddison, Tiago Ramalho, David Saxton, Murray Shana- han, Yee Whye Teh, Danilo Rezende, and SM Eslami. Conditional neural processes. arXiv preprint arXiv:1807.01611 2018
Anna Gaulton, Louisa Bellis, A. Patricia Bento, Jon Chambers, Mark Davies, Anne Hersey, Yvonne Light, Shaun McGlinchey, David Michalovich, Bissan Al-Lazikani, and John P. Over- ington. Chembl: a large-scale scale bioactivity database for drug discovery. Nucleic acids research, 40(Database is 1000--D1110-00000000000000 107, Jan 2012. ISSN 1362-4962. doi: 10. 10.10///aar// URL https:://dbi.rg//11 1093 /arrarr gkr777
Kam Hamidieh. A data-driven statistical model for predicting the critical temperature of super- conductor. Computational Materials Science, 154:346 354, 2018 ISSN 0927-0256. doi: https:
2018.07.057. 002. URL http:/ / www. scwencediiiie.ccec..cccc rect. com/
Geoffroy Hautier, Christopher C Fischer, Anubhav Jain, Tim Mueller, and Gerbrand Ceder. Finding nature's missing temary oxide compounds using machine learning and density functional theory. Chemistry of Materials, 22 2112):3762-376 2010.
Warren Hoburg and Pieter Abbeel. Geometric programming for aircraft design optimization. AIAA Journal, 22411)2444-2226, 2014.
Thorsten Joachims, Adith Swaminathan, and Maarten de Rijke. Deep learning with logged bandit feedback. In International Conference Learning Representations, 2018.
10

---

Published as a conference paper at ICLR 2021
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. arXiv preprint arXiv:2005.05951, 2020.
Nathan Killoran, Leo J Lee, Andrew Delong, David Duvenaud, and Brendan J Frey. Generating and designing dna with deep generative models. arXiv preprint arXiv:1712.061144 2017.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Gamelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint arXiv:1901.05700 2019.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1412.6441,.69800 2014.
Aviral Kumar and Sergey Levine. Model inversion networks for model-based optimization. arXiv preprint arXiv:1912, 13464, 2019.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0164 2020.
Aria Mansouri Tehrani, Anton O Oliynyk, Marcus Parry, Zeshan Rizvi, Samantha Couper, Feng Lin, Lowell Miyagi, Taylor D Sparks, and Jakoah Brgoch. Machine learning directed search for ultraincompressible, superhard materials. Journal of the American Chemical Society, 140(31): 9844-9853, 2018.
Jan Peters and Stefan Schaal. Reinforcement learning by swward-weighted regression for operational space control. In Proceedings of the 24th international conference on Machine learning, pp. 745- 750, 2007.
Mariya Popova, Olexandr Isayev, and Alexander Tropsha. Deep reinforcement learning for de novo drug design. Science advances, 4(7))eaap7885, 2018.
Jorma Rissanen. Modeling by shortest data description. Automatica, 14(5):465-471, 1978.
Jorma Rissanen and Teemu Roos. Conditional nml universal models. 2007 Information Theory and Applications Workshop, pp. 337-341. IEEE, 2007.
Jorma J Rissanen. Fisher information and stochastic complexity. IEEE transactions on information theory, 42(1):40-4) 1996.
Teemu Roos, Tomi Silander, Petri Kontkanen, and Petri Myllymaki. Bayesian network structure learning using factorized nml universal models. In 2008 Information Theory and Applications Workshop, pp. 272-276. IEEE, 2008.
Reuven Rubinstein. The cross-entropy method for combinatorial and continuous optimization. Methodology and computing in applied probability, 1(2):127-190, 1999.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelenn++: Improving the pixelenn with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv: 2017.
Karen S. Sarkisyan, Dmitry A. Bolotin, Margarita V. Meer, Dinara R. Usmanova, Alexander S. Mishin, George V. Sharonov, Dmitry N. Ivankov, Nina G. Bozhanova, Mikhail S. Baranov, Onu- ralp Soylemez, Natalya S. Bogatyreva, Peter K. Vlasov, Evgeny S. Egorov, Maria D. Logacheva, Alexey S. Kondrashov, Dmitry M. Chudakov, Ekaterina V. Putintseva, Ilgar Z. Mamedov, Dan S. Tawfik, Konstantin A. Lukyanov. and Fyodor A. Kondrashov. Local fitness landscape of the green fluorescent protein. Nature, 533(7603):397740) May 2016. ISSN 1476-4687. doi: 10. 100./038/aturel1 7995. 7995. URL TRLLhtpps:/ doi http://ooo.ooggg////// 038 maturee7999.
Glenn Shafer and Vladimir Vovk. A tutorial conformal prediction. Journal of Machine Learning Research. 9(Mar):37--211, 2008
Bobak Shahriari, Kevin Swersky, Ziyu Wang, Ryan P Adams, and Nando De Freitas. Taking the human out of the loop: A review of bayesian optimization. Proceedings of the IEEE, 104(1): 148-175, 2015.
11

---

Published as a conference paper at ICLR 2021
Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. In Advances in neural information processing systems. pp. 2951-2959, 2012.
Adith Swaminathan and Thorsten Joachims. The self-normalized estimator for counterfactual learn- ing. In advances in neural information processing systems, pp. 3231-3239, 2015.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229-255 1992.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. arXiv preprint arXiv:2005.13239, 2020.
Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning. arXiv preprint 1..05776,, 2016.
12

---

Published as a conference paper at ICLR 2021
A
APPENDIX
A.1 PROOFS
A.1.1 PROOF PROPOSITION 4.1
This statement directly from monotonicity. Because Umean is a sum of monotonic functions in μs, Mmxan must also be monotonic function of μ. This implies that the global maxima are the same, and that gradients must point in the same direction since (Valk ar,, V Vellmmaa))
V M,
dybmon Vali
dp
dyneee dp
V 2/ 2 O.
A.1.2 PROOF OF THEOREM 4.1
In this section we present the proof for Thm. 4.1 restated below:
V max Regret, (PNML. D, a, x,,") S 29max [(D,x)/2
y
There are two lemmas we will use in our proof. First, the difference in expected value between two distributions p(z) and g(z) can be bounded by the total variation distance TV(p, q) and the maximum function value fmxx = max, [(x):
[Ep()]((=)] Ef(=)f(())]
- (())]f(=))
S fmax Σlo(z) - {(>)]l
fmax 2TV
Second, Fogel & Feder (2018) show that the NMI distribution obtains the best possible minimax individual regret of
max Regrel (PNML, D, = log Σ plule, (Du(
def del I(D,2)
Using these two facts, we can show:
Regret, (PNMLL D,r,y")
E pmo (vlm) [g(v)] - E do lo(v)]
2gm TV (p(y) (DU((.D))) ). PNML(
s 29max
K KL(plylæ, Q ( PNML (N))(
< 29max max E (ogg (Du(,,))))))) log (yla()))))
9
2gmax / T(D,2)/2
Where we apply the total variation distance lemma from lines to 2. From lines 2 to 3, we used Pinsker's inequality to bound total variation with KI, and from lines 3 to we used the fact that the maximum regret is always greater than the KI,, i.e.
K L(pull, (D)(arr)) )), PNML (yl=))
E
y>r(vll [log logply/< hllp:/ll,,,,peee....... (pu(a,))) log PNMI. []
E y lloggllpplloo p(yl @pw(aw)) log PNML [3][r)]
max Ey log pogr (Du(e.)) log PNMI. (y]x)]
On the final step, we substituted the definition of T(D,x) as the individual regret the NML distribution PNML-
13

---

Published as a conference paper at ICLR 2021
A.2 EXPERIMENTAL DETAILS
A.2.1
DETAILED PSEUDOCODE
There are 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 0 in order initialize the procedure with an accurate initial model.
Optimizing over batch of M designs, in order to follow previous evaluation protocols.
Optimizing with x with respect to the internal scores instead of E [g()]- Using target networks for the NMI. model, originally proposed reinforcement learning algorithms, to improve the stability of the method.
We present pseudocode for the practical implementation of NEMO below:
Algorithm NEMO Practical Instantiation
Input Model class O, Dataset D (21.N, WINN), number of bins K, batch size M, learning rates
ag, Ctz, target update rate т Initialize K models 0g:K
Initialize batch of optimization iterates Bo = xb) I:M from the best performing æ in D.
Pretrain 0!:K using supervised learning D.
Initialize target networks 0g. t- 0
Quantize y1.N into K bins, denoted [})] [[n)," [vk]]. for iteration in I.. d' do for 1. I... do
Sample o; from batch B₁.
Construct Augmented dataset: D' -
(z DU luk
Compute gradient + Vep Logikclihoodd D')
Update model 0% using Adam and gø with learning rate ag.
end for
Update target networks BY 0++1 < T0 + +1 +
for m in I....A .M do
Compute internal values y²² (x=))) from target networks 0f for all k E e, K Compute gradient for ax!": V Σx (x:"")
Update x using Adam and gradient IN with learning rate ag- end for end for
A.2.2 HYPERPARAMETERS
The following table lists the hyperparameter settings we used for each task. We obtained our hy- perparameter settings by performing a grid search across different settings of agg Qgz and т. We used 2-layer neural networks with sofiplus activations for all experiments. We used smaller net- works for GFP, Hopper, Ant, and DKitty (64-dimensional layers) and lower discretization K for computational performance reasons, but did not tune over these parameters.
Superconductor MoleculeActivity
GFP
Hopper Ant
DKitty
Learning rate
0.05
0.005
Learning rate &g
0.1
0.01
0.001
Network Size
256,256
64,64
Discretization K
40
20
Batch size
128
Target update rate
0.05
For baseline methods, please refer Anonymous (2021) for hyperparameter settings.
14

---

Published as a conference paper at ICLR 2021
A.3 ABLATION STUDIES
In this section, 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 ratio of model optimization steps input optimization steps. Each logging iteration in these figures corresponds to 50 loop iterations as depicted Algorithm A.2.1
A.3.1 EFFECT OF NML
In this section, present learning curves for NML, as well as an ablation study comparing NML. (orange curve) against forward optimization algorithm without NML (blue curve), labeled "No NML". The "No NML" algorithm is identical the NML algorithm detailed in Alg. [..2.I except the NML learning rate Og is set to 0.0. This means that the only model training that happens is done during the pretraining step. For illustrative purposes, initialize the iterates from the worst- performing scores in the dataset to better visualize improvement, rather than initializing from the best scores which used in our final reported numbers.
The scores on the Superconductor task are shown the following figure. Removing NML training makes very difficult for training most designs, shown by the poor performance the 50th percentile metric.
100th Percentie Score
No NML
NML
1000
50th Percentle Score
The the MoleculeActivity task follow similar trend.
100th Percentle Scorr
NML
50th Percentle Score
100
150
200
250
300
Iteraion
15

---

Published as a conference paper at ICLR 2021
And finally, the scores the GFP task also display the same trend.
100th Percensle Score
NML
NML
1000
50th Percantle Spore
200
400
600
800
1000
Beration
16

---

Published as a conference paper at ICLR 2021
A.3.2 ARCHITECTURE CHOICE
In this ablation study, we investigate the efficacy of the discretized logistic architecture. As a base- line, we compared against a standard feedforward network, trained with softmax cross-entropy loss to predict the discretized output y. We label this network as "Categorical", because the out- put of the network is 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 metrie, but the Categorical architecture displays less consistency in optimizing designs, as given by poor performance the 50th percentile metric.
100th Percentle Scor
Categorical
Logistic
50th Percentle Score
20
200
400
600
800
1000
Beration
On the ActivileActivity task, the Categorical network performs comparatively better, but still un- derperforms the discretized logistic architecture.
100th Percentle Score
Categorical
Logistic
50t Percentle Score
20
400
800
1000
Beration
17

---

Published as a conference paper at ICLR 2021
A.3.3 RATIO OF OPTIMIZATION STEPS
In this ablation study, investigate the effect of the ratio of model optimization steps input optimization step. For this experiment, we fix the learning rate to the hyperparameter values in Appendix. A.2.2 fix the input optimization steps and vary the number of model optimization steps we take.
We first investigate the Superconductor task, using a small model learning rate as 0.0005. In this setting, we see a clear trend that additional model optimization steps are helpful and increase convergence speed.
100h 100th Percentle Score
60
50th Percentle Score
steps
20
steps
steps
Boration
Using higher learning rate of ag = 0.05, smaller amount of steps works better, which suggests that it easy to destabilize the learning process using a higher learning rate.
100th 100th Percentie Score
5 steps staps 20 steps
40 steps
200
1000
50th Percentile Score
200
400
600
800
1000
Meratio
18

---

Published as a conference paper at ICLR 2021
A similar trend holds true in the MoleculeActivity task, albeit less pronounced. The following figure uses learning rate of agg = 0.0005, and we once again see that more model optimization steps leads to increased performance.
100th Percentie Score
steps
10 steps
20 staps
40 steps
200
600
800
1000
50th Percentie Score
200
600
1000
Itraaion
And using a learning rate of ag = 0.05, the advantage becomes less clear.
100th Porcentlee Score
200
1000
50th Percenile Score
0.5
steps
steps
steps
steps
Overall, while performed grid search over the learning rates achieve the highest performance, using a large number of model optimization steps with a small learning rate oo appears to be consistent strategy which performs well.
19


Reviews:
[{"review_id": "zubXksL0NPx", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Updated review\n----\n----\n\n# Summary\n\nThis work proposes an approach for model-based optimization based on learning a density function through an approximation of the normalized maximum likelihood (NML). This is done by discretizing the space and fitting distinct model parameters for each value. To lower the computational cost, the authors propose optimizing the candidates concurrently with the model parameters. Each model's distribution is encoded as a neural net outputting a scalar which is then encoded using a thermometer approach using a series of shifted sigmoid. Candidates are optimized based on the average value of the scalar of each model evaluated using parameters obtained from an exponentially weighted average of its most recent parameters.\n\n# Reason for score\n\nThis work proposes a reasonable approximation to an interesting estimator and demonstrate it is capable of achieving good consistent performance. This is likely to be of interest to the community and, as far as I'm aware, is sufficiently novel. Given that I see no noteworthy issues and all of my major concerns have been addressed, I don't see any reason for rejection. I strongly support acceptance.\n\n# Pros\n\n* Using estimates of the NML for model-based optimization is an interesting idea.\n* This work shows that the NML can be successfully approximated with a relatively coarse discretization and that both the optimization of the candidate and the various model parameters can be optimized in tandem. This suggests that this type of approach is viable and possibly warrants further investigation.\n\n\nInitial review\n----\n----\n\n# Summary\n\nThis work proposes an approach for model-based optimization based on learning a density function through an approximation of the normalized maximum likelihood (NML). This is done by discretizing the space and fitting distinct model parameters for each value. To lower the computational cost, the authors propose optimizing the candidates concurrently with the model parameters. Each model's distribution is encoded as a neural net outputting a scalar which is then encoded using a thermometer approach using a series of shifted sigmoid. Candidates are optimized based on the average value of the scalar of each model evaluated using parameters obtained from an exponentially weighted average of its most recent parameters.\n\n# Reason for score\n\nThere are a lot of typos and issues with the notation which make this paper unfit for publication in its current state. Otherwise, the work seems interesting but I find the experiments don't provide much insight. Though the comparison with the selected method is favorable, it's hard to consider them significant when there is a notable difference in only one of the three datasets of an unpublished benchmark. Additionally, the fact that the reported results only cover half the datasets from this benchmark raises some questions. Despite the negative tone of this paragraph, I want to note that the severity of some of these issues is subjective while others can be easily fixed. My mind isn't made up and I hope the authors can clarify anything I might have missed.\n\n# Pros\n\n* Using estimates of the NML for model-based optimization is an interesting idea.\n* This work shows that the NML can be successfully approximated with a relatively coarse discretization and that both the optimization of the candidate and the various model parameters can be optimized in tandem. This suggests that this type of approach is viable and possibly warrants further investigation.\n\n# Cons\n\n* The current notation is often confusing and even ambiguous at times. This will probably transpire in some of my other comments, but I do consider these issues to be superficial and easily fixed. I've provided some suggestions below for how to improve the notation. I understand that notation preference is a very subjective thing so the authors should feel free to opt for something different, but I do think the notation needs to be improved.\n* The \"thermometer encoding\" of the output seems like an odd choice, especially given how it is done here. If I understood the approach correctly, this seems to be a hacky way of using the output of the NN, $o_{int}$, as parameters to a logistic distributions. Why not treat it this way directly? My interpretation of this approach is that the $o^k_{int}$s are parameters for logistic distributions and the mean is optimized through the unnormalized probs by optimizing each $o^k_{int}$ directly. Would this be a fair description of the approach? I was expecting the discretization to be used directly to approximate the integral in eq. (2). \n* The experimental results aren't very conclusive, only showing a clear benefit in a single case. Without additional results, it is difficult to say much about the behavior and properties of this method. A visualization of the distribution of the value of the candidates might help convey some additional information in favor of this method. It's possible that I am missing some context to appreciate these results. If that is the case, it would help if the authors could provide the context I need to appreciate these results.\n\n# Questions\n\n* I don't think I understand what makes the appendix results an ablation study. From what I understand, these results only compare with the case where there is no learning of the model parameters. What are the models initialized to? Where does the data come into play?\n* Have the authors considered using points selected from some fixed quadrature method instead of a uniform grid?\n* A common theme for the comparison methods is the idea of not diverging too much from the data. Was the validity of the outputs of this method evaluated in any way? How can I know that the method isn't just exploiting some quirk in the learned models used to evaluate it while some of the other methods avoid doing this?\n* Were comparisons with a simple approaches Bayesian optimization tried, e.g., Gaussian process?\n* What happens when doing more iterations on the log likelihood before updating $x$? How much do we lose when only doing a single update? Does using a more accurate approximation of the NLM improve/worsen performance? (I was hoping this would be part of the ablation study)\n* How do the run times compare?\n* Were the other datasets from the Design-bench benchmark tried?\n\n# Misc and typos\n\n* page 2, \"in that it has shown to be\", missing word?\n* page 2, \"to discuss how approximate this intractable\", missing word?\n* page 3, when $p_{NML}$ is formally written, the meaning of $y$ is ambiguous since it is on the LHS and also being redefined by $\\max_y$ in the RHS.\n* page 3, \"The notation $D \\cup (x, y)$ refers to an augmented dataset $D \\cup (x_{N+1}, y_{N+1})$\", this wasn't very informative and felt a bit tautological. I would recommend sticking with one of the two notations.\n* page 3, \"where $D$ is fixed to the given offline dataset, and $\\theta_{D \\cup (x,y)}$ [...]\", this sentence is a bit confusing as a whole. The start of the sentence talks only about $D$ so the $\\theta$ mention is unexpected when reading.\n* page 3, right after eq. (2), $(x, y)$ is on the LHS but then also part of the expectation on the RHS. What is the expectation being taken over? Are $x$ and $y$ being redefined?\n* page 4, \"for y in 1 ... K do\", is $y$ I don't believe that $y$ is assumed to be an integer. \n* page 4, algorithm 4?\n* page 4,  \"this would produce a distribution over output values $p(y|x)$\", this doesn't \"type check\" for me. If I understood correctly, $p$ or $p(\\bullet | x)$ represent distributions and $y$ are output values. Also, it might be good to reuse the \"$\\hat p_{NML}$\" notation to get the following sentence: \"this would produce a conditional distribution, $\\hat p_{NML}(\\bullet | x_t)$,  over output values.\n* page 4, \"we can optimize $x_t$ with respect to some evaluation function $g(y)$ of this distribution, such as the\nmean\", this is confusing since algorithm 1 has $\\mathbb{E}[ g(y)]$. What does it mean for the evaluation function $g$ to be the mean in this context? How should I interpret the expectation of mean(y)? Also, I assumed that what was meant is that $g$ is the evaluation function, not $g(y)$, or is $g$ meant to return a function given a $y$?\n* page 6, \"a straightforward choice for the evaluation function g is the identity function $g(x) = x$\", it might be best to stick to a consistent variable name, e.g., $g(y) = y$, to avoid confusion about what the domain of $g$ is.\n* Proof of thm 4.1, equation under \"using these two facts, we can show:\", $q$ should be replaced with $p_{NML}$ in the RHS.\n* Proof of thm 4.1, going from TV to KL, looking up the bound returns a $1/2$ factor inside the root rather than a $2$. I could have missed a detail but thought worth mentioning in case I haven't.\n* Proof of thm 4.1, missing a \"]\" when bounding the KL divergence.\n* Algorithm 2, there is some undefined superscript $k$ and $y$ in the loop over $x^m_t$.\n\n# Notation suggestions\n\n* When writing expectations, I would strongly recommend making it explicit over which variables they are, e.g., $\\mathbb{E}\\_{y \\sim p(\\bullet | x, \\theta)}$. Introducing some shorthand notation might help make this more concise, e.g., $p_{x, \\theta}(y) := p(y | x, \\theta)$.\n* When referring to a function, only mention its name/symbol and reserve the form that includes inputs to refer to the output of the function given those inputs, e.g., a function $g$, an evaluation score $g(y)$.\n* Avoid relying on the readers pattern matching abilities for assigning meaning to variables and make sure variables, e.g., $\\mathbf{x}$ and $y$, are always explicitly defined. By explicitly defined, I include defining $y$ compactly with something like $\\max_{y \\in \\mathcal{Y}} g(y)$, for example. There is not need to be overly verbose but it should never leave room for interpretation. This is related to the point about expectation notation.\n* I usually prefer explicit domains, e.g., $\\sum_{y \\in \\mathcal{Y}}$ instead of $\\sum_y$, but I consider it fine to omit it if variables names are always reserved to the same domain when reused. This was not the case for $x$ in this paper.\n* When writing pseudocode, either loop over integer indices or over elements of a set, it is confusing to use \"for y in 1 ... K\" when $y$ isn't an integer. Additionally, using both makes it difficult to tell that $k$ is associated with $y$ inside a loop.", "extra_scores": {}, "extra_text": {}}, {"review_id": "fpGZAusvPdJ", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.   \n\nTypo: \nIn section 4, “We outline the high-level pseudocode in Algorithm 4” -> “…. in Algorithm 1”\n\nQuestions: \n1. When sampling a batch of data points at each step of algorithm 1, is the sampling performed with or without replacements?  \n2. What’s the variance across the 16 random runs? Is the score of the best algorithm in the average performance across 16 random runs significantly different from the baseline algorithms?\n3. When estimating the CNML, what is the number of models in the experiments? Are the number of models differ from dataset to dataset? How to choose the number of models in practice?\n4. Since the output y needs to be discretized in the NEMO algorithm, how difficult for the algorithm to scale when y is multivariate?\n\n------------------------------------------------------------------------------------------------------------\nUpdate: \nI think the authors did a great job of addressing my concerns, I'm happy to increase my score to 6\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "TBhMYJxLjnG", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# summary\n\nThis paper proposed a method based on NML and provided a principled approach\nto estimate the uncertainty for OOD data.\n\n\n# pros\n\n1.  The method proposed in this work is a principled way to handle uncertainty\n    for novel points out of the original dataset compared with, for example,\n    deep ensemble.\n2.  One clear advantage of ths proposed approach is this method can scale to\n    large dataset, compared with GP, which scales cubically.\n\n\n# cons\n\n1.  The authors claim using a supervised approach is brittle and prone to\n    failure as the uncertainty of inputs outside of training data is\n    uncontroled. However, this is not true and uncertainty can be well\n    controlled depending on the model, which can be non-parametric or\n    parametric, distribution-free or distribution-dependent. For example, to\n    measure uncertainty on novel point, GP could be viewed as the ground truth\n    under some conditions. My question is why not directly compare your\n    approach with a GP approach, then combine the posterior with an acquisition\n    function, such as EI. The offline MBO problem presented in this work is\n    similiar to an online MBO, except we have only one online sample, and\n    we are tying to optimize this point. Comparing eq(1) of this paper with the\n    formula of EI, it is easy to see eq(1) is exactly EI, if we assume there\n    exists (x\\*,y\\*) such that y\\* is larger than objective values in the training\n    data set. Thus the problem presented in this work can be effectively solved\n    through **one step** of conventional BO. Given the datasets used in the\n    experiments of this work is are not of large scale, I think comparing with\n    a GP-based BO is necessary.\n2.  In Figure 3, although not a major concern, I don't think the comparison\n    with the ensemble is fair. Although this work uses bootstrap and ensemble,\n    MSE cannot capture uncertainty, thus it is not an ideal metric in this\n    setting. For example, to obtain a similar uncertainty estimation compared\n    with NML (middle column), we can use a deep ensemble, which optimizes NLL\n    instead of MSE.\n3.  The experimental results, in my opinion, are not sufficient and there is\n    only one table presenting empirical results. I don't want to judge\n    sufficiency by the quantity of tables or figures, but considering the\n    theoretical analysis is not strong enough, I think more empirical study\n    should be performed.\n4.  The uncertainty estimation seems too conservative, and this could make the\n    estimated uncertainty less useful, especially for high-dimensional\n    problems.\n\n\n# questions\n\nThe approach proposed in this paper seems very similiar with conformal\nprediction. In conformal prediction, the target value y\\* for a novel point x\\*\n(adversarial input) is chosen so that y\\* is compatible with the original\ndataset. As I am not familiar with the evaluation protocol in Brookles2019 and\nFannjiang2020, the metric used in Table 1 is not clear to me. Can the authors\nsay more about that?\n\nupdate:\n\n---\nOverall speaking, the added GP-BO results address my  concerns, and I've updated the score from 5->6. A final update will be given later.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work proposes a model-based optimization using an approximated normalized maximum likelihood (NML). It is an interesting idea and has the advantage of scaling to large datasets. The reviewers are generally positive and are satisfied with authors' response.  
Decision: Accept (Poster)