Bidirectional Learning for Offline Infinite-width Model-based Optimization
Authors: Can Chen
Yingxue Zhang
Jie Fu
Xue Liu
Mark Coates
Venue: NeurIPS 2022 Accept (2022)
Keywords: offline model-based optimization
infinite-width neural network
neural tangent kernel
bi-level optimization
out-of-distribution

Abstract: In offline model-based optimization, we strive to maximize a black-box objective function by only leveraging a static dataset of designs and their scores. This problem setting arises in numerous fields including the design of materials, robots, DNAs, proteins, etc. Recent approaches train a deep neural network (DNN) model on the static dataset to act as a proxy function, and then perform gradient ascent on the existing designs to obtain potentially high-scoring designs. This methodology frequently suffers from the out-of-distribution problem where the proxy function often returns adversarial designs. To mitigate this problem, we propose $\textit{\textbf{B}i\textbf{D}irectional learning for offline \textbf{I}nfinite-width model-based optimization}~(\textbf{BDI})$. BDI consists of two mappings: the forward mapping leverages the static dataset to predict the scores of the high-scoring designs, and the backward mapping leverages the high-scoring designs to predict the scores of the static dataset. The backward mapping, neglected in previous work, can distill more information of the static dataset into the high-scoring designs, which effectively mitigates the out-of-distribution problem. Yet, for a finite-width DNN model, the loss function of the backward mapping is intractable and only has an approximate form, which leads to a significant deterioration of the design quality. We thus adopt an infinite-width DNN model and propose to employ the corresponding neural tangent kernel to yield a closed-form loss for more accurate design updates. Experiments on various tasks verify the effectiveness of BDI. The code is available [here](https://github.com/GGchen1997/BDI).

Bidirectional Learning for Offline Infinitewidtth Model-based Optimization
Can (Sam) Chen'; Yingxue Zhang?, Jie Fu', Xue Liu', Mark Coates
McGill University," 2 Huawei Noah's ' Ark Lab, Beijing Academy of Artificial Intelligence chen@mail thee@maiil.mcilll.... 11. ca, yingxue wwwngghuaweii.ccooommm com, fujielbaai. ac. cn, xuel meeliudcs. mogill........ ca, mark www.ss@mcgglll.ccooommm
Abstract
In offline poddll-based optimization, we strive to maximize a black- box objective function by only leveraging a static dataset of designs and their scores. This prob- lem setting arises in numerous fields including the design of materials, robots, DNA sequences, and proteins. Recent approaches train a deep neural network (DNN) on the static dataset to act as proxy function, and then perform gradient ascent on the existing designs to obtain potentially high-scoring designs. This methodology frequently suffers from the nut-of-distribution problem where the proxy function often returns poor designs. To mitigate this problem, we propose BiDirectional learning for offline Infinite-width model-based optimization (BDI). BDI consists two mappings: the forward mapping leverages the static dataset to predict the scores of the high-scoring designs, and the backward mapping leverages the high- scoring designs to predict the scores of the static dataset. The backward mapping, neglected in previous work, distill more information from the static dataset into the high-scoring designs, which effectively mitigates the out-of-distributio problem. For finite-width DNN model, the loss function of the backward mapping is intractable and only has an approximate form, which leads to a significant dete- rioration of the design quality. We thus adopt infinite-width DNN model, and propose to employ the corresponding neural tangent kernel to yield a closed-form loss for more accurate design updates. Experiments on various tasks verify the effectiveness of BDI. The code is available here
Introduction
Designing new object or entity with desired properties is a fundamental problem in science and many works rely on interactive access with the unknown objective function to propose new designs. However, in real-world scenarios, evaluation of the objective function often expensive or dangerous (2 and thus it more reasonable to assume that we only have access to static (offline) dataset of designs and their scores. This setting is called offline model-based optimization. We aim to find a design to maximize unknown objective function by only leveraging the static dataset.
engineering [I] From material design [2] protein design
common approach to solve this problem to train deep neural network (DNN) model on the static dataset to act proxy function parameterized as f.(.). Then the new designs are obtained by performing gradient ascent on the existing designs with respect to fe(-). This approach attractive because it can leverage the gradient information of the DNN model to obtain improved designs. Unfortunately, the trained DNN model is only valid near the training distribution and performs poorly
"Corresponding author.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

outside of the distribution. More specifically, the designs generated by directly optimizing fa(-) are scored with erroneously high values 7 g As shown in Figure I the proxy function trained on the static dataset 21,2,3 overestimates the ground truth objective function, and the seemingly high-scoring design Pgrad obtained by gradient ascent has a low ground truth score Pgrad
problem by imposing effective inductive biases on the DNN model. The method in [8] Iearns a proxy function (- that lower bounds the ground truth scores on out-of-distie distiiiibiittoooon designs. The approach in [9] bounds the distance between the proxy function and the objective function by normalized maximum likelihood. In 7 Yu et al. use local smoothness prior to overcome the brittleness of the proxy function and thus avoid poor designs.
Recent works 7 address the out-of-distributi
As demonstrated in Figure [ these previous works try to better fit the proxy function to the ground truth function from model perspective, and then obtain high-scoring designs by gradient ascent regarding the proxy function. In this pa- per, we investigate method that does not focus on regularizing the model, but instead ensures that the proposed designs can be used to predict the available data. As illustrated in Figure [ by ensuring the high-scoring design Pours can pre- dict the scores of the static dataset P1,2,3 (back- ward) and vice versa (forward), Pours distills more information from P1,2,3- which makes Poura more aligned with P1,2,3, leading to Pours with a high ground truth score.
1.50
proxy
Pared
1.25
Dours previous work
backward
1.00
ground
Score
P Pours
0.75
P Pgrad
forward
0.50
P3
0.25
-0.5
0.0
0.5
Design
We propose BiDirectional learning for offtine Infinitewwidth model-based optimization (BDI)
Figure 1: Illustration of motivation.
between the high-scoring designs and the static dataset (a.k.a. low-scoring designs). BDI has forward mapping from low-scoring to high-scoring and a backward mapping from high-scoring to low-scoring. The backward mapping means the DNN model trained on the high-scoring designs expected to predict the scores of the static dataset, and vice versa for the forward mapping. To compute the prediction losses, we need to know the scores of the high-scoring designs. Inspired by [10], where a predefined reward (score) is used for reinforcement learning, we set a predefined target score for the high-scoring designs. By minimizing the prediction losses, the bidirectional mapping distills the knowledge of the static dataset into the high-scoring designs. This ensures that the high-scoring designs are more aligned with the static dataset and thus effectively mitigates the out- of-distribution problem. For a finite-width DNN model, the loss function of the backward mapping is intractable and can only yield an approximate form, which leads to significant deterioration of the design quality. We thus adopt a DNN model with infinite width, and propose to adopt the corresponding neural tangent kernel (NTK) [2 to produce closed-form loss function for more accurate design updates. This is discussed in Section 4.5 Experiments on multiple tasks in the
design-bench [ verify the effectiveness of BDI.
To summarize, our contributions are three-fold:
We propose bidirectional learning between the high-scoring designs and the static dataset, which effectively mitigates the out-of-distribution problem.
To enable a closed-form loss function, we adopt an infinite-widhh width DNN model and introduce NTK into the bidirectional learning, which leads to better designs.
We achieve state-of-the-an results various tasks, which demonstrates the effectiveness of BDI.
2 Preliminaries
2.1 Offline Model-based Optimization
Offline model-based optimization is formally defined as:
x arg max fxx), 25
(I)

---

where the objective function f() is unknown, but we have access to a size N dataset D {(x) yy)).... {(Ex, yN)], where æ; represents certain design and yi denotes the design score. The design could, for example, be drug, an aircraft or a robot morphology. Denote the feature dimension of design a: as D and then the static dataset 'D can also have a vector representation (X,: y1) where X, E RNNDD V1 E R..
common approach is fitting DNN model fo(.) with parameters 0 the static dataset:
N
0- = arg mim (fo(x:)) - -v:)
(2)
0 N
After that, the high-scoring design x" can be obtained by optimizing x against the proxy function Se- (a)) by gradient ascent steps:
L++1 = Dt + VV z=0, for t E [O, T 1], 1],
where the number of steps. The high-scoring design x" can be obtained as 2T. This simple gradient ascent method suffers from the out-of-distributie problem the proxy function fe (z) is not accurate for unseen designs, and its actual score is usually considerably lower than predicted g
2.2 Infinitely Wide DNN and Neural Tangent Kernel
Deep infinitely wide neural network models have recently received substantial attention. Previous works establish the correspondence between infinitely wide neural networks and kernel methods [[I- 14.. the infinite width assumption, wide neural networks trained by gradient descent with a squared loss yield a neural tangent kernel (NTK), which is specified by the network architecture. More specifically, given two samples (designs) æ; and æj, the NTK k(x,, measures the similarity between the two samples from the perspective of the DNN model.
The NTKs have outperformed their finite network counterparts in various tasks and achieved state-of- the-art performance for classification tasks such as CIFAR10 [15]. The approach in [16] leverages the NTK distill large dataset into a small one to retain as much of its information as possible and the method in 17 uses the NTK to learn black-box generalization attacks against DNN models.
3 Method
In this section, present BiDirectional learning for offline Infinite-widih model-based optimiza- tion (BDI). First, to mitigate the out-of-distriuttin problem, we introduce bidirectional learning between the high-scoring designs and the static dataset in Section For finite-width DNN model, the loss function is intractable. Approximate solutions lead to significantly poorer designs. We therefore propose the adoption of an infinite-width DNN model, and leverage its NTK to yield closed-form loss and more accurate updates in Section 3.2
3.1 Bidirectional Learning
We use X₁, MxD to represent the nigh-scoring designs, where M represents the number of high-scoring designs. The high-scoring designs X₁, are potentially outside of the static dataset distribution. To mitigate the out-of-distribution problem, we propose bidirectional learning between the designs and the static dataset, which leverages one to predict the other and vice versa. This can distill more information from the static dataset into the high-scoring designs. To compute the prediction losses, we need to know the scores of the high-scoring designs. Inspired by [10], in which a reward (score) is predefined for reinforcement learning, we set a predefined target score ys for every high-scoring design. More precisely, we normalize the scores of the static dataset to have unit Gaussian statistics following [ and then choose y/ to be larger than the maximal score in the static dataset. We set yı, 10 across all tasks but demonstrate that BDI is robust to different choices in Section 4.6 Thus, the high-scoring designs can be written as (X,, yr.). and the goal is to find X,.-
Forward mapping. Similar to traditional gradient ascent methods, the forward mapping leverages the static dataset to predict the scores of the high-scoring designs. To be specific, the DNN model (-) trained on (Xi, yt) is encouraged to predict the scores yo, given X Thus the forward loss
3

---

function can be written as:
Link (X1) = M 11.11 fb- (Xn))13;
(4)
where 0 is given by
8 1012.
0- arg min l1 (xx))1
(5)
where 3>0 a fixed regularization parameter. Then we can minimize Cr2h(X, against X, to update the high-scoring designs. Li2s ) is similar to the gradient ascent process in Eq.3) the only difference is that gradient ascent aims to gradually increase the score prediction whereas minimizing Lizh (x)) pushes the score prediction towards the (high) predefined target score !/--
Backward mapping. Similarly, the DNN model (+) trained on (Xx, yn. should be able to predict the scores yi given X₁. The backward loss function
(x))
N 1 fgr (x.) (x,))
(6)
where 0* (*xx) given by
1 M 3/,,
B M
0¹ ('h,,) = arg min
f0 f/ (1
1012.
Overall loss. The forward mapping and the backward mapping together align the high-scoring designs with the static dataset, and the BDI loss can be compactly written as:
(( L12n (Xx) Sn2)(Xn),
where we aim to optimize the high-scoring designs Xn.
3.2 Closed-form Solver via Neural Tangent Kernel
For finite-width DNN model fb (-). the high-scoring designs X,, of Eq. only exist in 0* (())) from Eq.(7). which is intractable and does not have closed-for form solution. The solution of Eq. is often approximately obtained by a gradient descent step [18 19),
0 yk
(X)))2
0 (x)) 0
M
00
where 7) denotes the learning rate. This approximate solution 0" (X,, leads to inaccurate updates to the designs X, through Eq.(6) and leads to much poorer designs, especially for high-dimensional settings, as illustrate in Section 4.5
We thus adopt DNN model (o(() with infinite width, and propose to leverage the corresponding
NTK to produce closed-form 0 (X,) [Л]22] and then closed-form loss function of 6q.
Neural tangent kernel. Denote by k( 2;,) the kernel function between æ; and Øj introduced by the infinite-wice width DNN fo(.). We use Kx,x, E RNNN Kx,x, RMM Kx,x, E NxM and K Kx,X,= E R MXN represent the corresponding covariance matrices induced by the kernel function k(a;;, (.))) By leveraging the closed-form 0* (X,) given by II 12] we can compute the closed-form L12 xx) in 6
Li (Xx)
1 Kxix, (Kx.x, + 2 N
(10)
Similar to 20 to focus on higher-scoring designs in the static dataset, we assign a weight to every design based its score and recompute the loss:
Liz1 (X.)
||w
(3 Kx.x, (xxx, + (1)
(11)
where Wi = / ////smaxxxxaappp (azi) represents the design weight vector and the weight parameter a:0 is a constant. Similarly, we have the forward loss function,
Cizn (X)) || (3h Ex,x((Xxx
(12)

---

where W) 1 V M
since every element in yn is the same in our paper.
Besides the advantage of the closed-form computation, introducing the NTK into the bidirectional learning can also avoid the expensive high-order derivative computations brought by Eq. (). which makes our BDI efficient first-order optimization method.
Analysis of M=1. When we aim for one high-scoring design (M 1), Eq.
has simpler form:
Lh21( = ||اا (y) k x.m,n (Kim,,
+ 3) +33)-'mm) yn.
N
A(Xi;, Bhh)
(13)
k + 3/.)
where k (xs 24)) measures the similarity between the ith design of X and the high-scoring design ah. Recall that yn represents the predefined target score and is generally much larger than yti- By minimizing La21 (x),, (xx), we aim to find high-scoring design I, such that any design in the static dataset similar to the high scoring design is encouraged to have high score prediction. In this way, 2y tries to incorporate as many high-scoring features from the static dataset as possible.
Optimization. Combining the two losses, the overall loss can be expressed as:
((x)) 2 (1lw, (y) K Xx,x, XKx,x, + +,31) "1)l
(14)
+ | (y: Kx xx ( Xx x, + (T) yı 113)
where only the high-scoring designs Xn are learnable parameters in the whole procedure of BDI, We optimize X against ((X;,) by gradient descent methods such as Adam 21
Experiments
We conduct extensive experiments on design-bench [. and aim to answer three research questions: (1) How does BDI compare with both recently proposed offline model-based algorithms and traditional algorithms? (2) Is every component necessary in BDI? (i.c., (12)(AN)) C121 (X,), NTK.) (3) Is BDI robust the hyperparameter choices including the predefined target score yh. the weight parameter and the number of steps T?
4.1 Dataset and Evaluation
To evaluate the effectiveness of BDI, we adopt the commonly used design-bench bench, including both continuous and discrete tasks, and the evaluation protocol as in the prior work [8].
Task overview. We conduct experiments on four continuous tasks: (a) Superconductor (SuperC) [2]. where the aim is to design superconductor with 86 continuous components to maximize critical temperature with 17010 designs; (b) Ant Morphology (Ant) I 22 where the goal is to design the morphology
quadrupedal ant with 60 continuous components to crawl quickly with 10004 designs, (c) D' Kitty Morphology (D' DD'Kitty) [ where the goal is to design the morphology of quadrupedal Kitty with continuous components to crawl quickly with 10004 designs; and (d) Hopper Controller (Hopper) D. where the aim is to find neural network policy parameterized by 5126 total weights to maximize return with 3200 designs. Besides the continuous tasks, we perform experiments three discrete tasks: (e) GFP
where the objective is find a length 238 protein sequence to maximize fluorescence with 5000 designs; (f) TF Bind (TFB) 5), where the aim is to find length 8 DNA sequence to maximize binding activity score with 32896 designs; and (g) UTR 6 where the goal is to find a length 50 DNA sequence to maximize expression level with 140, 000 designs. These tasks contain neither personally identifiable nor offensive information.
Evaluation. We follow the same evaluation protocol in { we choose the top 128 most promising designs for each method, and then report the 100th percentile normalized ground truth score as Vv ymae Jmiin
Wmin where Umin and Ymez represent the lowest score and the highest score in the full unobserved dataset, respectively. The additional t0th percentile (median) normalized ground truth scores, used in the prior work [8], are also provided in Appendix A. To better measure the performance across multiple tasks, we further report the mean and median ranks of the compared algorithms over all 7 tasks.
5

---

4.2 Comparison Methods
We compare BDI with two groups of baselines: (i) sampling via a generative model; and (ii) gradient updating from existing designs. The generative model based methods learn a distribution of the high-scoring designs and then sample from the learned distribution. This group of methods includes:
1) CbAS [24]: trains VAE model of the design distribution using designs with score above threshold and gradually adapts the distribution to the high-scoring part by increasing the threshold. 2) Auto.CbAS [25]: uses importance sampling to retrain a regression model using the design distribution introduced by CbAS; 3) MIN [20]: learns an inverse map from the score to the design and queries the inverse map by searching for the optimal score to obtain the optimal design. The latter group includes: 1) Grad: applies simple gradient ascent on existing designs to obtain new designs; 2) ROMA [7]: regularizes the smoothness of the DNN and then performs gradient ascent to obtain new designs; 3) COMs []: regularizes the DNN to assign lower scores to designs obtained during the gradient ascent process and leverages this model to obtain new designs by gradient ascent; 4) NEMO [9]: bounds the distance between the proxy and the objective function by normalized maximum likelihood and then performs gradient ascent. We discuss the NTK based Grad in Section 4.5
Besides the recent work, we also compare BDI with several traditional methods described in [: 1) BO-qEI [26]: performs Bayesian Optimization to maximize the proxy, proposes designs via the weepeeeeed-mpppoeeeeeee
acquisition function, and labels the designs via the proxy function; 2) CMA-ES |27]: an evolutionary algorithm gradually adapts the distribution via the covariance matrix towards the optimal design; 3) REINFORCE [28]: first learns a proxy function and then optimizes the distribution over the input space by leveraging the proxy and the policy-gradien estimator.
4.3 Training Details
We follow the training settings of (8] for all comparison methods if not specified. Recall that M represents the number of the high-scoring designs. We set M = 1 in our experiments, which achieves expressive results, and discuss the M >1 scenario in Appendix A.2. We set the regularization / as le following [16]. We set the predefined target score yı. as a constant 10 across all tasks, which proves to be effective and robust all cases. We set a as 1e for all continuous tasks and as 0.0 for all discrete tasks, and set the number of iterations T to 200 in all experiments. We discuss the choices of yh, a and
Sec. 4.6 in detail. We adopt a 6-layer MIP (Multil Layer Perceptron) followed by ReI.U for all gradient updating methods and the hidden size is set as 2048. We optimize L ((X,) with the Adam optimizer (1 with 1 learning rate for discrete tasks and a 1e-3 learning rate for continuous tasks. We cite the results from [8] for the non-gradient-asee ascent methods including BO-qEI, CMA-ES REINFORCE, CbAS, CbA. For other methods, we run every setting over trials and report the mean and the standard error. We use the NTK library 22) build on Jax 30) to conduct kereel-based experiments and use Pytorch (31] for other experiments. All Jax experiments are run on multiple CPUs within cluster and all Pytorch experiments run on one V100 GPU. We discuss the computational time details for all tasks in Appendix A.3.
4.4 Results and Analysis
We report experimental results in Table for continuous tasks and in Table for discrete tasks where D(best) represents the maximal score of the static dataset for each task. We bold the results within one standard deviation of the highest performance.
Results on continuous tasks. As shown in Table BDI achieves the best results over all tasks. Compared with the naive Grad method, BDI achieves consistent gain over all four tasks, which suggests that BDI can align the high-scoring design with the static dataset and thus mitigate the nut-of-distribution problem. As for COMs. ROMA and NEMO, which all impose a prior on the DNN model, they generally perform better than Grad but worse than BDI. This indicates that directly transforming the information of the static dataset into the high-scoring design is more effective than adding priors on the DNN in terms of mitigating the diitrr-distriuution problem. Another advantage of BDI is in its neural tangent kernel nature: 1) Real-world offline model- optimization often involves scall-scale datasets since the labeling cost of protein/DNA/robot is very high; 2) BDI is built on the neural tangent kernel, which generally performs better than finite neural networks on
2 *CMA-ES -ES performs very differently on Ant Morphology and D' D'Kitty Morphology and the reason is that Ant Morphology is simpler and much more sensitive initializations.

