Bidirectional Learning for Offline Model-based Biological Sequence Design
Authors: Can Chen
Yingxue Zhang
Xue Liu
Mark Coates
Venue: Submitted to ICLR 2023 (2023)
Keywords: offline model-based optimization
meta learning
biological sequence design.

Abstract: Offline model-based optimization aims to maximize a black-box objective function with a static dataset of designs and their scores. In this paper, we focus on biological sequence design to maximize some sequence score. A recent approach employs bidirectional learning, combining a forward mapping for exploitation and a backward mapping for constraint, and it relies on the neural tangent kernel (NTK) of an infinitely wide network to build a proxy model. Though effective, the NTK cannot learn features because of its parametrization, and its use prevents the incorporation of powerful pre-trained Language Models (LMs) that can capture the rich biophysical information in millions of biological sequences. We adopt an alternative proxy model, adding a linear head to a pre-trained LM, and propose a linearization scheme. This yields a closed-form loss and also takes into account the biophysical information in the pre-trained LM. In addition, the forward mapping and the backward mapping play different roles and thus deserve different weights during sequence optimization. To achieve this, we train an auxiliary model and leverage its weak supervision signal via a bi-level optimization framework to effectively learn how to balance the two mappings. Further, by extending the framework, we develop the first learning rate adaptation module Adaptive-$\eta$,  which is compatible with all gradient-based algorithms for offline model-based optimization. Experimental results on DNA/protein sequence design tasks verify the effectiveness of our algorithm. Our code is available https://anonymous.4open.science/r/BIB-ICLR2023-Submission/README.md.

Under review as a conference paper at ICLR 2023
BIDIRECTIONAL LEARNING FOR OFFLINE MODEL- BASED BIOLOGICAL SEQUENCE DESIGN
Anonymous authors Paper under double-blind review
ABSTRACT
Offline model-based optimization aims maximize a black-box objective func- tion with static dataset of designs and their scores. In this paper, we focus on biological sequence design to maximize some sequence score. recent approach employs bidirectional learning, combining forward mapping for exploitation and a backward mapping for constraint, and it relies on the neural tangent kernel (NTK) of an infinitely wide network to build proxy model. Though effective, the NTK cannot learn features because of its parametrization, and its use prevents the incorporation of powerful pre- pre-trained Language Models (LMs) that can capture the rich biophysical information in millions of biological sequences. We adopt an alternative proxy model, adding a linear head to pre-trained LM,, and propose a lincarization scheme. This yields a closed-form loss and also takes into account the biophysical information in the pre-trained LM. In addition, the forward mapping and the backward mapping play different roles and thus deserve different weights during sequence optimization. To achieve this, we train an auxiliary model and leverage its weak supervision signal via a bi-level optimization framework to effec- tively learn how to balance the two mappings. Further, by extending the framework, we develop the first learning rate adaptation module Adeptive-1), which is com- patible with all gradient-baat based algorithms for offiine model basedd optimization. Experimental results on DNA/protein sequence design tasks verify the effectiveness of our algorithm. Our code is available here.
INTRODUCTION
Offline model-based based optimization aims to maximize black-box objective function with static dataset designs and their scores. This offline setting is realistic since in many real-world scenarios we do not have interactive access to the ground-truth evaluation. The design tasks of interest include material, aircraft, and biological sequence (Trabucco et al. (202T) In this paper, we focus on biological sequence design, including DNA sequence and protein sequence, with the goal of maximizing some specified property of these sequences.
A wide variety of methods have been proposed for biological sequence design, including evolutionary algorithms (Sinai et al., 2202)] Ren et al., 2002), reinforcement learning methods (Angermueller al., 2021). search/sampling using generative models (Brookes et al. 2019 Chan et al. 2021 and GFlowNets (Jain et al., (022). Recently, bradient-based techniques have emerged as an effective alternative (Trabucco et al.] 202T). These
2019). Bayesian optimization (Terayama
approaches first train a deep neural network (DNN) on the static dataset as proxy and then obtain the new designs by directly performing gradient ascent steps on the existing designs. Such methods have been widely used in biological sequence design (Norn et al. 2021 Tischer et al. 2020 Linder Seelig, 2020). One obstacle the out-of-distribution issue, where the trained proxy model inaccurate for the newly generated sequences.
To mitigate the out-of-distribution issue, recent work proposes regularization of the model Trabucco et al. 2021 Yu et al. 2021 Fu & Levine 2021) the design itself (Chen et al. 2022). The first category focuses on training a better proxy by introducing inductive biases such as robustness YYu et al. 20221. The second category introduces bidirectional learning (Chen et al.] 2022). which consists of a forward mapping and backward mapping, to optimize the design directly. Specifically, the backward mapping leverages the high-scoring design predict the static dataset and vice versa for

---

