Data-Driven Model-Based Optimization via Invariant Representation Learning
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We study the problem of data-driven model-based optimization, where the goal isto ﬁnd the optimal design, provided access to only a static dataset, with no activedata collection.
The central challenge in data-driven model-based optimization isdistributional shift, where the optimizer is fooled into producing out-of-distribution(OOD) designs that erroneously appear promising under a model trained on theprovided data.
To address this issue, we formulate model-based optimization asdomain adaptation , where the goal is to make accurate predictions for the value ofdesigns during optimization (“target domain”), when training only on the dataset(“source domain”).
This perspective leads to 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 choices seem 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 [12] 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 [22], 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 [22].
To handle this, prior works have proposed to regularizethe learned model to be conservative by penalizing its values on unseen designs [38, 42, 18], orexplicitly constraining the optimizer (e.g., by using density estimation)
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.
[Figure]
Our main contribution is an ofﬂine, 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 ofﬂine 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 a least-squares generative adversarial network.
Thisprocedure trains a discriminator to discriminate betweenthe representations on the training dataset and optimizeddecisions, whereas the learned representations are trainedto be as invariant as possible to fool the discriminator.
Wetheoretically derive I OM from our domain adaptation formulation of model-based optimization and we also use thisformulation to derive tuning strategies for several commonhyperparameters in I OM .
Empirically, we evaluate I OMon several tasks from Design-Bench [37] and ﬁnd that it outperforms the best prior methods, andadditionally, admits appealing ofﬂine hyperparameter tuning strategies unlike the prior methods.
Related Work
Model-based optimization (MBO) or black-box optimization [8, 3, 35, 34, 30, 42] 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, 9, 22].
For methods that learn a model, this typically manifests itself as modelexploitation , where the optimizer can ﬁnd out-of-distribution designs for which the model erroneouslypredicts good values, analogously to adversarial examples.
To address distribution shift, several recent works have proposed to be conservative or pessimistic.
This is majorly the case in other decision-making problems ofﬂine reinforcement learning (RL) [27]and contextual bandits [31], where instead of optimizing for a design, we must ﬁnd the optimal actionat each state.
In ofﬂine MBO, prior works have tried to be conservative in two generally differentways: either by constraining the optimized designs to lie close to the training designs [8, 9, 22, 32, 23]using generative models [21, 13], or by learning a “robust” version of the learned model that istrained to be conservative [38, 25, 24] or calibrated [10, 42].
While the latter class of methodshas worked well in practice [37, 25], these methods can sometimes be overly conservative and arequite sensitive to hyperparameters.
Sensitivity to hyperparameters is acceptable as long as one canattain good performance using reasonable ofﬂine tuning strategies, without accessing the groundtruthfunction, but we ﬁnd in our experiments (Section 6.2) this is not possible for some of these priorconservative methods.
On the contrary, I OM not only does not require a generative model, attainsbetter performance, and also provides for a convenient hyperparameter selection rule.
In this work we formulate data-driven model-based optimization as domain adaptation [4, 11], andderive algorithms for model-based optimization from this perspective.
This perspective has beenused to devise decision-making methods for treatment effect estimation [19, 20, 41] and off-policyevaluation in RL [28, 26].
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 ofﬂine model-based optimization (MBO) [37] 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 ofﬂine MBO is formulated as ﬁnding 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 ﬁnd a distribution over designs µ O PT ( x ) using only the static dataset D , such thatit maximizes the expected value of f ( x ) .
This can be formalized as follows ( Alg is an algorithm):
[Equation]
The methods we consider in this paper ﬁt 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  2 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 [22, 8] represents µ O PT via a parametric density model over
X and optimizes this density model to obtain µ O PT , other prior work [38] 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 objective functionf ✓ (  ( x )) with respect to the input x via gradient descent, starting from i.i.d.
samples from the dataset.
Formally:
[Equation]
The optimized distribution µ O PT is then given by µ O PT =
{ x [Equation] x
[Equation] [Equation] The challenge of distributional shift in ofﬂine MBO.
When the learned model f ✓ (  ( x )) is trainedvia supervised regression on the ofﬂine dataset D : arg min ✓ 1 n
[-=
-]P ni =1 ( f ✓ ( x i )  y i ) 2 , the ERMprinciple states that f ✓ (  ( x )) will reasonably approximate f ( x ) in expectation only under thedistribution of the training dataset D , i.e., µ data , however this may not hold out of the trainingdistribution.
When optimizing the learned function, overestimation errors in f ✓ (  ( x )) will inevitablylead the optimizer to converge to points ( µ O PT 6 = µ data ) that have a high value under the learnedfunction (i.e., high J ✓ ( µ O PT ) ) but not a high ground truth value J ( µ O PT ).
On the other hand, we showthat when using our method I OM , this simple procedure is sufﬁcient in ﬁnding good designs.
Data-Driven Ofﬂine Model-based Optimization as Domain Adaptation
In this section, we will formulate data-driven ofﬂine 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 modiﬁes 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 [44], 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 ﬁrst present the intuition behind the method, and then formalize it bysetting up a bi-level optimization problem that we will also analyze theoretically.
Intuition: As the optimizer deviates from the training distribution µ , the invariance regularizerforces the distribution of optimized designs to resemble the distribution over representations  ( x )under µ data , i.e., P µ O PT (  ( x ))
⇡ P µ data (  ( x )) .
Such an invariance can be enforced when µ O PT issufﬁciently far away from the training distribution.
This form of representational invariance in turncauses the expected value of the learned function f ✓ under µ O PT to appear (roughly) equal to theaverage objective value in the training dataset, i.e., E x ⇠ µ O PT [ f ✓ (  ( x ))]
⇡ E x ⇠ µ data [ f ✓ (  ( x ))] , if µ O PTis too far away from the training distribution.
As a result, the expected value of the learned modelf ✓ (  ( x )) under µ O PT , J ✓ ( µ O PT ) , will be worse than the value of the best design within the trainingdistribution (since the best design input sampled from µ data will attain a higher objective value thanthe average value under µ data ).
This would pull the optimizer closer to the training distribution,thereby acting as a guard against diverging too far away from the dataset.
Iteratively regularizing  toenforce invariance between µ O PT and µ until convergence will restrict the optimizer from selectingerroneously overestimated out-of-distribution inputs.
We will now present a formalization of thisintuition in the form of a bi-level optimization problem that aims to maximize a lower-bound on theexpected objective value attained under µ O PT , and derive the training objectives for I OM using thisbi-level optimization formulation. O b x ⇤ = K
D k Run optimize f ) : x i x ⇤ i + ⌘ r x f ✓ k (  k ( x ⇤ i )) .
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 H 178 is the class of all 1 -Lipschitz functions, disc H 179 is the 1 -Wasserstein distance; if H is the classof linear functions, disc H is the maximum meandiscrepancy (MMD) [14] 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 b f ( · ) 186 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 justifythe 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 PT and the learned model f ✓ and  independently.
For training the learned model and the representation f ✓ and  , I OM solves the inner optimization byrunning one-step of gradient descent ( GD ), utilizing the current snapshot of µ O PT (denoted µ t O PT ):(  t +1 , ✓ t +1 ) = GD ✓ ,  ; t n
[Equation]
Simultaneously, we update the design particles representing µ t O PT :
=  { x 1 t , x 2 t , · · · } towards optimizing the learned function f ✓ t {++
+} (  t +1 ( x )) :
x it + x
[Equation]
+1 ( x ))
| x
= x it , 8 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) [29] 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) [14], 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 ﬁnd a better design than the dataset distributionµ data , provably.
To show this, we will combine theoretical tools for analyzing ofﬂine RL algorithmsand domain adaptation [44].
Our goal would be to show that the expected value of the ground truthobjective under any given distribution ⇡ ( x ) of design inputs can be lower bounded in terms of theexpected value under the learned function modulo some error terms.
Then we will show that theseerror terms are controlled tightly when solving the bi-level problem, and ﬁnally appeal to uniformconcentration to show that the performance of µ O PT is indeed better with high probability.
Notation and assumptions.
Following standard assumptions [26], we will analyze the setting whenthe learned representation lies in some space of functions,  2  , and the learned function f ✓ lies inf ✓ 2 F .
We will consider that the optimized distribution µ O PT belongs to a function class, µ O PT 2 ⇧ .
226 Our analysis will require that the ground truth model f is approximately realizable under somerepresentation  2  and for some g 2 F .
That is, we assume:
Assumption 4.1 ( "  realizability) .
For any distribution ⇡ over designs, there exists a representation 2  and a g 2 F such that min  ,g max ⇡ 2 ⇧ E x ⇠ ⇡ [ | f ( x )  g (  ( x )) | ]  " F ,  .
In addition, we assume that 8 f 2 F , || f || 1 < 1 and is Lipschitz continuous with constant C :
||
f || L  C .
This is true for most parameteric function classes, such as neural networks.
Then, we 232 can 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 ⇡ 2 ⇧ canbe lower bounded in terms of the learned objective model f ✓ 2 F and representation  2  , with 236 high 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 ✓ (  ( · )) minimize prediction 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 ﬁnds designs that improve the learned objective function f ✓ 2 F , such that (  ) is 247 positive, 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 ofﬂine 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 µ O PT (  ))  " 0 ,the learned function is bounded too close to the average learned value on the data i.e., | b J ( µ O PT
[Equation]
b
J ( µ data ) |  C 0 · " 0 .
Thus, the learned function simply reverts to “predicting the mean objective value” 258 of the training distribution as " 0 decreases, thereby losing any information about the identity of theoptimized distribution, when the optimizer goes too far out of µ .
Ofﬂine Workﬂow and Tuning for I OM
A central challenge in ofﬂine model-based optimization is that of ofﬂine workﬂow and hyperparametertuning: not only must an algorithm be able to produce designs using a static dataset, but it must alsoprescribe a fully ofﬂine workﬂow for tuning hyperparameters so that the produced designs are as goodas possible.
This includes both explicit hyperparameters (such as the coefﬁcient 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 ﬁnal 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 ofﬂine.
In this section, we discuss how I OM admits a particularly convenient and fully ofﬂine hyperparametertuning procedure that works well empirically.
The goal of our tuning problem is to ﬁnd 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 ﬁndthe  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 ofﬂine.
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 ﬁnd as good of alearned model on the dataset, that does not overﬁt or underﬁt as a result of the invariance regularizer.
This forms the basis of our tuning method.
Ofﬂine tuning of I OM :
To instantiate the idea practically, we ﬁrst run I OM with a set of hyperparameters ⇤ = {  1 ,  2 , · · · } .
Then for each run, we record the validation in-distribution error and 284 the value of the invariance regularizer on a validation set.
We then perform tuning in two steps:ﬁrst, we ﬁlter all  values that attain near-perfect invariance under a user-speciﬁed 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 pickingthe 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 a user-deﬁned margin " as a free parameter, but we ﬁnd 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 justiﬁed as aiming to maximize a lower bound on J ( µ O PT ) .
Practical implementation.
When implementing this procedure in practice, we must ﬁrst 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 ﬁlter.
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 ofﬂine tuning strategy, both against other alternative ofﬂine strategiesfor tuning I OM as well as ofﬂine strategies for tuning prior methods (COMs [38] and gradientdescent [37]).
We ﬁnd that our ofﬂine tuning strategy performs comparably to full online tuningwithout needing any online queries, whereas other strategies for I OM and other methods, lead tosigniﬁcantly 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 ofﬂine tuningstrategy for I OM better than other ofﬂine 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 [37].
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 [9] and MINs [22] that utilize generative models for constraining distributional shift; COMs [38] and RoMA [42] train robust models of the objective function by via conservatism andsmoothness, respectively; conventional Bayesian optimization methods such as REINFORCE [39],
CMA-ES [16] and BO-qEI [40].
We also compare the standard approach of learning a naïve model ofthe objective function, which is then optimized via gradient ascent.
To better understand the beneﬁtsof enforcing invariance, we also compare I OM to I OM -C, which additionally applies a conservatismterm on top of invariance that pushes up the learned values of designs in the dataset.
Tasks.
We evaluate on four tasks with continuous-valued input space from the Design-Bench [37]benchmark for ofﬂine model-based optimization (MBO).
Here we brieﬂy 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 2 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 2 R 56 and x 2 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 [37, 38, 42, 10], 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 reﬂects 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.
[37],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 efﬁcacyof our ofﬂine tuning approach in Section 6.2 extensively.
Results.
Following the convention proposed by Trabucco et al.
[37], 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 signiﬁcantlyoutperforms all of the prior methods, including state-of-the-art methods that employ conservatism,such as COMs [38].
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.
6.2 Analysis of the Ofﬂine Workﬂow for Tuning I OM
We will now present experiments to understand the effectiveness of our ofﬂine tuning scheme(Section 5) for tuning I OM .
It is important to highlight that, in general, ofﬂine tuning presentssigniﬁcant challenges for MBO.
Prior works in bandits and reinforcement learning [43] have suggestedthat ofﬂine tuning should not work at all and, to the best of our knowledge, no prior work in ofﬂine
MBO has proposed a principled and effective ofﬂine 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 ofﬂine tuning strategy to various ofﬂine 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 forpicking the checkpoint (“uniform checkpoint selection”) and two ofﬂine tuning strategies for tuningother prior methods: (a) tuning COMs [38] using validation-set prediction error (called “predictionerror only”) and conservatism loss (called ”conservatism loss only”); and (b) tuning an objectivemodel trained via gradient ascent [37] using validation MSE error (“gradient ascent”).
We reportthe drop in performance relative to the corresponding oracle online strategy, i.e., the relative dropin performance of the method (I OM , COM, gradient ascent) when tuned ofﬂine using the particularstrategy compared to oracle online tuning for the same method.
Results.
We present our results in Figure 3 (right) comparing the efﬁcacy of various tuning schemesacross various methods averaged over all tasks.
While ofﬂine tuning does lead to a reduction inperformance compared to oracle online tuning, we ﬁnd that the drop is smallest for our ofﬂine 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, ofﬂine tuning strategies that utilize a validation set with prior methods,such as COMs and gradient ascent, also lead to a signiﬁcant decrease in performance relative to theoracle.
This is quite appealing as it suggests that I OM is more amenable to ofﬂine tuning strategiesrelative to prior ofﬂine 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 ofﬂine.
Ablation analyses.
Next we aim to study the sensitivity of our ofﬂine tuning strategy with respectto the non-rigorous, user-speciﬁed parameter " in Equation 6 and checkpoint selection based on thevalidation prediction error.
Figure 2 (right) shows the ofﬂine 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 ofﬂine 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 ofﬂine selectionstrategy is quite robust to the hyperparameter ✏ .
Figure 3 (left) evaluates the efﬁcacy 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 efﬁcacy of early stopping for both I OM can be found in Appendix A.3.
Discussion
[Figure] In this work we study ofﬂine data-driven model based optimization (MBO), where the goal isto produce optimized designs, purely using a static dataset.
We proposed to view ofﬂine datadriven model-based optimization through the lens of domain adaptation, and proposed a methodthat trains a model of the objective with an additional regularizer to promote invariance betweenthe representations learned by the model in expectation under the training data distribution andthe distribution of optimized points.
We show that this leads to appealing theoretical bounds, apractically effective data-driven optimization method, and an appealing workﬂow for selecting thehyperparameters with ofﬂine data.
Our method outperforms prior approaches, which illustratesthe practical beneﬁts of invariance.
We also validate the effectiveness of the accompanying tuningstrategy.
While we evaluate our method on model-based optimization problems, our approach can inprinciple also be extended to contextual bandits and even to, sequential decision-making problemssuch as ofﬂine reinforcement learning. Exploring these broad applications of invariant representationlearning for addressing distributional shift is an exciting avenue for future work. A limitation of theproblem setting we study more generally is that it may not be practically feasible to just produceoptimized designs in a fully data-driven manner and some online evaluation may be inevitable. Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] We discuss limitations of ourwork in Section 7. (c) Did you discuss any potential negative societal impacts of your work? [No] Our work
pertains to fundamental algorithm design for model-based optimization, and it does not
have any particular negative societal impacts beyond those that apply more generallyto model-based optimization methods.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] We discuss the
full set of assumptions in Section 4.3.
(b) Did you include complete proofs of all theoretical results?
[Yes] Our proofs are
included in Appendix A.1.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] We include it
with the supplementary material
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how theywere chosen)?
[Yes] We include these details in Appendix A.2.
(c) Did you report error bars (e
g., with respect to the random seed after running experiments multiple times)?
[Yes] We reported error bars in our results, and also followedguidelines for rigorous statistical evaluation using the protocol from Agarwal et al.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] These details can be found inAppendix A.2.
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 Design Bench.
(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 identiﬁable
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]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[-Assumption A.
-]
spent on participant compensation?
[N/A]