---

Table 1: Experimental results on continuous tasks for comparison.
Method
Superconductor
Ant Morphology
D D'Kitty Morphology
Hopper Controller
D(best)
0.399
0.565
018884
1.0
BO-qEI
0.402 0.034
=.09:000000
0.896 + 0.000
0.350 + 0.118
CMA-ES
0.465 = 0.024
1.214 0.732
0.724 + 0.001
0.604 + 0.215
REINFORCE
0.481 0.013
0.266 + 0.032
0.562 + 0.196
-0.020
CbAS
0.503 0.069
0.876 + 0.031
0.892 = 0.008
0.141 = 0.012
Auto.CbAS
0.421 0.045
0.882 1 0.045
0.906 1 0.006
0.137 1 0.005
MIN
0.452 1 0.026
0.908 1 0.020
0.942 1 0,010
0.094 0.113
Grad
0.483 0.027
0.764 0.047
0.888 + 0.035
0.989 = 0.362
COMS
0.487 0.023
0.866 0.050
0.835 = 0.039
1.224 = 0.555
ROMA
0.476 = 0.024
0.905 + 0.018
1.849 = 0.110
NEMO
0.488 0.034
0.814 0.043
= 0.012
1.959 = 0.159
(ours)
0.520 0.005
0.962 E 0.000
0.941 = 0.000
1.989 = 0.050
Table 2: Experimental results on discrete tasks & ranking on all tasks for comparison.
Method
GFP
TF Bind
UTR
Rank Mean
Rank Median
D(best)
0.789
0.439
0.593
BO-qEI
0.254 1 0.352
0.798 1 0.083
0.684 1 0.000
8.6/11
9/11
CMA-ES
0.054 1 0.002
0.953 1 0.022
0.707 1 0.014
5.7/11
6/11
REINFORCE
0.865 1 0.000
0.948 1 0.028
0.688 1 0.010
7.4/11
9/11
CbAS
0.865 1 0.000
0.927 1 0.051
0.694 0.010
1.611
5/11
Auto.CbAS
0.865 1 0.000
0.910 1 0.044
0.691 1 0.012
6.0/11
7/11
MIN
0.865 1 0.001
0.882 1 0.020
0.694 0.017
5.3/11
4/11
Grad
0.864 = 0.001
0.491 + 0.042
#.0013
7.9/11
8/11
COMs
0.861 +: 0.009
0.920 1 0.043
0.699
1 0.011
5.6/11
6/11
ROMA
0.558 + 0.395
0.928 1 0.038
0.690 0.012
6.1/11
7/11
NEMO
0.150 + 0.270
0.905 + 0.048
0.694
+ 0.015
5.4/11
4/11
BDI (ours)
0.864 1: 0.000
0.973 1 0.000
0.760 1 0.000
1.9/11
1/11
3) COMs, ROMA, and NEMO built finite neural network and it is non-trivial to modify them to the neural tangent kernel version. The generative model based methods CbAS, Cbto. and MIN perform poorly for high-dimensional tasks like Hopper (D=5126) since the high-dimensional data distributions are harder to model. Compared with generative model based methods, BDI only achieves better results but is also much simpler.
small-scale datasets 32]
Results on discrete tasks. As shown in Table 2 BDI achieves the best or equal-best performances in two of the three tasks, and is only marginally inferior in the third. This suggests that BDI is also powerful approach in the discrete domain. For the GFP task, every design is length 238 sequence of 20-categorical one-hot vectors and gradient updates that do not take into account the sequential nature may be less beneficial. This may explain why BDI does not perform as well on GFP.
Overall, BDI achieves the best result in terms of the ranking as shown in Table 2 and Figure and attains the best performance on 6/7 tasks.
Table 3: Ablation studies on BDI components.
Task
D
BDI
w/o Liza
w/o Li21
NTK2NN
NTK2RBF
GFP
238
0.864 + 0.000
0.860 + 0.004
0.864
0.000
0.252 + 0.354
0.862
0.003
TFB
8
0.973 1 0.000
0.984 1 0.000
0.973
0.000
0.958 1 0.025
0.925
0.000
UTR
50
0.760 1 0.000
0.636 1 0.000
0.777 1 0.000
0.699 1 0.009
0.738 0.000
SuperC
86
0.520
0.005
0.511 = 0.007
0.502
0.007
0.450 0.036
0.510
0.012
Ant
60
0.962
0.000
0.914 1 0.000
0.933
0.000
0.861 1 0.019
0.927
0.000
D' D'Kitty
56
0.941
+ 0.000
0.920 1 0.000
0.942
0.001
0.929 1 0.010
0.938
0.008
Hopper
5126
1.989 + 0.050
1.935 = 0.553
1.821 + 0.223
0,503 + 0.039
0.965 0.103
7