Under review as a conference paper at ICLR 2023
the forward mapping, which distills the information of the static dataset into the high-scoring design. This approach achieves state-of-the-- performances on a variety of tasks. Though effective, the proposed bidirectional learning relies the neural tangent kernel (NTK) of an infinite-width model to yield a closed-form loss, which is a key component of its successful operation. The NTK cannot learn features due its parameterization (Yang & Hu. 202T and thus the bidirectional learning cannot incorporate the wealth of biophysical information from Language Models (LMs) pre-trained over vast corpus of unlabelled sequences (EInaggar al. 202 Ji al. 202T).
To solve this issue, we construct a proxy model by combining a finite-width pre-trained LM with an additional layer. We then linearize the resultant proxy model, inspired by the recent progress in deep linearization (Achille et al., 202T Dukler et al. (022). This scheme not only yields a closed-for loss but also exploits the rich biophysical information that has been distilled in the pre-trained L.M. In addition, the forward mapping encourages exploitation in the sequence space and the backward issue. is important to maintain
mapping serves as a constraint mitigate the an appropriate balance between exploitation and constraint, and this can vary across design tasks as well as during the optimization process. We introduce a hyperparameter y to control the balance, and develop bi-level optimization framework Adaptive-y. In this framework, we train auxiliary model and leverage its weak supervision signal to effectively update ry. To sum up, we propose BIdirectional learning for model-based Biological sequence design (BIB). Last but not least, since the offline nature prohibits standard cross-validation strategies for hyperparameter tuning, all gradient- offline model-based algorithms preset the learning rate 7). There danger of poor selection, and to address this, we propose to extend Adaptive-r to Adaptive-r), which effectively adapts the learning rate 77 via the weak supervision signal from the trained auxiliary model. To the best of our knowledge, Adaptive-n) is the first learning rate adaptation module for gradient- based algorithms on offline model-basee optimization. Experiments on DNA and protein sequence design tasks verify the effectiveness of BIB and Adaptive-n.
To summarize, our contributions are three-fold:
Instead of adopting the NTK, we propose to construct proxy model by combining pre-trained biological LM with an additional trainable layer. We then linearize the proxy model, leveraging the recent progress on deep linearization. This yields a closed-form loss computation in bidirectional learning and allows us to exploit the rich biophysical information distilled into the LM via pre- training over millions of biological sequences.
We propose bi-level optimization framework Adaptive-r where leverage weak signals from an auxiliary model to achieve a satisfactory trade-off between exploitation and constraint.
We further extend this bi-level optimization framework to Adaptive-n. As the first learning rate tuning scheme in offline model-based optimization, Adaptive-n allows learning rate adaptation for any gradient- based algorithm.
2
PRELIMINARIES
2.1 OFFLINE MODEL-BASED OPTIMIZATION
Offline model-based optimization aims to find a design X to maximize some unknown objective f(X). This can be formally written as,
X arg max f(X),
(1)
X
where we have access to a size-N dataset D = ((Xi, 11311110 31)
{(XN. kN)} with X; representing a certain design and ys denoting the design score. In this paper, X, represents biological sequence design, including DNA and protein sequences, and y represents property of the biological sequence such as the fluorescence level of the green fluorescent protein (Sarkisyan et al. 2016).
2.2 BIOLOGICAL SEQUENCE REPRESENTATION
Following (Nor al. 2021 Killoran al. 2017 Linder & Scelig] 2021), ), adopt the position-
specific scoring matrix to represent a length- L protein sequence as X x where 20 represents
20 different kinds of amino acids. For real-world protein sequence, X [1, (0 < < I 1) is
2

---

Under review as a conference paper at ICLR 2023
The static dataset: TFBindB(r)
Forward mapping
high scoring data
G
Liz (Xa) 11... (x,)).
ang min lIv. SX(X))] 0
X1
T
X Unknown
A
A
A
(x) - lan (x. (*()))) = ang min lia f (X,)ll 810pP.
G
A
A
C
G
T
T
yt
21
4.5
Backward mapping
Yh 10.0
Figure 1:
Illustration of bidirectional learning Chen et al. (002 where (X,, y1) denotes the static dataset, yh is large predefined target and X the high-scoring design aim to find.
one-hot vector denoting one kind of amino acid. During optimization, [[ is a continuous vector and softmaz (x [I,:) represents the probability distribution of all 20 amino acids in the position l. Similarly, for a DNA sequence, we have X RR41 where 4 represents 4 different DNA bases.
The protein sequence X is fed into the embedding layer of the LM, which produces the embedding,
(so
(2)
The main block of the LM takes e as input and outputs biophysical features. The DNA LM, which adopts the k-mer representation, little different from protein LMs. See Appendix A. for details.
2.3 GRADIENT ASCENT ON SEQUENCE
A common approach to the posed offline model-based optimization problem is to train proxy fo(X) on the offline dataset,
1 N
0- arg min
Σ (Jo(X,) -
(3)
N
Then obtain the
design X by T gradient ascent steps: X = X; + V (X) x=x,, for t [0. 1],
(4)
where the high-scoring design X can be obtained as Xx.
Considering the discrete nature of biological sequences, the input of sa(-) should be discrete one-hot vectors. Following (Norn et al., 2021). we can perform the conversion and predict the score via:
X; softmax X:),
(5)
Z, onehol (argaal X:))
Ú (Z.).
(7)
Then the gradient regarding X; can approximated as,
dfo(Z.) dx;
dfe Zi) dē.
dzz dx,
where we unroll the matrices X,, X, and Z; as vectors D., & and 2; for notational convenience. This approximation allows to use backpropagation directly from the proxy to the sequence design X,. For brevity, we will still use fo (X;) to represent the proxy.
2.4 BIDIRECTIONAL LEARNING
As shown in Figure D bidirectional learning Chen et al. 2022). consists of two mappings: the forward mapping leverages the static dataset (Xt, yt) predict the score yh. of the high-scoring

---

Under review as a conference paper at ICLR 2023
design X1, and the backward mapping leverages the high-scoring design data (Xi, ynd to predict the static dataset (Xt, yt). The forward mapping loss can be written as
Linn (X,, = l1an - f0. (Xn))l
(9)
where 0* is given by
0' arg min y 1 xx) 3100
(10)
where B > is regularization parameter. The backward mapping loss can be written as
(x21(Xn) = |1y/ (x. (Xi)ll
(11)
where (*)),,) is given by
0 (X)) = arg min ll - (xx)) +
(12)
The high-scoring design X can be optimized by minimizing the bidirectional learning loss ((Xx) C12n Ci24 (Xx) + Lx21(
=
METHOD
In this section, we first illustrate how leverage deep linearization compute the bidirectional learning loss in a closed form. Subsequently, we introduce a hyperparameter y to control the balance between the forward mapping and the backward mapping. We then develop novel bi-level optimization framework Adaptive-r, which leverages a weak supervision signal from an auxiliary model to effectively update 7. Last but not least, we extend this framework to Adaptive-n, which enables us to adapt the learning rate for all gradient-based offine model-basee algorithms. We summarize our method in Algorithm|
3.1 DEEP LINEARIZATION FOR BIDIRECTIONAL LEARNING
In bidirectional learning, the backward mapping loss is intractable for a finite neural network, so Chen al. (2022) employ a neural network with infinite width, which yields closed-form loss via the NTK. This however makes it impossible to incorporate the rich biophysical information that has been distilled into a pre-trained LM (Yang & Hu (02T). Considering this, we construct a proxy model by combining finite-width pre-trained LM with additional layer. We then linearize the resultant proxy model, inspired by the recent progress in deep linearization which has established that an overparameterized DNN model close to its linearization (Achille al. 2021 Dukler et al. 2022).
Denote by 0, = (0pt. olin.) E RDx1 the proxy model parameters derived by combining the parameters of the pre-trained LM 4pt and a random initialization of the linear layer 0lin In this paper, we adopt the pre-trained DNABERT (Ji al., 2021) and Prot-BERT (Elnaggar al. (2221) models, and compute the average of token embeddings as the extracted feature, which is fed into the linear layer to build the proxy. Then we construct a linear approximation for the proxy model:
fo(X) ~ foo ()) + Vole, (X) (0 00),
(13)
where fine-tuning does not significantly change fo, then this linearization is good approximation. By leveraging this linearization, we can obtain a closed-form solution for Eq 12 as:
E R. Vefe, (X) E RIxD and Vele, (X) E Rx1 Intuitively, the
8' = (Vefe. (X,) Ve feo (X,)) + 3,33) Ve fea T feo (X₂)) + 0b. (14)
Building on this result, we can compute the bidirectional learning loss as:
Las (X,, =
2 (1ly Kx,x, K Xx + (1133)1100 (y) foo (x:)))2
(15)
+ |1y1 Kx,x, K x x, + -'0'
(yn Sao (x)))
where K(X,, X,) fefe. (X,)T Ve fer (X,). Following (Dukler et al. 2022). we can also only linearize the last layer of the network for simplicity, which detines the tollowing kernel,
K(X;, X,) = BERT(X.) BERTT (E)T(XX)),
(16)
where (ERT(X) denotes the feature of the sequence X extracted by BERT. Its kernel nature makes this approach suitable for small-data tasks (Arora et al. 2020), especially in drug discovery where the labeling cost of DNA/proteins is high.
4

---

Under review as a conference paper at ICLR 2023
Algorithm 1 Biidirectional Learning for Offline Model-based Biological Sequence Design
Input: Static dataset D = (Xr, 3/., predefined target score y/, 10, iterations T', pre-trained biological LM parameterized by Oo. auxiliary model fauz (-), regularization 8. Output: High-scoring design X;
1: Initialize X as the sequence with the highest score 'D
2: for < to T do
3: Leverage Adaptive-7 in Sec 3.2 to update the balance by Eq.
21
4: if Adapt learning rate then
5: Leverage Adaptive-1) in Sec 3.3 to update the learning rate by Eq. 23
6: Optimize X by minimizing the bidirectional learning loss Lai (X,;7) 7) in Eq. (77)
7:
X++1 = X, -- n Lbs (X,, ()))))
8:
Return X =XT
3.2 ADAPTIVE-7
The forward mapping and the backward mapping play different roles in the sequence optimization process: the forward mapping encourages the high-scoring sequence to search for a higher target score (exploitation) and the backward mapping serves constraint. Since different sequences require different degrees constraint, we introduce an extra hyperparameter E balance between the corresponding terms in the loss function:
[0, 1] to control the
Lai (X,, 7) >/22n (x)) + (1 ()2 (xx)
(17)
Thus = 1.0 corresponds to the forward mapping alone, 0 results in backward mapping, and 7 0.5 leads to the bidirectional loss of (Chen et al. (022).
It is non-trivial to determine the most suitable value for since we do not know the ground-truth score for a new design. One possible solution to train an auxiliary fauæ(.) to serve as a proxy evaluation. A reasonable auxiliary is simple regression model fitted to the offline dataset. Although this auxiliary model cannot yield ground-truth scores, can provide weak supervision signals to update % since the auxiliary model and the bidirectional learning provide complementary information. This is similar to co-teaching (Han et al. 2018) where two models leverage each other's view.
Formally, introduce the Adaptive-y framework. Given good choice of 7, the produced X,, is expected to have a high score faux (x,). based on which we can choose 7. To make the search for more efficient, we can formulate this process bi-level optimization problem:
7" arg max fauz (x)(
(18)
s.t. X arg min Cai (X
(19)
X
0Sauz (X,(1) to update y. Specifically, the inner level solution can be approximated via gradient descent step with a learning rate ry:
We can then use the
xx(7) = X₁, d.sn, [X
(20)
dxx,T
For the outer level, we update by hyper-gradient ascent:
dfaux (7)
dfouw x₂,)
dhrn (X,:7)
(21)
dy
dxy dxy
where we unroll the matrix XA, vector æn for better illustration.
3.3 ADAPTIVE-11
We now extend the Adaptive-ry framework to Adaptive-n. As the first learning rate adaptation module for offline model-based optimization, Adaptive-n compatible with all gradient-based algorithms and can effectively finetune the learning rate r) via the auxiliary model's weak supervision signal. All gradient-based methods that maximize Ce(X) with respect to X have the following general form:
X X, + 1/OPT(V Lo x))x) for & [0, T 1],
(22)
5

---

Under review as a conference paper at ICLR 2023
where 1) represents the learning rate of the optimizer. For methods such simple gradient ascent
(Grad), COMs (Trabucco et al. 2021). ROMA (Yu et al. (021) and NEMO (Fu Levine (021))
Co( is related to the proxy model (-); for BDI Chen al. (2022) and our proposed method, BIB, Lo (+) is the negative of the bidirectional learning loss, i.e., Lo Lbi-
Though the learning rate 7 can be adapted in some optimizers such as Adam (Kingma & Ba 2013).
these adaptations rely on only the past optimization history and do not consider the weak supervision signal from the auxiliary model. Our Adaptive-n) optimizes 1) by solving:
arg max faua (7))
(23)
where can be updated via gradient ascent methods. Considering the sequence optimization procedure is highly sensitive to the learning rate 17, we reset v to f/o at each iteration and update vj from 7jo.
dfauu (n))
7 = %/o
(24)
dr
In general, this serves stabilize the optimization procedure.
4
EXPERIMENTS
We conduct extensive experiments DNA and protein design tasks, and aim to answer three research questions: (1) How does BIB compare with state-of-the-art algorithms? (2) Is every design component necessary in BIB? (3) Does the Adaptive-r) module improve gradient-based methods?
4.1 BENCHMARK
We conduct experiments on two DNA tasks: TFBind8(r) and TFBind following (Chen et al., 2022) and three protein tasks: avGFP, AAV and E4B, in (Ren et al. 2022) which have the most data points. See See Appendix A.2 for details.
Following (Trbbucco et al. 2021 we select the top N 128 most promising sequences for each comparison method. Among these sequences, we report the maximum normalized ground truth score as the evaluation metric following (Ren et al. 2022).
4.2 COMPARISON METHODS
We compare BIB with two groups of baselines: the gradient-based methods and the non- non-gradient-
based methods. For fair comparison, the pre-trained LM is used for all methods involving
proxy and we don't 't finetune the L.M. The gradient-base methods include: 1) Grad: gradient ascent
on existing sequences to obtain new sequences; 2) COMs (Trabucco et al.) 202T): lower bounds the
DNN model by the ground-truth values and then applies gradient ascent; 3) ROMA (Yu et al. 202T): incorporates smoothness prior into the DNN model before gradient ascent steps; 4) NEMO (Fu
Levine 202T): leverages the normalized maximum maximum-likelihood estimator to bound the distance
between the DNN model and the ground-truth values; 5) BDI (Chen et al. 2022): ; adopts the infinitely
wide neural network and its NTK to yield a closed-form bidirectional learning loss.
The non-gradient-based methods include: 1) BO-qEI (Wilson et al. 2017 y builds acquisition function for sequence exploration; 2) CMA-ES (Hansen) 2006): estimates the covariance matrix to adjust the sequence distribution towards the high-scoring region; 3) AdaLead (Sinai al. 2020):
performs a hill-climbing search on the proxy and then queries the sequences with high predictions; 4) CbAS (Brookes al. 2019): builds generative model for sequences above property threshold and
gradually adapts the distribution by increasing the threshold: 5) PEX (Ren et al. 2022): prioritizes the evolutionary search for protein sequences with low mutation counts; 6) GENH (Chan et al. (02T):
enhances the score through a learned latent space.
4.3 TRAINING DETAILS
We follow the training setting in (Chen et al., (022) if not specified. We choose OPT as the Adam optimizer (Kingma & Ba 2015) for all gradient-based methods. We implement the auxiliary model

---

Under review as a conference paper at ICLR 2023
Table 1: Experimental results (maximum normalized ground truth score) for comparison.
Method
LFBind8())
TFBindTOie)
avGFP
AAV
E4B
Rank Mean
Rank Median
D[best)
0.242
0,248
0.742
0.452
0.224
B0-qEI
0.940 + 0.032
01585 + 0.028
1.700
0.020
0.591 + 0.002
0.436 + 0.004
6.0/12
7.0/12
CMA-ES
0.930 +: 0.034
0,617 1:117::00 0.031
5.488
= 0.006
0.470 - 0.006
0.748 + 0.009
6.0/12
6.0/12
AdaLead
0.941 + 0.032
0.602
+.002:00022
1.611
0.008
0.581 0.002
0:433 +: 0.003
6.2/12
8.0/12
CbAS
0.878 = 0.049
0.030
1.371
0.016
0.543 4 0.002
0.349 + 0.003
8.6/12
10.0/ 12
PEX
0.924 + 0.041
0.612 +: 0.026
0.019 1.516
0.388 + 0.002
0.397 + 0.004
7,2/12
8.0/12
GENH
0.323
0.000
0.448
0.000
0.835
0.000
0.452 0.000
0.228
0.000
11.4/12
11.0/12
Grad
0.941
0.026
0.630
0.029
4.869
0,042
0.463 + 0.005
1219 + 0.061
4.8/12
4.0/12
COMs
0.921 + 0.039
0.637
+ 0.065
3.873 = 0.080
0.511 + 0.005
0.829 + 0.026
5.6/12
5.0/12
ROMA
0.926 0.032
0.634
5.621
0.143
0.471 0.005
1.198 0.042
4.8/12
4.0/12
NEMO
0.930 + 0.038
0.632
+ 0.024
4.624 0.087
0.505 + 0.005
1.036
+ 0.046
5.0/12
5.0/12
BDI
0.823 + 0.000
0.678 1: 0.000
0.712 0.000
0.452 + 0.000
0.221 + 0.000
9.4/12
11.0/12
BIB (oan)
0.952 = 0.033
0.639
0.032
8.084
0.224
0,501
0.007
1,255
Σ 0.029
2,4/12
1.0/12
as linear layer with the feature from the pre-trained LM. We set the number of iterations T 25 for all experiments following (Norn et al. [02T]] and as 0.1 following (Chen et al.) (002)). We run every setting over 16 trials and report the average score. See Appendix A..3 for other details.
4.4 RESULTS AND ANALYSIS
We report all experimental results in Table and plot the ranking statisties in Figure 2 We make the following observations. (1) As shown in Table D BIB consistently outperforms the Grad method on all tasks, which demonstrates that our BIB can effectively mitigate the autoo--distribution issue. (2) Furthermore, BIB outperforms BDI on 4 out of 5 tasks, which demonstrates the ef- fectiveness of the pre-trained biological LM over NTK. The reason why BDI outperforms BIB on TFBind10(r) may be that short sequences do not rely much the rich sequential informa- tion from the pre-trained LM. (3) As shown in Figure [ the gradient- based methods generally perform better than the non-gradient-basee methods, as also observed by Trabucco et al. (2021) (4) The gradient-basee methods are inferior for the AAV task.
One possible reason is that the design space of AAV (2028) is much smaller than those of avGFP (20) and E4B (2), which makes the generative modeling and evolutionary algo- rithms more suitable. (5) This conjecture is also supported by the experimental results on two DNA design tasks. We compute the average ranking of gradient-
80-4B)
CHALES
Acstesd:
cbss
rtx
GENU.
Gmd-
methods and non-gre non-grad-ban based methods on TFBind10G1 O(r) as 3.5 and 9.5, re- spectively, and the average ranking of gradient-based methods and non-gran con-grad-bur -radd-adeddd methods on TFBind8(r) as 5.8 and 6.8, re- spectively. The advantage of gradient-basee based methods are larger (9.5- 3.5 = 6.0) TFBind than that (6.8 5.8 = 1.0) in TFBind8(r). (6) The generative modeling methods CbAS and GENH yield poor results on all tasks, probably because the high-dimensional data distribution is very hard model.
101112
Rank
Figure 2: Rank minima and max- ima are represented by whiskers; vertical lines and black triangles de- note medians and means.
(7) Overall, BIB attains the best performance in 3 out of 5 tasks and achieves the best ranking results as shown in Table Figure
We also visualize the trend of performance (the maximum normalized ground truth score) and trade- ([a) a) and GvGFP in Figure Bb). The performance generally increases with the time step T and then stabilizes, which demonstrates the effectiveness and robustness of BIB. Furthermore, we find that the 7 values of TFBind8(r) and GvGFP generally increase first. This means that BIB reduces the impact of the constraint to encourage a more aggressive search for a high target value during the initial phase. Then 7 of TFBind8(r) continues to increase while the 7 of avGFP decreases. We conjecture that the difference is caused by the sequence length. Small mutations of a biological sequence are enough to yield a good candidate (Ren et al.) 2022). For the length-239 protein av GVGFP, dramatic mutations 1) are not necessary and 2) can easily lead to out-of-distribution points. The weak supervision signal from the auxiliary model therefore encourages a tighter constraint towards the static dataset. By contrast, the DNA sequence is relatively short and more widespread search of the sequence space can yield better results. To investigate this conjecture, we further visualize the trend of E4B in Figure 3). E4B also has long sequences (102) and can observe its similar sest-increas-thhe-deeee
off y as a function of T on TFBind8(r) in Figure
trend, although is not pronounced.
7

---

Under review as a conference paper at ICLR 2023
1.0
1.0
1.0
Performannt
0.0
Performance
Performanze
1.5
0.8
0.6
Trade-off y
Perfomance
6
Performance
Trade
1.0
Perforance Tradeoff
Trade-
0.4
2
0.5
10
20
30
40
(a) (FBind8(r).
(h) avGFP.
(c) E4B.
Figure 3: Trend of performance and trade-off as function of T.
Table 2: Ablation studies on BIB components.
Task
y= 0.0
7= 1.0
? = 0.5
Y = 0.5 Joint
BIB
BIB Ada-1)
TFBind8(r)
0.936
0.933
0.947
0.935
0.952
0.956
100r
0.611
0.637
0.616
0.622
0.639
0.639
avGFP
6.051
7.588
7.940
7.920
8.084
8.197
AAV
0.449
0.458
0.480
0.420
0.501
0.525
E4B
0.778
0.903
1.198
1.176
1.255
1.301
4.5 ABLATION STUDIES
In this subsection, we conduct ablation studies to verify the effectiveness the forward mapping. the backward mapping, and the Adaptive-y module of BIB. We report the experimental results in Table 2
Forward mapping & Backward mapping. We can observe that bidirectional learning (? = 0.5) performs better than both forward mapping (y 1.0) and backward mapping (y 0.0) alone in most tasks, which demonstrates the effectiveness of forward mapping and backward mapping. The advantage of bidirectional mappings the forward mapping is larger in the long-sequence tasks like av GvGFP (238) and E4B (102) compared with the short-sequence tasks. A possible explanation is that the constraint is more important for long sequence tasks than short sequence design since the search space is large and many mutations can easily go out of distribution.
Adaptive-y. BIB learns y and this leads to improvements over bidirectional mappings (7 0.5) for all tasks, verifying the effectiveness of Adaptive-ry We also consider the following variant,
X arg min Las (X,, 0.5) Sanx (Xn), X
(25)
which jointly optimizes the bidirectional learning loss Lx (X,, 0.5) and the auxiliary term (X).. We found this yields similar even worse results than pure bidirectional learning. The reason may be that the weak supervision signal from fauz (X,) can serve as a guide to update the scalar ?y but not component of the main optimization objective that directly updates the sequence.
In the final column of Table we examine the performance of the Adaptive-n module. Adding this module leads improvements all five tasks, which demonstrates its effectiveness.
4.6 ADAPTIVE-11
In this subsection, we aim to further demonstrate the effectiveness of the Adaptive-1) module on all six gradient-based methods. We conduct experiments on two tasks: TFBind8(r) and avGFP. Since the use of the infinitely wide neural network Icads to poor performance for BDI, modify its implementation via deep lincarization so that it can make use of the pre-trained LM.
As shown in Table Adaptive-1 provides consistent gain for all scenarios, which demonstrates the widespread applicability and effectiveness of the module. Furthermore, Adaptive-n) leads to a maximum improvement of 1.4% in TFBind8(r) and 12.5% in GvGFP. ROMA is the algorithm that benefits the most. One possible explanation is that ROMA incorporates local smoothness prior that leads to stable gradients, with which Adaptive-1) can be more effective. Similar Sec 4.5
8

---

Under review as a conference paper at ICLR 2023
Table Adaptive-y on all gradient-based methods.
TFBind8(r)
avGFP
Method
Grad
COMs
ROMA
NEMO
BDI
BIB
Grad
COMs
ROMA
NEMO
BDI
BIB
Normal
0.941
0.921
0.926
0.930
0.347
0.952
4.869
3.873
5.621
4.624
7,940
8.084
Joint
0.941
0.921
0.931
0.932
0.935
0.925
4.869
3.836
6.438
3.078
7.920
7.823
Gain
0.000
0.000
0.005
0.002
0.008
0.027
0.000
-0.037
0.116
-1546
-0.020
-02261
Ada-1
0.941
0.928
0.939
0.935
0.951
0.956
5.235
4.027
6.322
4.658
7.960
8,197
Gain
0.000
0.007
0.013
0.005
0.004
0.004
0.366
0.154
0.701
0.034
0.026
0.113
consider the following variant,
X arg max (X)) + faux (Xi)) X1.
(26)
which performs joint optimization instead of bi-level optimization on two objectives. As shown in T joint optimization generally deteriorates the performance. This again verifies that the auxiliary model can only serve guide instead of contributing to the main objective.
5 RELATED WORK
Biological sequence design. There has been wide range of algorithms for biological sequence design. Evolutionary algorithms (Sinai al. 20200 Ren al. (022) leverage the learned surrogate model to provide evolution guidance towards the high-scoring region. Angermueller et al. (2019) propose a flexible reinforcement learning framework where sequence design a sequential decision- making problem. Bayesian optimization methods propose candidate solutions via acquisition function (Terayama et al.) 2021). Deep generative model methods design sequences in the latent space (Chan et al. 22221) p or gradually adapt the distribution towards the high-scoring region (Brookes) ett al. 2019). GelowNets (Jain et al.] (022) amortize the cost of search over learning and encourage methods leverage surrogate model and its gradient information to maximize the desired property (Chen et al. 2022 Norn et al. [202]] Tischer et al. 2020 Linder &
diversity. Gradient-basee
Scelig 2020). Our proposed BIB belongs the last category and leverages the rich biophysical
information
(ji et al.
2021 Elnaggar et al. 2021 to directly optimize the biological sequence.
Offline model-based optimization. A majority of sequence design algorithms (Angermueller et al. 2019; (Siai et al. 2020; Ren et al., 2022) focus on the online setting where wet-lab experimental results in the current round are analyzed to propose candidates in the next round. The problem of this setting is that wet-lab experiments are often very expensive, and thus a pure data-driven, offline
approach is attractive and has received substantial research attention recently (Trabucco et al. 2022
Kolli et al.] 2022). Gradient-based methods have proven to be effective (Trabucco et al. |021] Yu ct al. 2021 Fu & Levine 202T] Chen et al., 2022). Among these algorithms, CChn et al. (2022)
propose bidirectional mappings to distill information from the static dataset into a high-scoring design, which achieves
art performances on a variety of tasks. However, this bidirectional learning is designed for general tasks, like robot and material design, and the rich biophysical
information in millions of biological sequences is ignored. In this paper, we leverage recent advances in deep linearization to incorporate the rich biophy information into bidirectional learning.
CONCLUSION
In this paper, we propose bidirectional learning for offline model-based biological sequence. Our work is built on the recently proposed bidirectional learning approach (Chen al., (022), which is designed for general inputs and relies on the NTK of an infinitely wide network to yield a closed-form loss computation. Though effective, the NTK cannot learn features. We build a proxy model using the pre-trained LM model with linear head and apply the deep linearization scheme to the proxy, which can yield a closed-form loss and incorporate the wealth biophysical information at the same time. In addition, we propose Adaptive-y to maintain a proper balance between the forward mapping and the backward mapping by leveraging the weak supervision signal from an auxiliary model. Based on this framework, we further propose Adaptive-n, the first learning rate adaptation strategy compatible with all gradient-basee based offline model-based algorithms. Experimental results on DNA and protein sequence design tasks verify the effectiveness of BIB and Adaptive-n)

---

Under review as a conference paper at ICLR 2023
7 ETHICS STATEMENT
Protein sequence design aims to find a protein sequence with a particular biological function, which has broad application scope. This can lead to improved drugs that are highly beneficial to society. For instance, designing the antibody protein for SARS-COV-2 can potentially save millions of human lives (Kumar et al., 202T) and designing novel anti-microbial peptides (short protein sequences) is central to tackling the growing public health risks caused by anti-microbial resistance (Murray et al. 2022). Unfortunately, is possible to direct the research results towards harmful purposes such as the design of biochemical weapons. As researchers, we believe that must be aware of the potential harm of any research outcomes, and carefully consider whether the possible benefits outweigh the risks of harmful consequences. We also must recognize that cannot control how the research may be used. In the case of this paper, we are confident that there is much greater chance that the research outcomes will have beneficial effect. We do not consider that there are any immediate ethical concerns with the research endeavour.
REPRODUCIBILITY STATEMENT
We provide the code implementation of BIB and Adaptive-n here and also attach the code in the supplementary material. We describe the DNA/protein benchmarks in Sec. and the training details in Sec. 4.3 We also explain how obtain the sequence embedding from the pre-trained LM and how to perform gradient ascent steps on the sequence in Sec.
REFERENCES
Alessandro Achille, Aditya Golatkar, Avinash Ravichandran, Marzia Polito, and Stefano Soatto. Lqf: linear quadratic fine-tuning, In Proc. Comp. Vision. Pattern. Rec.(CVPR), 2021.
Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. In Proc. Inf. Conf. Learning Rep. (ICLR), 2019.
Sanjeev Arora, Simon S. Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the power of infinitely wide deep nets on small-data tasks. Proc. Int. Conf. Learning Rep. (ICLR), 2020.
Luis Barrera et al. Survey of variation in human transcription factors reveals prevalent DNA binding changes. Science, 2016.
David Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. In Proc. Int. Conf. Machine Learning (ICML), 2019.
Drew H Bryant, Ali Bashir, Sam Sinai, Nina K. Jain, Pierce J Ogden, Patrick F Riley, George M Church, Lucy Colwell, and Eric D Kelsic, Deep diversification of an AAV capsid protein by machine learning, Nature Biotechnology. 2021.
Alvin Chan, Ali Madani, Ben Krause, and Nikhil Naik. Deep extrapolation for attribute-enhance generation. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2021.
Can Chen, Yingxue Zhang, Jie Fu, Xue Liu, and Mark Coates. Bidirectional learning for offline infinite-width model-based optimization. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2022.
Yonatan Dukler, Alessandro Achille, Giovanni Paolini, Avinash Ravichandran, Marzia Polito, and Stefano Soatto. DIVA: Dataset derivative of a learning task. In Proc. Int. Conf. Learning Rep. (ICLR), 2022.
Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Yu Wang. Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, et al. ProtTrans: towards cracking the language of lifes code through self-super deep learning and high performance computing. IEEE Trans. Pattern Analysis and Machine Intelligence, 2021.
10

---

Under review as a conference paper at ICLR 2023
Justin Fu and Sergey Levine. Offline model-based optimization via normalized maximum likelihood estimation. Proc. Int. Conf. Learning Rep. (ICLR), 2021.
Bo Han, Quanming Yao, Xingrui Yu, Gang Niu, Miao Xu, Weihua Hu, Ivor Tsang, and Masashi Sugiyama. Co-teaching: robust training of deep neural networks with extremely noisy labels. Proc. Adv. Neur. Inf. Pruc. Sys/ (NeurIPS), 2018.
Nikolaus Hansen. The CMA evolution strategy: a comparing review. Towards A New Evolutionary Computation, 2006.
Moksh Jain, Emmanuel Bengio, Alex Grmmandez-Garcica Jarrid Rector- Rector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with GFlowNets. In Proc. Int. Conf. Machine Learning (ICML), 2022.
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.
Nathan Killoran, Leo J Lee, Andrew Delong, David Duvenaud, and Brendan J Frey. Generating and designing DNA with deep generative models. arXiv preprint
2017.
Diederik Kingma and Jimmy Ba. Adam: a method for stochastic optimization. In Proc. lnt. Conf. Learning Rep. (ICLR), 2015.
Sathvik Kolli, Amy X. Lu, Xinyang Geng, Aviral Kumar, and Sergey Levine. Data-driven optimization for protein design: workflows, algorithms and metrics. In ICLR2022 Machine Learning for Drug Discovery, 2022.
Vikash Kumar, Haiguang Liu, and Chun Wu. Drug repurposing against SARS-CoV-2 receptor binding domain using ensemble-based virtual screening and molecular dynamics simulations. Computers in Biology and Medicine, 2021.
Johannes Linder and Georg Seelig. Fast differentiable DNA and protein sequence optimization for molecular design. arXiv preprint arXiv:2005.11275 2020.
Johannes Linder and Georg Seelig. Fast activation maximization for molecular sequence design. BMC Bioinformatics, 2021.
Christopher JL. Murray, Kevin Shunji Ikuta, Fablina Sharara, Lucien Swetschinski, Gisela Robles Aguilar, Authia Gray, Chieh Han, Catherine Bisignano, Puja Rao, Eve Wool, et al. Global burden of bacterial antimicrobial resistance in 2019; systematic analysis. The Lancet, 2022.
Christoffer Norn, Basile IM Wicky, David Juergens, Sirui Liu, David Kim, Doug Tischer, Brian Koepnick, Ivan Anishchenko, David Baker, and Sergey Ovchinnikov. Protein sequence design by conformational landscape optimization. Proceedings of the National Academy of Sciences, 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, al. Pytorch: an imperative style, high-performance deep learning library. Proc. Adv. Neur. Inf. Proc. Syst (NeurIPS), 2019.
Zhizhou Ren, Jiahan Li, Fan Ding, Yuan Zhou, Jianzhu Ma, and Jian Peng. Proximal exploration for model-guided protein sequence design. Proc. Int. Conf. Machine Learning (ICML), 2022.
Karen S Sarkisyan et al. Local fitness landscape of the green fluorescent protein. Nature, 2016.
Sam Sinai, Richard Wang, Alexander Whatley, Stewart Slocum, Elina Locane, and Eric D Kelsic. Adalead: a simple and robust adaptive greedy search algorithm for sequence design. arXiv preprini arXiv:2010.02144: 41, 2020.
Lea M Starita, Jonathan N Pruneda, Russell Lo, Douglas M Fowler, Helen J Kim, Joseph B Hiatt, Jay Shendure, Peter S Brzovic, Stanley Fields, and Rachel E Klevit. in an E3 ubiquitin ligase identified by high- thghthhroughhut mutagenesis. Proceedings of the National Academy of Sciences, 2013.
mutations
11

---

Under review as a conference paper at ICLR 2023
Kei Terayama, Masato Sumita, Ryo Tamura, and Koji Tsuda. Black-box optimization for automated discovery. Accounts of Chemical Research, 2021.
Doug Tischer, Sidney Lisanza, Jue Wang, Runze Dong, Ivan Anishchenko, Lukas F Milles, Sergey Ovchinnikov, and David Baker. Design of proteins presenting discontinuous functional sites using deep learning. Biorxiv, 2020.
Brandon Trabucco, Aviral Kumar, Xinyang Geng, and Sergey Levine. Conservative objective models for effective offline model-based optimization. In Proc. Int. Conf. Machine Learning (ICML), 2021.
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-Bench: benchmarks for data-driven offline model- based optimization. arXiv preprint arXiv:2202.08450, 2022.
James T Wilson, Riccardo Moriconi, Frank Hutter, and Mare Peter Deisenroth. The reparameterizztion trick for acquisition functions. arXiv preprint arXiv:/712.00422 2017.
Greg Yang and Edward J Hu. Tensor programs iv: feature learning in infinite-width neural networks.
In Proc. Int. Conf. Machine Learning (ICML), 2021.
Sihyun Yu, Sungsoo Ahn, Le Song, and Jinwoo Shin. Roma: robust model adaptation for offline model- based optimization. Proc. Adv. Neur: Inf. Proc. Syst (NeurIPS), 2021.
A APPENDIX
A.1 DNA EMBEDDING
To incorporate richer contextual information, the DNA LM Ji et al. (2021) adopts the k-mer sequence representation, which is widely used in DNA sequence analysis. For example, the sequence ATGGCT has its 3-mer representation as {ATG, TGG, GGC, GCT). In this paper, we adopt its 3-mer representation and compute the probability of the 3-mer token by multiplying the probabilities of the three individual bases. The 3-mer representation is then sent to the pre-trained DNA LM.
A.2 DATASET DETAILS
We conduct experiments on two DNA tasks following (Chen et al. 2022 and three protein tasks in 2022 which have the most data points. We report the dataset details in Table
(Ren et al.
DNA Task TFBind8(r). The goal is find length-8 DNA sequence maximize the binding activity score with a particular transcription factor, SIX6REFRI
Barrera et al. 2016). We sample
5000 data points for the offline algorithms following (Chnn al. 2022).
DNA Task 2 1FBind10(r). The task TFBind10(r) is the same as TFBind8(r) except that the goal is to find length- 10 DNA sequence. Both DNA tasks measure the entire search space and we adopt these measurements as the approximate ground-truth evaluation.
Protein Task 1 GvGFP. This task aims to find a protein sequence with approximately 239 amino acids to maximize the fluorescence level of Green Fluorescent Proteins (Sarkisyan et al., 2016). The task oracle is constructed by using the full unobserved dataset (around 52,000 points) following (Rnn) jet al. 2022). The oracle passes the average of the residue embeddings from the pre-trained Prot-T5 (Elnaggar et al., 202T into a linear layer and then fits the dataset. The following two task oracles take the same form. The offline algorithms
only access the lowest-scoring 26,000 data points.
Protein Task 2 AAV. The goal is to engineer a 28-amino acid segment (positions 561-588) of the VPI protein to remain viable for gene therapy (Bryant et al. 2021). We use the entire 284, 000 data points to build the oracle and the fowest-scoring 142, 000 points for the offline algorithms.
Protein Task 3 E4B. This task aims to design protein (around 102 amino acids) to maximize the ubiquitination rate to the target protein (Starita et al.) 2013). The full dataset consisting of around 100, 000 points is used to build the oracle and the bottom half is used for the offline algorithms.
The parameterization the oracle is different from that the regression model from two aspects: 1) model architecture; 2) pre-trained information source. First, the oracle adopts the Prot-T5 model
12

---

Under review as a conference paper at ICLR 2023
Table 4: Dataset details.
Task
Metric
Min of D
Max of D
Min of Dentire
Max of Dentire
TFBind8(r) TFBind10(r)
binding activity
0.000
0.242
0.000
1.000
binding activity
- -1.859
-0.869
-1.859
2.129
avGFP
fluorescence level
1.283
3.553
1.283
4.123
AAV
viruses viability
111.176
-1.399
-11.176
9.536
E4B
ubiquitination rate
-3.89
-0.841
-3.589
8.998
Table 5: Experimental results on different pre-traine pre-trained LMs for comparison.
Pre-trained LM
avGFP
AAV
E4B
ProtAlbert
3.367 0.456
0.418
E 0.004
0.552 = 0.023
(rodBert(
8.084 0.224
0.501
+ 0.007
1.255 = 0.029
ProtBert BFD
8,240 = 0.094
0.549 + 0.009
1.880 = 0.054
which consists of an encoder and a decoder, while the regression model adopts the Prot-BERT model which only has an encoder. Second, Prot-T5 is trained on the BFD and UniRef100 datasets and ProtBert trained on the UniRef50 dataset. These two points demonstrate that the oracle and the regression model are different function classes. We choose the Prot-T5 model the oracle because this is the state-of-the-an protein LM to extract features and recent work Elnaggar et al. (202T) has demonstrated its effectiveness. In order to test how related the Prot-T5 Brracle/ProttBBBT[PPpee models are, we trained them on a sampled training dataset and compared the test predictions of the testing set. By evaluating the Pearson correlation coefficient (PCC) between the two prediction errors PCC(ProUTS predictions test labels, ProtBERT predictions -test labels), we obtain - -..053 on avGFP, -0.0005 on AAV, and -0.0062 on E4B. These results suggest that the two models are not strongly related in terms of the predictions they form.
Following TTabucco et 2021 we select the top 128 most promising sequences for each comparison method. Among these sequences, we report the maximum normalized ground truth score as the evaluation metric following (Ren et al. 2022).
A.3 TRAINING DETAILS
We use Pytorch Paszke et al. 2019) to run all experiments on one V100 GPU. Following the setting in Norn et al. (021). l). we introduce length-L protein sequence continuous random matrix X E 2 (X E R for DNA), initialized using a normal distribution with the mean 0 and the standard deviation of 0.01. To make this sequence correspond correctly to the candidate sequence, we exchange the largest value in X [[,,=] with the value in the amino acid index.
A.4 DIFFERENT PRETRAINED LMS
As shown in Table 5 we have tested the ProtBERT, ProtAlbert, and ProtBert-BBD models and found that better-quality models generally work better. The publicly available pre-trained DNA models are limited and thus we only perform experiments on the protein tasks. Elnaggar et al. 2021 demonstrate that the language model performances follow the ordering: BroBBert-BFD >> ProtBert > ProtAlbert. We can see that the performance ranks over the three protein tasks avGFP, AAV, and E4B are the same.
A.5 DIFFERENT DATASET SIZE
As shown in Table we have tested the performance of BDI as function of dataset size (N= 20, 40, 60, 80, 100) in TFBind8(r) and TFBind10(r) since they have exact oracle evaluations. We see that performance is already good for N=20 for TFBind8(r) and N=40 for TFBind
A.6 RANKING PERFORMANCE
As for prediction performances, the rank should be: a NN > linearized pre- pre-trained LM > NTK. We have conducted experiments to verify this. We sample half of the data, train model predict another
13

---

Under review as a conference paper at ICLR 2023
Table Experimental results on different size datasets for comparison.
Dataset size
20
40
60
80
100
TFBind8rr)
0.849
0.027
0.853
0.036
0.890
0.033
0.042
0.923 0.049
TFBind10(r)
0.248
0.000
0.586
0.035
0.602 + 0.023
0.616 + 0.024
0.632 0.036
Table 7: Mean squared prediction losses for comparison.
Method
(FBind8r)
avGFP
AAV
E4B
Fanetuned NN
0.001
1.130
0.411 0.197
5.148 1 0.074
0.683
Linearized NN
0.107 = 0.000
1.618 + 0.000
0.735 + 0.000
23,041 0.000
1.050
NTK
0.111 0.000
1.840 0.000
0.807 0.000
24.451 0.000
1.075 0.000
half data, and report the mean squared loss here and in Appendix A.6 Table A small mean squared loss indicates good prediction performance; thus, we have verified the above ranking order.
14


Reviews:
[{"review_id": "_zNfeVcokNo", "reviewer": "Reviewer_bELc", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper proposes a new method for offline model-based design, which is then applied to biological sequences. The method is able to use a kernel based on pretrained LM features, which allows the method to leverage additional knowledge. In addition, a learning rate adaptation method is proposed for bidirectional learning, with results showing improvement on the chosen benchmarks.", "summary_of_the_review": "The method is well presented and novel, and would be a reasonable contribution to ICLR. There are some issues with the evaluation and the numbers in the table appear not to be correct (or they are not correctly explained in the paper). If this issue is fixed, the paper would be a good candidate for acceptance, but it must be fixed first.", "strength_and_weaknesses": "Strengths\n* New gradient based optimization scheme that can leverage pretrained language models\n* Interesting approach for adapting the learning rate, which helps to deal with varying protein landscapes\n* Math is correct and is fairly easy to follow\n* Results appear relatively strong\n\nWeaknesses\n* Something is odd with the scale of the experiments. When I download the avGFP dataset, the maximum fitness value in the dataset is ~4.12, while the paper reports a value >8. The 4.12 number also appears consistent with the results in Ren et al. 2022. There also seems to be a scale mismatch in the AAV and E4B datasets between this paper and Ren et al. 2022 (in this paper, AAV fitness varies between 0.4 and 0.6, in Ren et al., it varies between -3.5 and 4.5; something similar is true for E4B).\n    - It would be useful if the authors listed the exact fitness metric for each task and perhaps the maximum value achievable in the dataset. Additionally authors should explain or correct any scale mismatch when comparing with previous work.\n    - This is the largest issue that must be addressed before publication - I would increase my score if this is addressed satisfactorily.\n* The authors argue that the incorporation of pretrained language model features is a point in favor of their method. It would be interesting to see how the method performs as a function of language model quality. You could use different models from the Elnaggar et al. or the various sizes of the ESM language models and see if improvements in language modeling result in improvements in optimization, which would more directly show that these features have a significant effect on performance.\n* Hard to follow what metrics are being reported in the table - could this be added to the table caption? It’s actually never explicitly said that the reported number is maximum fitness achieved, but I’m assuming it is\n* It would be nice to motivate the problem setting better. I have a hard time understanding when this class of method would be useful. It purports to very quickly optimally sample a set of designs (# iterations is set to 25 in this evaluation). But it also requires a dataset of >1000 training points to start out (even if this is limited to low scoring points). That suggests that you 1) must already have a high throughput way of screening fitness values for this task but also 2) want to further optimize a very constrained number of samples. In the case where you are able to generate the low-scoring dataset of size 1000+, does it matter if you find the optimal solution after 25 iterations or after 500 iterations?\n* Related to the above point, it would be good to see the performance of the method as a function of the size of the initial training dataset. How well does the method function in a true low-N setting (e.g. < 100 total samples).\n\nMinor Comments\n* Section 4.4, points (4) and (5) argue that the length of the sequence is a factor in whether gradient-based methods have a large advantage  over other methods. While I agree this makes sense intuitively, I'm not sure the evidence in the paper supports it. Gradient-based methods still seem to outperform on the DNA tasks, which have a much smaller search space. Additionally, the gap is larger in TFBind10 vs. TFBind8, but not clear exactly how much larger or whether this is actually significant.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "CYXJ4LYSNZ", "reviewer": "Reviewer_q1Be", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "A recent 'bidirectional learning' approach was proposed for offline model-based optimization. This paper extends it with a few tweaks: instead of using an NTK kernel, use the linearization of a pretrained foundation model, and tune optimizer hyper-parameters automatically by gradient descent using an auxiliary model to provide pseudo-labels for held-out data. The experiments focus on protein/dna design tasks and consider a number of methods from recent work.", "summary_of_the_review": "A have some important reservations about the paper below.\n \n==Parametrization of the oracle vs. the regression model==\nFor datasets where the entire search space hasn't been measured (i.e., everything other than the tfbind datasets), you define the ground truth f(x) by fitting a regression model (the 'oracle') on the entire dataset, such that this model can be queried on any x. I'm very concerned that the same overall model architecture is used for the oracle, the methods' function approximators f_theta, and f_aux: \n*Oracle: \"The oracle passes the average of the residue embeddings from the pre-trained Prot-T5 (Elnaggar et al., 2021) into a linear layer and then fits the dataset. The following two task oracles take the same form.\"\n*Function approximator: \"In this paper, we adopt the pre-trained DNABERT (Ji et al., 2021) and Prot-BERT (Elnaggar et al., 2021) models, and compute the average of token hidden embeddings as the extracted feature, which is fed into the linear layer to build the proxy. … Following (Dukler et al., 2022), we can also only linearize the last layer of the network for simplicity.\"\n*Auxiliary model f_aux: \"\"We implement the auxiliary model as a linear layer with the feature from the pre-trained LM?\"\n\nThis setup will unfairly favor methods that don't rely on a function approximator (i.e., it will penalize methods like CMA-ES), since the function approximator will have low approximation error for fitting the oracle, since they're from the same function class. It will additionally favor methods that use f_aux and model linearization, since these closely mirror the structure of the oracle. This is not a realistic setup. Model-misspecification will always be present when trying to model real experimental data with a neural network.\n\nAn alternative approach would avoid the oracle by evaluating optimization methods' ability to rank a held-out set of examples from the dataset (ranking them by the methods' associated acquisition function). I would have preferred to see some experiments seeing the second approach, since it removes the above issues.\n\n==Presentation of the paper's method==\nI found section 3.1 extremely confusing. There should be a background section that explains the bidirectional learning approach fully. The paper should be self-contained. Right now, readers will not be able to understand the explanation of how it works. The datasets Xh and Xl are never defined. I was completely unclear to me where yh comes from. I assumed that Xh were prospective designs and yh was unobserved. It was also never explained how the two models are combined into an acquisition function.\n\n==Complex model parameterization==\nI work in this exact research field. In my experience, simple linear models are often quite competitive when modeling these datasets. I would have appreciated some demonstration that the particular modeling choice you employ is sensible.\n\n==Presentation of experimental results==\nIt was very difficult for me to understand if any of the difference between methods' performances in Table 1 are statistically significant. As far as I can tell, all of these optimization methods are randomized. How much variance in performance is due to this randomness? I would have appreciated some error bars that convey this variance.\n\n==Role/Importance of f_aux==\nsec 3.2 is something you could do for basically ML system: just train two models and use M1 to provide aux labels for hyper tuning. for this to work, the aux model needs to be accurate, but if you're able to fit a good aux model, then isn't the overall problem easy?\n\n==Additional issues==\n\n\"Both DNA tasks have exact oracles for ground-truth evaluation.\"\nThis isn't quite true. The datasets measure the entire search space, but these measurements come from a noisy wet-lab experiment. We never observe the true noiseless value.\n\nFig 2 was not compelling to me. Where in sec 3.2 does gamma depend on T? Are you simply plotting the trajectory of the learned scalar gamma over the course of hyper-tuning. The shape of this trajectory should not be important. \n", "strength_and_weaknesses": "=Strengths=\nPaper contributes to an existing active thread on offline model-based optimization, which is an interesting and important area.\nPaper tackles offline biological sequence design, which is a practical application with large potential impact.\nPaper combines a number of modern techniques, such as neural network linearization, gradient-based hyper-parameter tuning using bilevel optimization, etc.\n\n=Weakness=\nI found the exposition of the bidirectional learning method extremely confusing.\nSee my full review below regarding my hesitations around the setup of the experiments.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "3eeB1cLeLt7", "reviewer": "Reviewer_Ay2G", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper tackles the model-based optimization problem for biological sequences. They extend the NTK-based bidirectional learning approach from previous work which essentially attempts to regularize according to features that transfer between low-fitness and high-fitness sequences. Their main contribution is twofold (1) they extend Chen et al 2022 by replacing the NTK representation with a pretrained language model with a linear layer atop and (2) they propose adaptive hyperparameter tuning strategies to improve the optimization scheme. They then benchmark their method against biological sequence datasets.", "summary_of_the_review": "The paper has a nice set of ideas that could plausibly provide an improvement upon bidirectional learning proposed in (Chen et al 2022). However, the lack of care taken toward the presentation and experimental section in particular significantly hampers the ability of the paper to be published in its current form.", "strength_and_weaknesses": "Strengths\n- The paper puts together an interesting set of ideas(bilevel optimization, deep linearization) that help resolve a weakness of prior work.\n- The paper seems to perform better on benchmark for biological sequence design tasks.\n- Adaptive hyperparameter tuning is extremely useful in model-based optimization settings due to the difficulty of tuning hyperparameters on out-of-distribution datasets.\n\nMajor Weaknesses\n- Lack of clarity on major experimental details makes interpretation of results difficult: Is using the pretrained LM for all proxy methods shown to perform better than a standard CNN-type model as this has been shown to work well in FLIP (Dallago et al 2021)? Is there finetuning done on the LM used as a proxy for other methods? What is a \"task oracle\"? How were X_h and y_h initialized?  \n- The central premise of the paper is that the following rank order is expected for surrogate model models: a NN(perhaps finetuned version of pretrained LM ) which is intractable to optimize against > linearized pretrained LM > NTK. However, this is never shown in any ablation studies rather intuition about biophysical features are referenced. The paper would be strengthened a lot if it was shown that such a representation indicates a better surrogate model by demonstrating improved prediction performance(rather than optimization) on the biological sequence datasets shown later. This is because there is significant stochasticity in the ranking of the MBO methods depending on the experimental setting and hyperparameter choices so a better understanding of the central claim of the paper would\n\n\nMinor Weaknesses:\n- Confidence intervals should be provided for the scores.\n- How were indels handled in the AAV task (given the biological sequence representation assumes fixed length $L$)?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject