O FFLINE
M ODEL -B ASED
O PTIMIZATION
VIA
N OR -
MALIZED M AXIMUM L IKELIHOOD E STIMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
In this work we consider data-driven optimization problems where one must maximize a function given only queries at a ﬁxed 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.
I NTRODUCTION
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 ofﬂine model-based optimization (MBO) problem, where a static dataset is available but function queries are not allowed.
A straightforward method to solving ofﬂine 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-ﬁtting 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 provides a ﬂexible 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 ofﬂine MBO algorithm that utilizes a novel approximation to the NML distribution to obtain an uncertainty-aware forward model for optimization, which we call NEMO ( N ormalized maximum likelihood E stimation for M odel-based
O ptimization).
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 conﬁdence 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 conﬁdent and wrong in low-data regimes.
R ELATED W ORK
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 ofﬂine 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 ﬁnetuning 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 ofﬂine 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 conﬁdence 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 ofﬂine 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 ofﬂine 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 ofﬂine 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.
P RELIMINARIES
We begin by reviewing the problem formulation for ofﬂine model-based optimization, as well as necessary background on the normalized maximum likelihood estimator.
Problem statement .
We deﬁne the ofﬂine 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 ﬁnd the input value that maximizes the true function:
[Equation]
However, in ofﬂine MBO, the algorithm is not allowed to query the true function f ( y | x ) , and must ﬁnd the best possible point x ∗ using only the guidance of a ﬁxed dataset D = { x 1: N , y 1: 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 ﬁnd the optimum of the proxy function, ˆ x ∗ =
argmax x E y ∼ ˆ 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 p NML ( 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 deﬁned as Regret ind ( h, y )
= log p ( y | x , ˆ θ D∪ ( x ,y ) )
− log h ( y | x ) , and p NML ( y | x ) =
arg min h max y (cid:48)
Regret ind ( h, y (cid:48) )
(Fogel &
Feder, 2018).
The notation D ∪ ( x , y ) refers to an augmented dataset by appending a query point and label ( x , y ) , to a ﬁxed ofﬂine 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):
[Equation]
where ˆ θ D∪ ( x ,y ) =
arg max θ
[Equation]
( x ,y ) ∈D∪ ( x ,y ) log p ( y | x , θ ) is the maximum likelihood estimate for p using the dataset D augmented with ( x , y ) .
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).
Algorithm
Input Model class { f θ : θ ∈ Θ } , Dataset D =
( x 1: N , y 1: N ) , number of bins K , evaluation function g ( y ) , learning rates α θ , α x .
Initialize K models θ 1: K, optimization iterate x
Quantize y 1: N into K bins, denoted as (cid:98)Y(cid:99) = {(cid:98) y 1 (cid:99) , · · · (cid:98) y k (cid:99)} .
for iteration t in 1 . .
T do for k in 1 . . .
K do construct augmented dataset: D (cid:48) ← D ∪ ( x t , (cid:98)
y k (cid:99) ) .
update model: θ kt +1 ← θ kt + α θ ∇ θ kt LogLikelihood ( θ kt , D (cid:48) ) end for estimate CNML distribution: ˆ p NML ( y | x t ) ∝ p ( y | x t , θ yt ) /
(cid:80) k p ( (cid:98) y k
(cid:99)| x t , θ kt )
Update x : x t +1 ← x t + α x ∇ x E y ∼
ˆ p NML ( y | x ) [ g ( y )] end for
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 sufﬁcient for optimization on challenging problems, which we discuss in Section 4.
NEMO: N ORMALIZED M AXIMUM L IKELIHOOD E STIMATION FOR M ODEL -B ASED O PTIMIZATION
We now present NEMO, our proposed algorithm for high-dimensional ofﬂine 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 difﬁcult 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.
4. A N I TERATIVE A LGORITHM FOR M ODEL -B ASED O PTIMIZATION
We ﬁrst 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.
at each optimization step, conditioned on the current optimization iterate x t .
This would produce a conditional distribution p NML ( y | x ) over output values, and we can optimize x t 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 signiﬁcantly slower than standard optimization methods, because it requires ﬁnding 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 x t .
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 x t as well as a label y t,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 p NML ( y t | x t ) .
We then compute a score from the NML distribution, such as the mean, and perform one step of gradient ascent on x t .
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.
4. Q UANTIZATION AND A RCHITECTURE
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 ﬁnite.
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 ﬂooring each y value to the nearest bin (cid:98) y k (cid:99) , with the size of each interval deﬁned as
B = ( y max − y min ) /K .
While quantization has potential to induce additional rounding errors to the optimization process, we ﬁnd 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:
[Equation]
Discretized logistic architecture .
Quantization introduces unique difﬁculties into the optimization process for MBO.
In particular, quantization results in ﬂat 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 architecture, depicted in Fig.
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 ﬁnal 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 ﬁrst encoding a label y as a vector y disc , where y disc [ k ≤ bin ( y )]
=and elements y disc
[ 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 y mean ( x ) =
E y ∼ p ( y | x ) [ g ( y )]
= (cid:80) 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, y mean
= 1 K (cid:80) k
o [ k ]
Optimization The beneﬁt 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 ﬂat gradients introduced by the quantization procedure.
Optimizing with respect to µ is sensible as it shares the same global optimum as y mean , and gradients with respect to µ and y mean 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 y mean ( x ) denote the predicted mean.
Then,
If x ∈ arg max x µ ( x ) , then x ∈ arg max x y mean ( x ) .
For any x , (cid:104)∇ x µ ( x ) , ∇ x y mean ( x ) (cid:105) ≥ 0 .
Proof.
See Appendix.
A.1.2.
T HEORETICAL R ESULTS
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 difﬁcult 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 Θ .
Deﬁnition 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 deﬁne the functional regret of a distribution q as:
[Equation]
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 p NML be the conditional NML distribution deﬁned in Eqn.
Then,
[Equation]
Γ( D , x ) = log { (cid:80)
y p ( y | x, ˆ θ D∪ ( x ,y ) ) } is the minimax individual regret, and g max = max y ∈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 p NML .
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.
E XPERIMENTS
In our experimental evaluation, we aim to 1) evaluate how well the proposed quantized NML estimator estimates uncertainty in an ofﬂine setting, and 2) compare the performance of NEMO to a number of recently proposed ofﬂine MBO algorithms on high-dimensional ofﬂine MBO benchmark problems.
Our code is available at https://sites.google.com/view/nemo-anonymous
5. M ODELING WITH Q UANTIZED NML
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.
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 conﬁdent 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 conﬁdent estimates, even though they may be wrong.
5. H IGH -D IMENSIONAL M ODEL -B ASED O PTIMIZATION
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.
5.2. T ASKS
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 ﬂuorescent protein (GFP) task is to design a protein with high ﬂuorescence, based on work proposed by Sarkisyan et al.
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 ﬂuorescence 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 ﬁngerprints.
This dataset contains a total of 4216 data points.
The ﬁnal 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.
5.2. B ASELINES
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 training K 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 ﬁt 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 ofﬂine 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.
5. R ESULTS AND D ISCUSSION
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 signiﬁcant 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 beneﬁcial 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. [Table] [Table]
C ONCLUSION
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 ofﬂine 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 ofﬂine reinforcement learning.
A.1. P ROOF OF P ROPOSITION 4.
This statement directly from monotonicity.
Because y mean is a sum of monotonic functions in µ , y mean 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 (cid:104)∇ x µ, ∇ x y mean (cid:105) =
(cid:68) ∇ x µ
[Equation]
dy mean dµ
[Equation]
x µ (cid:69) =
dy mean dµ ||∇ x µ || 22 ≥ 0 .
A.1.
P ROOF OF T HEOREM 4.
In this section we present the proof for Thm.
4.1, restated below:
[Equation]
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 f max = max x f ( x ) :
[Equation]
Second, Fogel & Feder (2018) show that the NML distribution obtains the best possible minimax individual regret of
[Equation]
Using these two facts, we can show:
[Equation]
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.
[Equation]
On the ﬁnal step, we substituted the deﬁnition of Γ( D , x ) as the individual regret of the NML distribution p NML .
A. E XPERIMENTAL D ETAILS A.2. D ETAILED P SEUDOCODE
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 E p NML [ 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:
Input Model class Θ , Dataset D
= ( x 1: N , y 1: N ) , number of bins K , batch size M , learning rates α θ , α x , target update rate τ
Initialize K models θ 1: K
Initialize batch of optimization iterates B 0 = x 1:
Mfrom the best performing x in D .
Pretrain θ 1: Kusing supervised learning on D .
Initialize target networks ¯ θ 1: K← θ 1: K.
Quantize y 1: N into K bins, denoted as (cid:98)Y(cid:99) = {(cid:98) y 1 (cid:99) , · · · (cid:98) y k (cid:99)} .
for iteration t in 1 . .
T do for k in 1 . . .
K do
Sample x (cid:48) t from batch B t .
Construct Augmented dataset: D (cid:48) ← D ∪ ( x (cid:48) t , (cid:98) y k (cid:99) )
Compute gradient g θ ← ∇ θ yt LogLikelihood ( θ yt , D (cid:48) )
Update model θ yt using Adam and g θ with learning rate α θ .
end for
Update target networks ¯ θ yt
+
A.2. H YPERPARAMETERS
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.
[Table]
Discretization
[Table] [Table]
[Table]
For baseline methods, please refer to Anonymous (2021) for hyperparameter settings.
A. A BLATION S TUDIES
In this section, we present 3 ablation studies.
The ﬁrst 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 ﬁnal ablation study investigates the ratio of model optimization steps to input optimization steps.
Each logging iteration in these ﬁgures corresponds to 50 loop iterations as depicted in Algorithm A.2.1.
A.3. E FFECT OF NML
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 ﬁnal reported numbers.
The scores on the Superconductor task are shown in the following ﬁgure.
Removing NML training makes it very difﬁcult for training on most designs, as shown by the poor performance on the 50 th percentile metric.
The scores on the MoleculeActivity task follow a similar trend.
[Table]
A.3. A RCHITECTURE C HOICE
In this ablation study, we investigate the efﬁcacy 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 100 th percentile metric, but the Categorical architecture displays less consistency in optimizing designs, as given by poor performance on the 50 th percentile metric.
On the MoleculeActivity task, the Categorical network performs comparatively better, but still underperforms the discretized logistic architecture.
A.3. R ATIO OF O PTIMIZATION S TEPS
In this ablation study, we investigate the effect of the ratio of model optimization steps to input optimization step.
For this experiment, we ﬁx the learning rate α x to the hyperparameter values in
Appendix.
A.2.2, ﬁx the input optimization steps to 1, and vary the number of model optimization steps we take.
We ﬁrst 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 ﬁgure 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.