---

1.0
GFP 178
REMFORCE
CMA-ES
CAAS
Score rabi
1.2
SuPC
Hopper
Diitty
0.9
C14S
HNN
1.1
0.8
d= 16-1
g 0.0
d= 5p
0- 40
Gred
1.0
8= 2e
0.7
KOMA
NENO
0.6
6 7 Rank
8
9 1011
0.8
10 15
20
25
30
50
100
150
200
Predefined target score
Step
Figure 2: Rank minima and max- Figure 3: The ratio of the ima are indicated by whiskers; ver- ground truth score of the de- tical lines and black triangles rep- sign with y/. to the ground resent medians and means.
Figure 4: The ground truth score the design as a func- tion of the step T for different weight parameters a on Ant.
truth score with yı = 10.
4.5 Ablation Studies
In this subsection, we conduct ablation studies on the components of BDI and aim to verify the effectiveness of Ca2 (X.), Li2s (Xx),, and NTK. The baseline method is proposed BDI and remove each component to verify its effectiveness. We replace the NTK with its corresponding finite DNN (6-layer MLP followed by ReLU), which is trained on (X1, 31) for Eq.(5) and (Xn, U/h for Eq. Following [16,, another alternative is to replace the NTK with RBF.
As shown in Table 3
BDI generally performs the best (at least second-best across tasks, which indicates the importance of all three components.
Forward mapping. Removing L/2x (12)(Xn) yields the best result 0.984 in the TFB task, but worsens the designs in other cases, which verifies the importance of the forward mapping. Incorporating the fi2A (X,) loss term is similar to using the Grad method as we discuss in Section 3.1 We also run experiments on gradient ascent with NTK. This approach first builds regression model using the NTK and then performs gradient ascent on existing designs to obtain new designs. Results are very similar BDI without Ca21( n,), ), so we do not report them here.
Backward mapping, Removing Lizi( (x,) decreases model performances in three tasks, and only improves it for the UTR task. This demonstrates the importance of the backward mapping. Although removing C121 leads to the best result 0.777 in UTR compared with 0.760 of BDI, this does not necessarily demonstrate that the removal of L1.21 is desirable. The 50th percentile of BDI is higher than that of BDI without L122 (0.664 > 0.649). Furthermore, we observe that the backward mapping is more effective for continuous tasks, possibly because do model the sequential nature of DNA and protein sequences. We follow the default procedure in [8] which maps the sequence design real-valued logits of categorical distribution for fair comparison. Given more effective discrete space modeling, the backward mapping might yield better results; we leave this to future work.
Neural tangent kernel. We replace the NTK with the corresponding finite DNN (NTK2NN) and the RBF kernel (NTK2RBF), and both replacements degrade the performance significantly. which verifies the effectiveness of the NTK. In addition,
observe that NTK2NN performs relatively well on the low-dimensional task TFB (D=8) but poorly on high-dimensional tasks, especially for GFP (D 238) and Hopper (D) 5126). It is likely that the high-dimensional designs X,, are more sensitive to the approximation in Eq. (9) Furthermore, BDI performs better than NTK2RBF and this can be explained follows. Compared with the RBF kernel, the NTK enjoys the high expressiveness of a DNN [15] and represents a broad class of DNNs [16] [7], which can better extract the design features and enhance the generalization of the high-scoring designs.
4.6 Hyperparameter Sensitivity
In this subsection, we first study the sensitivity of BDI to the predefined target score yh- Note that all algorithms in this subsection do not have access to the ground truth scores, which are only used for analysis. We expect this score to be larger than the maximal score of the static dataset since we want to obtain designs better than those in the static dataset. We first normalize all scores to have unit Gaussian statistics following g small target score may not be able to identify high-scoring
design and large score is too hard to reach. Thus we choose to set the target score yh 10 across all
8

---

Table 4: Comparison between data distillation and backward mapping.
Comparison
Data distillation
Backward mapping
Data
images in the training datasct D
designs in the offline dataset D some measurement of protein/dna/robot larger than the max of the offline dataset D mean squared error for regression distill æ to incorporate high-scoring features
Label content
0-9 for the MNIST dataset
Label value
within the training dataset
Loss function
cross-entropy for classification
Task objective
distill D into small subset
tasks. In Figure 3 we report the ratio of the ground truth score of the design with yh to the ground truth score of the design with yı, 10 and the score details are in Appendix A.4. All tasks are robust to the change of yh- Although we set yn, as 10, which is smaller than the maximal score 10.2 for Hopper, BDI still yields a good result, which demonstrates the robustness of BDI. For Hopper, after we set yh to larger value, we can observe that the corresponding results become better.
Besides yh. we also explore robustness to the choice of the weight parameter a and the number of steps T. We study the continuous task Ant and the discrete task TFB. As shown in Figure 4 we visualize the 100th percentile ground truth score of the Ant task as function of T for 0.0, 5e 4, le 2e ,e 3 We can observe that the behavior of BDI is robust to change of a within the range of evaluated values. BDI yields the high-scoring designs at early steps and then the solutions remain stable. This demonstrates the robustness of BDI to the choice of T. For the TFB task, we find the behavior of BDI is identical for different and the behavior with respect to is very similar to that of Ant. Thus we do not report details here. One possible reason for the insensitivity to is that TFB is a discrete task and the discrete solution is more robust to a.
5 Related Work
Offline model-based optimization. Recent methods for offline model-based optimization can be divided into two categories: (i) obtaining designs from generative model and (ii) applying gradient ascent on existing designs. In the first category, the methods in [24 25 gradually adapt VAE-based generative model towards the optimized design by leveraging a proxy function. Kumar et al. adopt an alternative approach in (20), learning a generator that maps the score to the design. The generative model based methods often require careful tuning to model the high-scoring designs [].
Gradient-based methods have received substantial attention recently because they can leverage DNN models to produce improved designs. A challenge in the naive application of such methods is out-of -ift--dddittitttionn designs, for which the trained DNN model produces inaccurate score predictions. Several approaches have been proposed to address this. Trabucco al. add a regularizing term for conservative objective modeling [8] 33 In [9], Fu et al. leverage normalized maximum likelihood to bound the distance between the proxy function and the ground truth. Yu et overcome the brittleness of the proxy function by utilizing a local smoothness prior [7]. Our work falls into the gradient-baat based line of research, but rather than imposing a prior on the model, we propose the bidirectional mapping to ensure that the proposed designs can be used to predict the scores of the data and vice versa, thus ensuring that the high-scoring designs are more aligned with the static dataset. Similar to our backward mapping constraining designs, BCQ [d constrains the state-action pairs to be contained in the support of the static datasct. Our proposed backward mapping is different in its implementation and also enables the model to generate designs outside the training distribution.
aims to distill large training set into a small one, which enables efficient retraining. Building upon [18. the methods in 39 40) propose to distill labels instead of images. Approaches in [6 41 introduce a meta-learning algorithm powered by NTK, which significantly improves the performance. The backward mapping of BDI is inspired by 16 41 and the difference is that BDI is handling a regression problem with a predefined target score larger than the maximal score of the static dataset. In Table we provided a detailed comparison of backward mapping and data distillation; both adopt a bi-level framework
Data distillation. Data distillation [T8
35 38
42
This bi-level
framework can also be used to learn other hyperaparameters [1 43.

---

Conclusion and discussion
In this paper, we address offline model-based optimization and propose a novel approach that involves bidirectional score predictions. BDI requires the proposed high-scoring designs and the static offline dataset to predict each other, which effectively ameliorates the out-of-distribution problem. The finite- width DNN can only yield approximate loss function, which leads to significant deterioration of the design quality, especially for high-dimense -ineessinnii tasks. We thus adopt an infinite-widhh DNN and employ the NTK yield closed foosed-form solution, leading better designs. Experimental results verify the effectiveness of BDI, demonstrating that outperforms state-of-the-art approaches.
Limitations. Although BDI has been demonstrated to be effective in tasks from wide range of domains, some evaluations are not performed with realistic setup. For example, some of our experiments, such the GFP [3] task, we follow the prior work and use transformer model which is pre-trained on a large dataset as the evaluation oracle. However, this may not reflect real- world behavior and the evaluation protocol can be further improved by close collaboration with domain experts [0). Overall, we still believe BDI can generalize well to these cases since BDI has simple formulation and proves to be effective and robust for the wide range of tasks in the design-bench.
Negative impacts. Designing a new object or entity with desired properties is a double-edged sword. On one hand, can be beneficial to society, with an example being drug discovery to cure unsolved diseases. On the other hand, if these techniques are acquired by dangerous people, they can also be used to design biochemicals to harm our society. Researchers should pay vigilant attention to ensure their research does end up being used positively for the social good.
7
Acknowledgement
We thank Aristide Baratin from Mila and Greg Yang from Microsoft Research for their helpful discussions on the neural tangent kernel and thank Zixuan Liu from the University of Washington for his helpful suggestions the paper writing. This research supported in part by funding from the Fonds de recherche du Québec Nature et technologies.
References
[1] Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench: Bench- marks for data-driven offline model-based optimization. arXiv preprint arXiv:2202.08450. 2022.
[2] Kam Hamidieh. data-driven statistical model for predicting the critical temperature of a superconductor. Computational Materials Science, 2018.
[3] Karen S Sarkisyan et al. Local fitness landscape of the green fluorescent protein. Nature, 2016.
[4] Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model- reinforcement learning for biological sequence design. In Proc. Int. Conf. Learning Rep. (ICLR), 2019.
[5] Luis A Barrera et al. Survey of variation in human transcription factors reveals prevalent dna binding changes. Science, 2016.
[6] Paul Sample, Ban Wang, David W Reid, Vlad Presnyak, lain J McFadyen, David R Morris, and Georg Seelig. Human 5 UTR design and variant effect prediction from a massively parallel translation assay. Nature Biotechnology, 2019.
[7] Sihyun Yu, Sungsoo Ahn, Le Song, and Jinwoo Shin. Roma: Robust model adaptation for offline basel optimization. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2021.
[8] Brandon Trabucco, Aviral Kumar, Xinyang Geng, and Sergey Levine. Conservative objective models for effective offline model-based optimization. In Proc. Int. Conf. Learning Rep. (ICLR), 2021.
10

---

[9] Justin Fu and Sergey Levine. Offline model based optimization via normalized maximum likelihood estimation. Proc. Int. Conf. Learning Rep. (ICLR), 2021.
[10] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2021.
[11] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2018.
[12] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2019.
[13] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks Gaussian processes. Proc. Int. Conf. Learning Rep. (ICLR), 2017.
[14] Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics. Proc. Int. Conf. Machine Lea. (ICMI.), 2021.
[15] Jachoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2020.
[16] Timothy Nguyen. Zhourong Chen, and Jachoon Lee. Dataset meta-learning from kernel ridge-regression. Proc. Int. Conf. Learning Rep. (ICLR), 2020.
[17] Chia-Hung Yuan and Shan-Hung Wu, Neural tangent generalization attacks. In Proc. Int. Conf. Machine Lea. (ICML), 2021.
[18] Tongzhou Wang, Jun- Yan Zhu, Antonio Torralba, and Alexei A Efros. Dataset distillation. arXiv preprint Xrviv: 1811.19959 2018
[19] Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-learning. Proc. Int. Conf. Machine Lea. (ICML), 2018
[20] Aviral Kumar and Sergey Levine. Model inversion networks for model-based optimization. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2020.
[21] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proc. Int. Conf. Learning Rep. (ICLR), 2015.
[22] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv: 160600/60000 2016.
[23] Michael Ahn, Henry Zhu, Kristian Hartikainen, Hugo Ponte, Abhishek Gupta, Sergey Levine, and Vikash Kumar. Robel: Robotics benchmarks for learning with low-cost robots. In Conf. on Robot Lea. (CoRL), 2020.
[24] David Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. In Proc. Int. Conf. Machine Lea. (ICML), 2019.
[25] Clara Fannjiang and Jennifer Listgarten. Autofocused oracles for model-based design. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2020.
[26] James T Wilson, Riccardo Moriconi, Frank Hutter, and Marc Peter Deisenroth. The reparame- terization trick for acquisition functions. arXiv preprint arXiv:1712.00422, 2017.
[27] Nikolaus Hansen. The CMA evolution strategy: A comparing review. In Jose A. Lozano, Pedro Larrañaga, Iñaki Inza, and Endika Bengoetxea, editors, Towards a New Evolutionary Computation: Advances in the Estimation of Distribution Algorithms,, pages 75-102. Springer Berlin Heidelberg, Berlin, Heidelberg, 2006.

---

[28] Ronald J Williams. Simple statistical gradient-followwee algorithms for connectionist reinforce- ment learning. Machine Learning. 1992.
[29] Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl- Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In Proc. Int. Conf. Learning Rep. (ICLR), 2020.
[30] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milne. Jax: composable transformations of python+ numpy programs. URL http::/gglhub..oo/oggggee/e/ 2020.
[31] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-peric perffommrrm deep learning library. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2019.
[32] Sanjeev Arora, Simon S. Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the power infinitely wide deep nets on small-data tasks. Proc. Int. Conf. Machine Lea. (ICML), 2020.
[33] Dinghuai Zhang, Tianyuan Zhang, Yiping Lu, Zhanxing Zhu, and Bin Dong. You only propagate once: Accelerating adversarial training via maximal principle. Proc. Adv. Neur Inf. Proc. Syst (NeurIPS), 2019.
[34] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Proc. Int. Conf. Machine Lea. (ICML), 2019.
[35] Yongchao Zhou, Ehsan Nezhadarya, and Jimmy Ba. Dataset distillation using neural feature regression. arXiv preprint (xXXi:2206600719, 2022.
[36] Justin Cui, Ruochen Wang, Si Si, and Cho-Jui Hsieh. DC-BENCH: Dataset condensation benchmark. arXiv preprint ar 2rxiv:2277.09639 2022.
[37] Noveen Sachdeva, Mehak Preet Dhaliwal, Carole-Jean Wu, and Julian McAuley. Infinite recommendation networks: A data-centric approach. arXiv preprint arXiv:2206. NiNi::2206.02666, 2022.
[38] Songhua Liu, Kai Wang, Xingyi Yang. Jingwen Ye, and Xinchao Wang. Dataset distillation via factorization. Proc. Adv. Neur. Inf. Proc. Syst (NeurlPS), 2022.
[39] Ondrej Bohdal, Yongxin Yang, and Timothy Hospedales. Flexible dataset distillation: Learn labels instead of images. arXiv preprint arXiv:2006. 08572, 2020.
[40] Ilia Sucholutsky and Matthias Schonlau. Soft-label dataset distillation and text dataset distilla- tion. In Proc. Int. Joint Conf. on Neur. Net. (IJCNN), 2021.
[41] Timothy Nguyen, Roman Novak, Lechao Xiao, and Jachoon Lee. Dataset distillation with infinitely wide convolutional networks. Proc. Ady. Neur. Inf. Proc. Syst (NeurIPS), 2021.
[42] Can Chen, Xi Chen, Chen Ma, Zixuan Liu, and Xue Liu. Gradient- based bi-level optimization for deep learning: survey, arXiv preprint arXiv:2207. 17719, 2022.
[43] Zhiting Hu, Bowen Tan, Russ R Salakhutdinov, Tom M Mitchell. and Eric P Xing. Learning data manipulation for augmentation and weighting. Proc, Adv. Neur. Inf. Proc. Syst (NeurIPS), 2019.
[44] Mengye Ren, Wenyuan Zeng, Bin Yang, and Raquel Urtasun. Learning to reweight examples for robust deep learning. In Proc. Int Conf. Machine Lea. (ICML), 2018.
[45] Can Chen, Chen Ma, Xi Chen, Sirui Song, Hao Liu, and Xue Liu. Unbiased implicit feedback via bi-level optimization. arXiv preprint arXiv:2206.00147 2022.
[46] Can Chen, Shuhao Zheng, Xi Chen, Erqun Dong, Xue Steve Liu, Hao Liu, and Dejing Dou. Generalized data weighting via class-level gradient manipulation. Proc, Ady. Neur. Inf. Proc. Syst (NeurIPS), 2021.
12

---

[47] Can Chen, Jingbo Zhou, Fan Wang, Xue Liu, and Dejing Dou. Structure-aware protein self- supervised learning. arXiv preprint vrXiv:2204.042111 2022.
[48] Kexin Huang, Tianfan Fu, Wenhao Gao, Yue Zhao, Yusuf Roohani, Jure Leskovec, Connor W Coley. Cao Xiao, Jimeng Sun, and Marinka Zitnik. Artificial intelligence foundation for therapeutic science. Nature Chemical Biology, 2022.
[49] Changlin Wan. Machine Learning Approaches to Reveal Discrete Signals Gene Expression. PhD thesis, Purdue University Graduate School, 2022.
[50] Xiao Luo, Xinming Tu, Yang Ding, Ge Gao, and Minghua Deng. Expectation pooling: an effective and interpretable pooling method for predicting dna-protein binding. Bioinformatics, 2020.
[51] Yanrong Ji, Zhihan Zhou, Han Liu, and Ramana V Davuluri. Dnabert: pre-trained bidirectional encoder representations from transformers model for dna-language in genome. Bioinformatics 2021.
[52] Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rihawi, Yu Wang, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, et al. ProtTrans: towards cracking the language of life's code through self-supervised deep learning and high performance computing. arXiv preprint arXiv:2007.0622 2020.
[53] Courtney E Gonzalez and Marc Ostermeier. Pervasive pairwise intragenic epistasis among sequential mutations in tem-1 3-lactamase. Journal of molecular biology, 2019.
[54] Elad Firbeerg, Jason W Labonte, Jeffrey Gray, and Mare Ostermeier. A comprehensive, high-resolution map of gene's fitness landscape. Molecular biology and evolution, 2014.
[55] Dinghuai Zhang, Jie Fu, Yoshua Bengio, and Aaron Courville. Unifying likelihood-frie inference with black-box optimization and beyond. In Proc. Int. Conf. Learning Rep. (ICLR), 2021.
[56] Emily E Wrenbeck, Laura R Azouz, and Timothy A Whitehead. Single- mutation fitness landscapes for an enzyme on multiple substrates reveal specificity is globally encoded. Nature Communications, 2017.
[57] Justin R Klesmith, John- Bhul Bacik, Ryszard Michalezyk, and Timothy A Whitehead. Compre- hensive sequence-flux mapping of levoglucosan utilization pathway e. coli. ACS Synthetic Biology, 2015.
[58] Jochen Weile, Song Sun, Atina G Cote, Jennifer Knapp, Marta Verby, Joseph C Mellor, Yingzhou Wu, Carles Pons, Cassandra Wong, Natascha van Lieshout, et al. A framework for exhaustively mapping functional missense variants. Molecular Systems Biology, 2017.
[59] Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining, arXiv preprint arXiv:2203.06115,, 2022.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] We discuss limitations in Sec-
tion G
(c) Did you discuss any potential negative societal impacts of your work? [Yes] We discuss negative impacts in Section 6
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
13

---

(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proof' of all theoretical results? [N/A]
3. 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 URL.)? [Yes] They in the supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] All of these can in found in Section
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] We run every setting 8 times and report the mean and standard deviation in the tables.
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes] is in Section 4.3
4. If you using cxisting assets (c.g., code, data, models) curating/releasing
assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cite them [
(b) Did you mention the license of the assets? [Yes] We mention the license in the README.mu of our code the supplemental material.
(c) Did you include any new assets either in the supplemental material or as URL? [Yes] We include our code in the supplemental material new asset and we will release this new asset upon paper acceptance.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [Yes] The assets are publicly available and free for research use. (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [Yes] We use design-bench, which does not contain personally identifiable information nor offensive content, as illustrated in Section 4.1
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "bWVQzEOkk_e", "reviewer": "Reviewer_HR4n", "summary": "This paper proposes bidirectional learning for offline infinite-width model-based optimization (BDI) to address the out-of-distribution problem. BDI consists of a forward mapping and a backward mapping. The authors claim that the backward mapping can distill more information into high-scoring designs. Experiments show that BDI achieves state-of-the-art performance in both continuous and discrete tasks.", "questions": "Please refer to the \"Strengths And Weaknesses\" part.", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality:\n\nThe idea of BDI is novel. This paper introduces the idea of neural tangent kernels into offline model-based optimization. Moreover, the authors propose a novel backward mapping, significantly improving the performance of model-based optimization.\n\nQuality:\n\nThe comparisons are thorough. Results show that BDI outperforms baselines in most tasks, and all component in BDI is useful. The use of neural tangent kernels in BDI is based on solid theoretical motivation. Nevertheless, I am confused about the backward mapping.\n1. What is the motivation for using the backward mapping? The authors may want to provide an intuitive explanation of why we need a backward mapping besides the forward mapping.\n2. How does the backward mapping help to distill information into the high-scoring designs? Though the authors provide ablation studies to show that the backward mapping does work, the authors may want to provide more analyses/visualization to show why the backward mapping works.\n3. I can not understand the predefined score $y_h$ used in the backward mapping. In forward mapping, a high $y_h$ can encourage the algorithm to generate $X_h$ with high scores. In backward mapping, what is the purpose of training $f_\\theta^h$ with a virtual score $y_h$?\n\n\n\nClarity:\n\nMost of this paper is easy to follow. Nevertheless, I am confused about the backward mapping. The authors may want to discuss the backward mapping more, especially its motivation and how it works.\n\nSignificance:\n\nThe backward mapping is novel and can provide a significant performance improvement. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "_M9owUpMXf", "reviewer": "Reviewer_qG1d", "summary": "Offline model-based optimization could suffer from the distributional shift issue where the model bias could be exploited by the optimization process to output poor designs. Most previous offline model-based optimization aims to solve the problem by introducing loss function terms that promote distribution matching between the learned model outputs and the ground-truth. Unlike previous works, this paper introduces a directional learning approach based on neural tangent kernel to facilitate the distribution matching. By encouraging learning accurate mappings from low-scoring designs to high-scoring designs and vice versa, the proposed method could alleviate the issue of outputting unseen designs with significantly overestimated scores. The proposed method achieved the best performance on six out of seven widely studied offline model-based optimization task.", "questions": "Could the authors explain why using a constant $y_h$ could work? Would it be possible to further improve the proposed method with a more sophisticated design of $y_h$ where the values are different for each high-scoring sample?", "limitations": "The authors have addressed the limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n1. The proposed approach is well motivated and the writing is easy to follow.\n2. The authors provide a good insight on why the bidirectional learning helps mitigate the distributional shift issue in 3.3.\n3. Experiment results show strong performance of the proposed method.\n\n### Weakness\nThis paper lacks a discussion on the relationship between the directional learning and a conservative regularization. I believe that adding the backward mapping loss is similar to a conservative regularizer in the sense that they both aim to alleviate the distributional shift between the predicted optimal designs and the dataset designs. I hope that the authors could provide some discussions on their similarity/difference and why learning the backward mapping could be a better alternative than conservative regularization methods.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "2KuycUBCl-m", "reviewer": "Reviewer_QNRb", "summary": "This paper investigates the subject of Offline Model-Based Optimization, which involves processing a static dataset of designs X and function evaluations Y in order to solve the optimization problem X* = \\argmax_{X} f(X) for an unknown function f. This paper proposes a gradient-based method for solving Offline MBO problems using infinite-width DNN models. The approach facilitates efficiently optimizing designs X via a proposed bidirectional objective that encourages designs to be found that not only achieve high performance under an approximation of the objective function f_{\\theta} (X), which is referred to as the forward mapping in the paper, but jointly maximize a distillation objective that encourages optimized design-score pairs to be informative about the characteristics of the true objective function (the backward mapping).\n\nOne key novelty in the paper is the proposed use of a distillation objective to constrain the optimized design-score pairs, rather than regularizing the DNN model used for optimization. Results in the paper appear to show the effectiveness of this principle on a standard benchmarking framework for Offline MBO using a comprehensive set of baselines.", "questions": "(1) There appears to be a discrepancy between a small fraction of the results presented in the main performance table, and those reported in other recent papers. In particular, the performance of MIN and COMs are different from those reported in design-bench (Trabucco, B. et al., 2022) on the following tasks: Ant Morphology, DKitty Morphology, and Hopper Controller. Could the authors investigate this and explain the cause of this difference?\n\n(2) The optimization objective in Equation (4) is chosen as a regression objective, rather than purely maximizing the predictions of the learned objective function f_{\\theta} (X). How are the results affected when this equation is changed to pure maximization, rather than regression?\n\n(3) The objective functions in Equation (8) appear to be assigned equal weight. How are the results affected by ablating the weight of the distillation term (Equation 6) in the loss function?", "limitations": "The authors have sufficiently addressed these.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality:\n\nThe subject of mitigating out-of-distribution samples (ie, design-score pairs) using gradient-based Offline Model-Based Optimization techniques has been the subject of multiple papers. The subject itself is not novel, but the solution proposed by the authors (regularizing the optimized samples rather than the model) is an original proposition in the Offline MBO domain. Though different in execution, a similarly inspired technique has been proposed in the Offline RL setting, namely BCQ (Fujimoto, S. et al., 2019), which constrains the (s’, a’) pairs sampled during bellman backups \\max_{a’} Q(s’, a’) to be contained in the support of the static dataset.\n\nBoth papers explore constraints (either implicitly or explicitly) on the data being optimized (actions in the case of BCQ, and designs in the case of BDI), rather than model regularization. However, it should be noted that the implementation of the constraint is fundamentally different between these two approaches, and BDI’s implementation is original. Dataset distillation enforces a constraint on the data that is an appealing alternative to the support-based constraint seen in BCQ because it lets the model generate designs outside the training distribution.\n\n\nQuality:\n\n\nThe experimental setup in the paper adheres to the evaluation protocol established in prior benchmarking papers, which aids in the reproducibility of the paper. Relevant baselines from the literature, including several new methods, and multiple classical methods, adapted to the Offline MBO setting, are reported. These factors improve the quality of the paper’s main evaluation. Certain relevant ablations are not included (see my questions below), and a small number of claims and speculations made in the paper are not fully supported by evidence. For example, on line 278-279 of the manuscript, the authors write “we observe that the backward mapping is more effective for continuous tasks, possibly because we do not model the sequential nature of DNA and protein sequences.” This is a helpful intuition, and could be improved if the authors tested this hypothesis with a model that does “model the sequential nature of DNA and protein sequences.” Quality may also be improved by addressing the questions listed below.\n\nClarity:\n\n\nCertain experimental details are not clear from the manuscript. For example, on line 194-195 the authors write “we choose the top N = 128 most promising designs for each method,” which could suggest more than N = 128 designs are generated, and the number is downsampled to 128 using some indicator for how promising the designs are. If such a post-selection mechanism is used, additional details should be included in the paper for reproducibility. Such a mechanism should query the ground truth objective function only N = 128 times to compute 100th percentile and 50th percentile statistics as per design-bench (Trabucco, B. et al., 2022).\n\nSignificance:\n\n\nMitigating the out-of-distribution problem in Offline MBO has significance as a research area, due to its compelling applications in real-world design problems. The results in this paper show the proposed approach is quite effective in practice at finding high-performing designs.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept