Data-Driven Model-Based Optimization via Invariant Representation Learning
Anonymous Author(s)
AffiliationAddress
Abstract
We study the problem of data-driven model-based optimization (MBO), where thegoal is to find the optimal design, provided access to only a static dataset, with noactive data collection.
The central challenge in data-driven MBO is distributionalshift, where the optimizer is fooled into producing out-of-distribution (OOD)designs that erroneously appear promising under a model trained on the provideddata.
To address this issue, we formulate MBO as domain adaptation , where thegoal is to make accurate predictions for the value of designs during optimization(“target domain”), when training only on the dataset (“source domain”).
Thisperspective leads to invariant objective models (I OM ), our approach for addressingdistributional shift by enforcing invariance between the learned representations ofthe training dataset and optimized designs.
In I OM , if the optimized designs are toodifferent from the training dataset, the representation will be forced to lose much ofthe information that distinguishes good designs from bad ones, making all choicesseem mediocre.
Critically, when the optimizer is aware of this representationaltradeoff, it should choose not to stray too far from the training distribution, leadingto a natural trade-off between distributional shift and learning performance.
Introduction
Many real-world applications of machine learning involve learning how to make better designsfrom data.
This is formulated as data-driven model-based optimization (also known as black-boxoptimization [2]) in which an algorithm is provided with data from previous experiments consistingof ( x , y ) tuples, where x represents a design (e.g., a molecule [14] or robot controller [5]) and yrepresents its utility (e.g., certain property of the molecule or the effectiveness of a controller).
Thegoal is to generate the best possible design, typically one that is better than the best one in the dataset.
All of these problem domains are united by a common challenge: in order to improve, the designslearned by the algorithm must differ from the distribution of the designs in the training data, inducingdistributional shift [26], which must be carefully handled.
Typical methods for solving such data-driven model-based optimization problems try to learn someform of a proxy model ˆ f ( · ) that maps a design x to its objective value y via supervised learning onthe static dataset, and then optimize the design x against this learned model.
However, overestimationerrors in the learned model will erroneously drive the optimizer towards low-value designs that “fool”the proxy into producing a high values [26].
To handle this, prior works have proposed to regularizethe learned model to be conservative by penalizing its values on unseen designs [46, 51, 20], orexplicitly constraining the optimizer (e.g., by using density estimation) [8, 26].
In this paper, we take a different perspective for the distributional shift challenges in the model-basedoptimization, and formulate it as domain adaptation, where the “target” domain consists of thedistribution of designs that are chosen by the optimizer, and the source domain corresponds to thetraining set.
Framing the problem in this way, we derive a class of methods that regulate distributional shift at the representation level , and admit effective hyperparameter tuning strategies.
The intuitionbehind our approach is that, if the internal representations of the learned model are made invariantto the differences between the designs seen in the training data vs. the optimized distribution, thenexcessive distributional shift is naturally discouraged because excessively invariance would lead theout-of-distribution points to attain mediocre values on average.
That is, as the optimized decisionsdeviate more from the training data, the requirement to maintain invariance will cause the model’srepresentation to lose information, and will hence force the predicted value to become closer to theaverage y value in the dataset.
This will force the optimizer to stay within the distribution of thedataset, since out-of-distribution points will only be mediocre under the learned model.
Of course,the representation cannot be perfectly invariant in general, because the optimized points are differentfrom the training data, and hence the method must strike a balance.
In practice, one can instantiatethis idea using an adversarial training procedure.
Our main contribution is an offline, data-driven modelbased optimization method based on a novel connectionwith domain adaptation.
Our approach, I OM , optimizesthe design against a learned model that admits invariantrepresentations between optimized designs and the dataset.
This enables controlling distributional shift and also provides a way to derive effective offline hyperparametertuning strategies that we show actually work well in practice.
I OM can be implemented simply by combining anysupervised regression procedure with a regularizer thatenforces invariance.
The regularizer is implemented using any discrepancy measure between two distributions,and we instantiate I OM using the χ 2 -discrepancy measurevia the discriminator of a least-squares generative adversarial network.
This procedure trains a discriminator todiscriminate between the representations on the trainingdataset and optimized designs, whereas the learned representations are trained to be as invariant as possible to foolthe discriminator.
We theoretically derive I OM from ourdomain adaptation formulation of MBO and we also usethis formulation to derive tuning strategies for several common hyperparameters.
Empirically, weevaluate I OM on several tasks from Design-Bench [45] and find that it outperforms the best priormethods, and additionally, admits appealing offline tuning strategies unlike the prior methods.
Related Work
Model-based optimization (MBO) or black-box optimization [8, 3, 43, 42, 35, 51] has been studiedwidely in prior works.
Among these works, the central challenge is in handling the distributionalshift that arises between the distribution of designs seen in the dataset and those taken by theoptimizer [8, 11, 26].
For methods that learn a model, this typically manifests itself as modelexploitation , where the optimizer can find out-of-distribution designs for which the model erroneouslypredicts good values, analogously to adversarial examples.
This distributional shift is distinct fromwhat is studied in prior works that tackle misspecification in GPs and bandits [36, 37, 24] or methodsthat tackle robustness across contexts [25] in the online setting.
To address distribution shift, several recent works have proposed to be conservative.
This is majorlythe case in other decision-making problems offline reinforcement learning (RL) [32] and contextualbandits [38], where instead of optimizing for a design, we must find the optimal action at each state.
Inoffline MBO, prior works have tried to be conservative in two different ways: either by constraining theoptimized designs to lie close to the training designs [8, 11, 26, 39, 27] using generative models [23,15], or by learning a conservative learned model [46, 29, 28] or calibrated [12, 51].
While thelatter class of methods has worked well in practice [45, 29], these methods can sometimes beoverly conservative and are quite sensitive to hyperparameters.
Sensitivity to hyperparameters isacceptable as long as one can attain good performance using reasonable offline tuning strategies,without accessing the groundtruth function, but we find in our experiments (Section 6.2) this isnot easy for conservative methods.
On the contrary, I OM not only does not require a generativemodel over the design space, which can be intractable with high-dimensions (e.g., 5000-d designs in HopperController), attains better performance, and provides a convenient tuning rule.
In this work we formulate data-driven model-based optimization as domain adaptation [4, 13], andderive algorithms for model-based optimization from this perspective.
This perspective has beenused to devise decision-making methods for treatment effect estimation [21, 22, 50] and off-policyevaluation in RL [33, 31].
While these prior works do use an invariance regularizer, they primarilyfocus on obtaining more accurate estimations, whereas our goal is to utilize invariance to combatdistributional shift during optimization.
Thus, our usage of invariant representations is tailoredtowards the optimizer, instead of aiming for accurate predictions everywhere like these prior works,which is impossible in general.
Furthermore, our analysis uses invariant representations to deriveprocedures for tuning and early stopping, while these prior works do not explore these perspectives.
Preliminaries
The goal in offline model-based optimization (MBO) [45] is to produce designs that maximizesome objective function f ( x ) , by only utilizing a dataset D
= { ( x i , y i ) } ni =1 of designs x i andtheir corresponding objective values
y i .
No active queries to the groundtruth function are allowed.
Typically offline MBO is formulated as finding the optimal design x ⋆ .
However, we will use a slightlymore general formulation, where the goal is to optimize a distribution over designs µ ( x ) .
We willuse this formulation largely for convenience of analysis, where the classic case is recovered whenµ ( x ) is a Dirac delta function, though we also note that optimizing distributions over designs is oftenstudied in settings such as synthetic biology, where the goal is to explicitly diversify the designs [6].
Thus, our goal is to find a distribution over designs µ O PT ( x ) using only the static dataset D , such thatit maximizes the expected value of f ( x ) (the expected value under distribution µ is referred to as J ( µ ) ).
This can be formalized as follows ( Alg is an algorithm):
[Equation]
The methods we consider in this paper fit a learned model f θ ( x ) to ( x i , y i ) pairs in the dataset via astandard supervised regression objective, with an additional regularizer.
We will assume that f θ iscomposed of two components: a representation ϕ ∈ R d that maps a given x into a d -dimensionalvector, and a forward model, f θ that converts the representational output into a scalar objective value,i.e., f ϕ ( x ) :
= f θ ( ϕ ( x )) .
We will use the notation f θ and f ϕ interchangeably.
Given a learned modelf θ ( ϕ ( x )) , one approach to obtain the optimized distribution µ O PT is by optimizing the expected value
J θ ( µ ) : =
E x ∼ µ [ f θ ( ϕ ( x ))] under the learned model.
This distribution µ O PT ( x ) can be instantiatedin many ways: while some prior work [26, 8] represents µ O PT via a parametric density model over
X and optimizes this density model to obtain µ O PT , other prior work [46] uses a non-parametricrepresentation for µ O PT using a set of optimized “design particles”.
In this paper, we represent µ O PTusing a set of design particles, each of which is obtained by optimizing the learned function f θ ( ϕ ( x ))with respect to the input x via grad. ascent, starting from i.i.d.
samples from the dataset.
Formally:
[Equation]
The optimized distribution µ O PT is then given by µ O PT =
δ { x 1 T , x 2 T , · · · } .
The challenge of distributional shift in offline MBO [26, 9, 11, 46].
When the learned modelf θ ( ϕ ( x )) is trained via supervised regression on the offline dataset D : arg min θ 1 n (cid:80) ni
{+=
+} ( f θ ( ϕ ( x i )) −y i ) 2 , the ERM principle states that f θ ( ϕ ( x )) will reasonably approximate f ( x ) in expectation onlyunder the distribution of the training dataset D , i.e., µ data , however f θ ( ϕ ( x i ))
̸
= y i in general onpoints not in the training distribution. Then, when optimizing the learned function, overestimationerrors in the value of f θ ( ϕ ( x )) will inevitably lead the optimizer to converge to a distribution
µ O PTover designs that have a high value under the learned function (i.e., high J θ ( µ O PT ) ) but not a highground truth value J ( µ O PT ). In general, this distribution µ O PT will be different from µ data , since thelearned function is likely to be more correct in expectation under µ data .
Data-Driven Offline Model-based Optimization as Domain Adaptation
In this section, we will formulate data-driven offline MBO as domain adaptation, which will allow usto derive our method, invariant objective models (I OM ).
Instead of attempting to simply constrainthe optimizer to directly avoid distributional shift, our approach modifies the representations learnedby the model ˆ f θ by adding an invariance regularizer during training such that there is minimaldistribution shift under the learned invariant representation during optimization.
If we can ensure that the learned model f θ is accurate on the distribution found by the optimizer, µ O PT ,then we can ensure that we will obtain good designs.
Therefore, to handle the distributional shift, we can treat µ data as the “source domain” and the optimized distribution µ O PT as the “target domain.”
Just as the goal in domain adaptation is to train on the source domain and make accurate predictionson the target domain, we will aim to train f θ ( x ) on µ and with the goal of making more accuratepredictions under µ O PT .
Of course, learning to make accurate predictions on any unseen distributionis impossible for any domain adaptation method [55], but in the case of MBO, the choice of the targetdistribution µ O PT is also made by the algorithm itself, such that invariance can be enforced both bychanging the representation and changing the target distribution µ O PT .
4.1 Invariant Objective Models (I OM )
I OM controls distributional shift between the optimized distribution µ O PT and the training distributionµ data at the representational level, by regularizing the distribution of the features ϕ ( x ) under µ O PT tomatch that under µ data .
We will first present the intuition behind the method, and then formalize it bysetting up a bi-level optimization problem that we will also analyze theoretically.
Intuition: We supplement this intuition with an illustration shown in Figure 10. As the optimizerdeviates from the training distribution µ , the invariance regularizer forces the distribution of optimizeddesigns to resemble the distribution over representations ϕ ( x ) under µ data , i.e., P µ O PT ( ϕ ( x ))
≈ 160 P µ data ( ϕ ( x )) .
When µ O PT is far from the training distribution (i.e., if D KL ( µ O PT , µ data ) is very high),it is easier to enforce invariance as it does not conflict with fitting f θ on the training distribution.
Thisrepresentational invariance in turn causes the expected value of the learned function f θ under µ O PT toappear (roughly) equal to the average objective value in the training dataset, i.e., E x ∼ µ O PT [ f θ ( ϕ ( x ))]
≈ 164 E x ∼ µ data [ f θ ( ϕ ( x ))] , if µ O PT is too far away from the training distribution.
As a result, the expectedvalue of the learned model f θ ( ϕ ( x )) under µ O PT , J θ ( µ O PT ) , will be worse than the value of the bestdesign within the training distribution (since the best design input sampled from µ data will attain ahigher objective value than the average value under µ data ).
This would act as a guard to prevent theoptimizer diverging too far away from the dataset, as out-of-distribution designs no longer appearpromising.
Iteratively regularizing ϕ to enforce invariance between µ O PT and µ until convergencewill restrict the optimizer from selecting erroneously overestimated out-of-distribution inputs.
Wewill now present a formalization of this intuition in the form of a bi-level optimization problem thataims to maximize a lower-bound on the expected objective value attained under µ O PT , and derive thetraining objectives for I OM using this bi-level optimization formulation.
However, do note that, ourpractical method does not aim to solve a bi-level optimization exactly.
Bi-level optimization for invariant objectivemodels (I OM ).
To enforce invariance at a representational level, we utilize an invariance regularizer.
Denoted as disc H ( p, q ) , our invarianceregularizer is a discrepancy measure betweentwo probability distributions p and q , w.r.t.
acertain hypothesis class H .
For example, if His the class of all 1 -Lipschitz functions, disc His the 1 -Wasserstein distance; if H is the classof linear functions, disc H is the maximum meandiscrepancy (MMD) [16] distance.
In practice,we implement this via a generative adversarialnetwork as we discuss in the next subsection.
Inaddition to training f θ via standard supervisedregression on the labels in the training data, I OM trains the representation ϕ ( x ) and the model (cid:98) f ( · )to additionally minimize the discrepancy disc H ( P µ ( ϕ ( x )) , P µ O PT ( ϕ ( x ))) between distributions ofrepresentations under the optimized distribution, µ O PT and the training distribution µ .
This objectivecan be formalized as:
[Equation]
where λ denotes a weighting hyperparameter.
In Section 4.2, we will describe how we can convertthis bi-level optimization problem into a practical MBO algorithm that can be implemented withhigh capacity deep neural networks.
This is followed by a theoretical analysis to formally justify the objective in Section 4.3, and a discussion of how the hyperparameter λ can be tuned against avalidation set, in Section 5, without any online access to the ground truth function.
4.2 Optimizing the Bi
Level Problem and the Practical I OM Algorithm
In this section, we will describe how we can convert the abstract bi-level optimization problem aboveinto an objective that can be trained practically and discuss some practical implementation details.
I OM alternates between solving the bi-level optimization problem with respect to the distribution µ O PTand the learned model f θ and ϕ independently.
For training the learned model and the representationf θ and ϕ , I OM solves the inner optimization by running one-step of gradient descent ( GD ), utilizingthe current snapshot of µ O PT (denoted µ t O PT ):
[Equation]
Simultaneously, we update the design particles representing µ t O PT :
= δ { x 1 t , x 2 t , · · · } towards optimizing the learned function f θ t [-+
-]+1 ( ϕ t +1 ( x )) :
x it +1 ← x it + η ∇ x f
[Equation]
1 ( ϕ t + [Equation] x ))
| x
= x it , ∀ i. Pseudocodefor I OM is shown in Algorithm 1.
Practical implementation details: We model the representation ϕ ( x ) and the learned functionf θ ( · ) each as two-hidden layer ReLU networks with sizes 2048 and 1024, respectively.
We utilizethe χ 2 -discrepancy, disc H in our experiments: disc ϕ ( p, q ) :
= χ 2 ([ ϕ ( x )] x ∼ p , [ ϕ ( x )] x ∼ q ) .
Thisdivergence can be instantiated in its variational form via a least-squares generative adversarial network(LS-GAN) [34] on the representations ϕ ( x ) .
Concretely, we train an LS-GAN discriminator todiscriminate between the representations ϕ ( x ) under the training distribution µ ( x ) and the optimizeddistribution µ O PT ( x ) , whereas the feature representations under these distributions are optimizedto be as invariant to each other as possible.
We also evaluated a variety of discrepancy measures,including the maximum mean discrepancy (MMD) [16], but found the χ 2 -divergence to be the bestperforming version across the board.
More implementation details can be found in Appendix A.2.
4.3 Theoretical Analysis of Invariant Representation Learning for MBO
In this section, we will formally show that under certain standard assumptions, solving the bi-leveloptimization problem (Equation 3) enables us to find a better design than the dataset distributionµ data , similar to analyses of offline RL algorithms.
To show this, we will combine theoretical tools foranalyzing offline RL algorithms [10, 28, 52] and domain adaptation [55].
We show that the expectedvalue of the ground truth objective under any given distribution π ( x ) of design inputs can be lowerbounded in terms of the expected value under the learned function modulo some error terms.
Then wewill show that these error terms are controlled tightly when solving the bi-level problem, and finallyappeal to uniform concentration to show that the performance of µ O PT is better with high probability.
Our goal isn’t to derive the tightest possible bound for I OM , but show that I OM can attain reasonableperformance guarantees, and this domain adaptation perspective can tackle distributional shift.
Notation and assumptions.
Following standard assumptions [31], we will analyze the setting whenthe learned representation lies in some space of functions, ϕ ∈ Φ , and the learned function f θ lies inf θ ∈ F .
We will consider that the optimized distribution µ O PT belongs to a function class, µ O PT ∈ Π .
Our analysis will require that the ground truth model f is approximately realizable under somerepresentation ϕ ∈ Φ and for some g ∈ F .
That is, we assume:
Assumption 4.1 ( ε − realizability) .
For any distribution π over designs, there exists a representationϕ ∈ Φ and a g ∈ F such that min g max π ∈ Π E x ∼ π [ | f ( x ) − g ( ϕ ( x )) | ] ≤ ε F , Φ .
In addition, we assume that ∀ f ∈ F , || f || ∞ < ∞ and is Lipschitz continuous with constant C :
||
f || L ≤ C .
This is true for most parameteric function classes, such as neural networks.
Then, wecan lower bound the average objective value under any distribution in terms of the invariance:
Proposition 4.2 ((Informal) Lower-bounding the ground truth value under π ) .
Under Assumption 4.1and the Lipschitz continuity on the function f θ , the ground truth objective for any given π ∈ Π canbe lower bounded in terms of the learned objective model f θ ∈ F and representation ϕ ∈ Φ , withhigh probability ≥ 1 − δ that:
[Equation]
where C F is a uniform constant depending on the function class F and ε stat refers to statistical errorthat decays inversely with the dataset size, |D| .
A proof for Proposition 4.2 is in Appendix A.1.
The ■ term can further be lower bounded using astandard concentration error bound for ERM, and can be controlled if we train f θ ( ϕ ( · )) minimizeprediction error against ground truth values of f .
Now we utilize Proposition 4.2 and a uniformconcentration argument to lower-bound the value under the optimized distribution µ O PT .
Proposition 4.3 ((Informal) Performance guarantee for I OM ) .
Under Assumption 4.1, the expectedvalue of the ground truth objective under µ O PT , J ( µ O PT ) is lower bounded by:
[Equation]
A proof of Proposition 4.3 can be found in Appendix A.1.
This proposition implies that as longas the optimizer finds designs that improve the learned objective function f θ ∈ F , such that ( ◦ ) ispositive, and the discrepancy term ( ⋆ ) is minimized, then the distribution over designs found by I OM ,µ O PT , will compete favorably against the distributions of designs in the dataset.
Of course, thesebounds present the worst-case guarantees for I OM , but as we will show in our experiments, I OM doesclearly improve over the best designs in the offline dataset.
In addition, this theoretical analysis alsomotivates the design of our scheme for performing hyperparameter tuning and model selection.
Thisscheme will be based directly on the result from Proposition 4.2.
We will discuss this connectionwhen discussing tuning in Section 5.
Finally, we would like to highlight how invariance leads to “average” values if µ O PT deviates too faraway from the data distribution µ .
If µ O PT deviates too far, such that disc H ( P µ data ( ϕ ) , P
[Equation]
( ϕ )) ≤ ε ′ ,the learned function is bounded too close to the average learned value on the data i.e., | (cid:98) J ( µ O PT ) −(cid:98)
J ( µ data ) | ≤ C ′ · ε ′ .
Thus, the learned function simply reverts to “predicting the mean objective value”of the training distribution as ε ′ decreases, thereby losing any information about the identity of theoptimized distribution, when the optimizer goes too far out of µ .
Offline Workflow and Tuning for I OM
A central challenge in offline model-based optimization is that of offline workflow and hyperparametertuning: not only must an algorithm be able to produce designs using a static dataset, but it must alsoprescribe a fully offline workflow for tuning hyperparameters so that the produced designs are as goodas possible.
This includes both explicit hyperparameters (such as the coefficient of conservatism orinvariance, λ for I OM ) and implicit hyperparameters (such as the early stopping point).
The primarychallenge stems from the fact that unlike supervised learning, the final performance in MBO isdetermined by the optimized distribution µ O PT , which is different from the training distribution, µ data ,and we must reason about this performance fully offline.
In this section, we discuss how I OM admits a particularly convenient and fully offline hyperparametertuning procedure that works well empirically.
The goal of our tuning problem is to find a tupleconsisting of the learned representation ϕ λ , the learned model, f λθ , and the optimized distributionµ λ O PT , out of a given set of tuples obtained for different values of λ .
In other words, we wish to findthe λ from the set of models M λ :=
( f λθ , ϕ λ , µ λµ O PT ) , such that the corresponding µ λ O PT attains thehighest value under the ground truth objective.
Crucially note that this must be done fully offline.
Our key idea: The primary question we wish to answer is: how can we identify if a given tuple( f λθ , ϕ λ , µ λ O PT ) leads to good performance?
In order for a tuple trained via I OM to perform well, weneed two primary conditions: (1) it should be nearly invariant, and as a result, robust to distributionalshift, and (2) the learned f λθ must attain good generalization performance within the training distribution.
Intuitively (1) guards us against distributional shift, and (2) enables us to find as good of alearned model on the dataset, that does not overfit or underfit as a result of the invariance regularizer.
This forms the basis of our tuning method.
Offline tuning of I OM :
To instantiate the idea practically, we first run I OM with a set of hyperparameters Λ = { λ 1 , λ 2 , · · · } .
Then for each run, we record the validation in-distribution error andthe value of the invariance regularizer on a validation set.
We then perform tuning in two steps:first, we filter all λ values that attain near-perfect invariance under a user-specified margin ε , i.e.,disc H ( P µ data ( ϕ λ ( x )) , P µ ω O PT ( ϕ λ ( x ))) ≤ ϵ , and are hence guaranteed to be robust to distributional shift.
Second, we now pick models that attain good performance within the training distribution by selectingλ values that attain the smallest validation prediction error: ( f θ ( ϕ ( x )) − y ) 2 , in addition to picking the early stopping point based on the smallest validation error.
This enables us to pick λ that givesrise to a generalizing model f λθ , ϕ λ , while being most robust to distributional shift.
This still leaves auser-defined margin ε as a free parameter, but we find that it is comparatively insensitive and canbe selected heuristically, as we show in our experiments in Section 6.2.
Formally, this procedurecorresponds to solving the following optimization:
[Equation]
Finally, for a more formal explanation of the soundness of this strategy, we note that applying
Equation 6 provides a tight lower bound on the ground truth objective value in Proposition 4.2.
Upto statistical error (which decays as |D| increases), the prediction error on a held-out validation setestimates ( ■ ) in Proposition 4.2, and a smaller discrepancy (i.e., disc H ≤ ε ) controls the ( ⋆ ) termin Prop. 4.2.
ε stat is statistical error, and ε F , Φ only depends on the function classes involved ( F , Φ ).
This tuning λ per Equation 6 can be justified as aiming to maximize a lower bound on J ( µ O PT ) .
Practical implementation.
When implementing this procedure in practice, we must first selectthe value of ε .
For obtaining perfect invariance, the value of ε must be equal to 0.25 (when thediscriminator is trained using a χ 2 -divergence least-squares GAN objective), and therefore, ε must beset close to 0.25.
However, in practice, instead of selecting ε , we can select a fraction of all possibleλ values that pass the filter.
In our experiments, we sort the runs based on the absolute differencebetween the value of invariance regularizer (i.e., the discriminator training objective) and 0.5, and letthe top 15% (top 3 out of the 7 λ values we tune on) of the runs be selected for the second step.
In Section 6, we validate this offline tuning strategy, both against other alternative offline strategiesfor tuning I OM as well as offline strategies for tuning prior methods (COMs [46] and gradientdescent [45]).
We find that our offline tuning strategy performs comparably to full online tuningwithout needing any online queries, whereas other strategies for I OM and other methods, lead tosignificantly poor performance compared to online tuning.
Experimental Evaluation
The goal of our empirical evaluation is to answer the following questions: (1) Does I OM outperformprior state-of-the-art methods for data-driven model-based optimization?
(2) Is our offline tuningstrategy for I OM better than other offline tuning strategies for I OM or other prior methods?
(3) Howsensitive is I OM to the choice of the hyperparameter λ ?
In this section, we will answer these questionsvia a comparative evaluation of I OM on several continuous data-driven design tasks from the standardtasks from the Design-Bench suite [45].
6.1 Empirical Evaluation on Benchmark Tasks
We compare I OM to a variety of prior methods for model-based optimization: CbAS [8], Autofocused CbAS [11] and MINs [26] that utilize generative models for constraining distributionalshift; COMs [46] and RoMA [51] train robust models of the objective function by via conservatism and smoothness, respectively; gradient-free optimization methods such as REINFORCE [48],
CMA-ES [18] and conventional Bayesian optimization BO-qEI [49].
We also compare the standardapproach of learning a naïve model of the objective function, which is then optimized via gradientascent.
To better understand the benefits of enforcing invariance, we also compare I OM to I OM -C,which additionally applies a conservatism term on top of invariance that pushes up the learned valuesof designs in the dataset.
Additionally, we add a prior method that optimizes designs against alower-confidence bound estimate computed using a Gaussian process posterior in Appendix A.3.
Tasks.
We evaluate on four tasks with continuous-valued input space from the Design-Bench [45]benchmark for offline model-based optimization (MBO).
Here we briefly summarize these tasks:(1)
Superconductor: this task requires optimizing the properties of superconducting materials tomaximize the operating temperature; the dataset consists of 17010 training points, and each inputx ∈ R 86 represents a serialization of the chemical formula.
(2) D’Kitty and (3) Ant robot morphologydesign: these tasks require selecting the optimal morphology of a quadrupedal robot, in order tomaximize how quickly it can walk or crawl in a particular direction; the morphologies are representedby x ∈ R 56 and x ∈ R 60 , respectively.
(4) Hopper controller: this task reframes the classic hoppertask [7] as a bandit problem, where instead of sequentially controlling the hopper, the goal is todirectly output the parameters of a neural network controller that maximizes hopping speed; notethat this task, as set up in Design-Bench, is not a RL problem, but rather a bandit problem with a5126-dimensional continuous action space corresponding to neural network weights.
Evaluation protocol.
Following prior work [45, 46, 51, 12], for each method, we query the learnedmodel to obtain the top performing N points, i.e., x 1 , x 2 , · · · , x N .
We report the maximum objectivevalue among these N samples.
This protocol reasonably reflects a real-world design process, where anumber of computationally produced designs are tested and the best one is used for deployment.
Hyperparameter tuning.
For fair comparison against prior methods studied in Trabucco et al.
[45],we follow the uniform tuning strategy for reporting results in this section: we run I OM on each taskfor a given set of λ values, and pick a single value of λ across every task.
Note that this does requireaccess to online evaluation, but prior works do use the same protocol.
We will evaluate the efficacyof our offline tuning approach in Section 6.2 extensively.
Results.
Following the convention proposed by Trabucco et al.
[45], we evaluate all methods in termsof normalized score (higher is better).
In Figure 2 (left), we show the median and interquantile mean(IQM)
[1] for the aggregated scores across all tasks.
Due to space limits, we only report the scores forthe top-performing baselines (See Appendix A.2 for the results for all baselines).
I OM significantlyoutperforms all of the prior methods, including state-of-the-art methods that employ conservatism,such as COMs [46].
This indicates that invariance alone (I OM ) can effectively handle distributionshifts and avoid over-estimation for out-of-distribution actions.
Additionally, a comparison betweenthe performance of I OM and I OM with conservatism (I OM -C) suggests that adding conservatism doesnot improve performance further.
This further corroborates the fact that invariance is effective inpreventing the optimizer from going out of distribution and conservatism is not needed.
That said,we would like to clarify that this result does not mean that invariance will always be better thanconservatism on any given offline MBO problem.
In order to understand why I OM outperforms COMson these tasks, we examine the in-distribution validation error of I OM and COMs in Appendix A.6and find that I OM learns less distorted functions that generalize better.
6.2 Analysis of the Offline Workflow for Tuning I OM
We will now present experiments to understand the effectiveness of our offline tuning scheme(Section 5) for tuning I OM .
It is important to highlight that, in general, offline tuning presentssignificant challenges for MBO.
Prior works in bandits and reinforcement learning [54] have suggestedthat offline tuning should not work at all and, to the best of our knowledge, no prior work in offline
MBO has proposed a principled and effective offline tuning method.
In this section, we empiricallyvalidate our tuning strategy in comparison to other potential strategies, and compare results from ourapproach to an oracle strategy that uses online evaluation of the true function.
Comparisons.
We compare our offline tuning strategy to various offline strategies: (1).
only tunes
I OM based on validation-set prediction error, analogous to supervised learning (called “predictionerror only”); (2).
a strategy that tunes I OM based on only the value of the invariance regularizer (called“invariance only”); (3).
a method that utilizes our tuning strategy for picking λ , but not for picking thecheckpoint (“uniform checkpoint selection”) and two offline tuning strategies for tuning other priormethods: (a) tuning COMs [46] using validation-set prediction error (called “prediction error only”),conservatism loss (called ”conservatism loss only”) and both prediction error and conservatismloss, one by one (called “two-step tuning”); and (b) tuning an objective model trained via gradientascent [45] using validation MSE error (“gradient ascent”).
We report the drop in performancerelative to the corresponding oracle online strategy that actually computes the “groundtruth” valuesof optimized designs in the simulator to choose hyperparameters, and is therefore not fully offline, i.e., the relative drop in performance of the method (I OM , COM, grad. ascent) when tuned offlineusing the particular strategy compared to oracle tuning for the same method.
Results.
We present our results in Figure 3 (right) comparing the efficacy of various tuning schemesacross various methods averaged over all tasks.
While offline tuning does lead to a reduction inperformance compared to oracle online tuning, we find that the drop is smallest for our offline tuningstrategy applied to I OM .
All other tuning strategies based on either only prediction error or theinvariance regularizer lead to worse performance.
This indicates that our tuning strategy is effectivein tuning I OM .
Furthermore, offline tuning strategies that utilize a validation set with prior methods,such as COMs and gradient ascent, also lead to a significant decrease in performance relative to theoracle.
This is quite appealing as it suggests that I OM is more amenable to offline tuning strategiesrelative to prior offline MBO algorithms.
Even if these prior methods can perform well when providedaccess to online evaluations, they might start to fail in real-world MBO problems when one needs totune them fully offline.
We discuss computational complexity of tuning the methods in Appendix A.6.
Ablation analyses.
Next we aim to study the sensitivity of our offline tuning strategy with respectto the non-rigorous, user-specified parameter ε in Equation 6 and checkpoint selection based on thevalidation prediction error.
Figure 2 (right) shows the offline tuning performance heatmap for all tasks(along their aggregated performance in the bottom row) under different choice of quantiles (wherequantile x % means the top x % of the models according to the value of the invariance regularizer, andthe given value of ϵ were chosen).
The color in each position of the matrix corresponds to the ratio ofperformance of our offline tuning strategy for a given quantile value, compared to that of uniformoracle tuning.
As shown in the bottom row, aggregated over all tasks, we can see our offline selectionstrategy is quite robust to the hyperparameter ϵ .
Figure 3 (left) evaluates the efficacy of our checkpointselection strategy, we plot the objective value J ( µ O PT ) attained as a function of the number of trainingepochs (i.e., checkpoints), alongside the values for the prediction MSE in the validation set for thecorresponding epochs.
The checkpoints selected by our early stopping scheme are shown via thevertical line.
Observe that indeed the checkpoint for the lowest validation MSE corresponds to arather good optimized distribution µ O PT for different tasks, no matter the optimal training epochshappen at the beginning or near the end, which validates our early stopping strategy.
More detailedresults showing the efficacy of early stopping for both I OM can be found in Appendix A.3.
Discussion
In this work we study offline data-driven model based optimization (MBO), where the goal is toproduce optimized designs, purely using a static dataset.
We proposed to view offline MBO throughthe lens of domain adaptation, and proposed a method that trains a model of the objective with anadditional regularizer to promote invariance between the representations learned by the model inexpectation under the training data distribution and the distribution of optimized designs.
We showthat this leads to a practically effective data-driven optimization method, and an appealing workflowfor selecting the hyperparameters with offline data.
Our method outperforms prior approaches, whichillustrates the practical benefits of invariance.
We also validate the effectiveness of the accompanyingtuning strategy.
While we evaluate our method on MBO problems, our approach can in principle alsobe extended to contextual bandits and even to, sequential reinforcement learning.
Limitations : (1)fully data-driven offline optimization may on its own be insufficient (due to limitations of stayingclose to the dataset, hardness of tuning) and
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] We cited DesignBench.
(b) Did you mention the license of the assets?
[Yes] Details are found in Appendix A.2.(c)
Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
[Table]
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
[Table]