OFFLINE MODEL-BASED OPTIMIZATION VIA NOR-MALIZED MAXIMUM LIKELIHOOD ESTIMATION

Abstract

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

1. INTRODUCTION

Many real-world optimization problems involve function evaluations that are the result of expensive or time-consuming process. Examples occur in the design of materials (Mansouri Tehrani et al., 2018 ), proteins (Brookes et al., 2019; Kumar & Levine, 2019) , neural network architectures (Zoph & Le, 2016) , or vehicles (Hoburg & Abbeel, 2014) . Rather than settling for a slow and expensive optimization process through repeated function evaluations, one may instead adopt a data-driven approach, where a large dataset of previously collected input-output pairs is given in lieu of running expensive function queries. Not only could this approach be more economical, but in some domains, such as in the design of drugs or vehicles, function evaluations pose safety concerns and an online method may simply be impractical. We refer to this setting as the offline model-based optimization (MBO) problem, where a static dataset is available but function queries are not allowed. A straightforward method to solving offline MBO problems would be to estimate a proxy of the ground truth function fθ using supervised learning, and to optimize the input x with respect to this proxy. However, this approach is brittle and prone to failure, because the model-fitting process often has little control over the values of the proxy function on inputs outside of the training set. An algorithm that directly optimizes fθ could easily exploit the proxy to produce adversarial inputs that nevertheless are scored highly under fθ (Kumar & Levine, 2019; Fannjiang & Listgarten, 2020) . In order to counteract the effects of model exploitation, we propose to use the normalized maximum likelihood framework (NML) (Barron et al., 1998) . The NML estimator produces the distribution closest to the MLE assuming an adversarial output label, and has been shown to be effective for resisting adversarial attacks (Bibas et al., 2019) . Moreover, NML provides a principled approach to generating uncertainty estimates which allows it to reason about out-of-distribution queries. However, because NML is typically intractable except for a handful of special cases (Roos et al., 2008) , we show in this work how we can circumvent intractability issues with NML in order to construct a reliable and robust method for MBO. Because of its general formulation, the NML distribution pro-

