Protein Sequence Design in a Latent Space via Model-based Reinforcement Learning
Authors: Minji Lee
Luiz Felipe Vecchietti
Hyunkyu Jung
Hyunjoo Ro
Meeyoung Cha
Ho Min Kim
Venue: Submitted to ICLR 2023 (2023)
Keywords: Biological sequence design
Model-based reinforcement learning
Protein design
Representation learning

Abstract: Proteins are complex molecules responsible for different functions in the human body. Enhancing the functionality of a protein and/or cellular fitness can significantly impact various industries. However, their optimization remains challenging, and sequences generated by data-driven methods often fail in wet lab experiments. This study investigates the limitations of existing model-based sequence design methods and presents a novel optimization framework that can efficiently traverse the latent representation space instead of the protein sequence space. Our framework generates proteins with higher functionality and cellular fitness by modeling the sequence design task as a Markov decision process and applying model-based reinforcement learning. We discuss the results in a comprehensive evaluation of two distinct proteins, GPF and His3, along with the predicted structure of optimized sequences using deep learning-based structure prediction.

Under review as a conference paper at ICLR 2023
PROTEIN SEQUENCE DESIGN IN A LATENT SPACE VIA MODEL-BASED REINFORCEMENT LEARNING
Anonymous authors
Paper under double-blind review
ABSTRACT
Proteins are complex molecules responsible for different functions in the human body. Enhancing the functionality of a protein and/or cellular fitness can signif- icantly impact various industries. However, their optimization remains challeng- ing, and sequences generated by data-driven methods often fail in wet lab exper- iments. This study investigates the limitations of existing model-based sequence design methods and presents a novel optimization framework that can efficiently traverse the latent representation space instead of the protein sequence space. Our framework generates proteins with higher functionality and cellular fitness by modeling the sequence design task Markov decision process and applying model-based reinforcement learning. We discuss the results in a comprehensive evaluation of two distinct proteins, GFP and His3, along with the predicted struc- ture of optimized sequences using deep learning-basee structure prediction.
INTRODUCTION
Proteins mediate the fundamental processes cellular fitness and life. Iterated mutations on various proteins and natural selection during the biological evolution diversify traits, eventually accumulat- ing beneficial phenotypes. Similarly, in protein engineering and design, the directed evolution of proteins has proved to be an effective strategy for improving or altering the proteins' functions or cellular fitness for industrial, research, and therapeutic applications (Yang et al.) 2019 Huang et al.) 2016). However, the protein sequence space of possible combinations of 20 amino acids is too large to search exhaustively in the laboratory, even with high-throughput screening from the diversified library (Huang et al. 2016). In other words, directed evolution becomes trapped at local fitness max- ima where library diversification is insufficient to cross fitness valleys and access neighboring fitness peaks. Moreover, functional sequences this vast space of sequences are rare and overwhelmed by nonfunctional sequences.
To tackle the limitations, data-driven methods have been applied to protein sequence design. They used reinforcement learning (RL) (Angermueller et al. 2019), Bayesian optimization (Wu et al. 2017| Belanger et al.) [2019] Terayama et al. [021] Stanton et al. 2022). ) and generative models Jain al. [022]] Kumar& & Levine 2020 Hoffman et al. 2022) in model fashion, i.e., using a protein functionality predictor trained on experimental data to model the local landscape. De- spite the advances obtained by these methods, it is still challenging to generate optimized sequences that are experimentally validated. We suggest that the cause for this is two-fold. The first cause that the optimization process is usually performed by generating candidate sequences directly through amino acid substitutions (Belanger et al., 2019) additions (Angermueller et al. 2019). Given the vast search space, these methodologies are highly computationally inefficient and com- monly lead to the exploration of parts of the space with low chance of having functional proteins. In designing biological sequences, previous literature explored optimizing a learned latent repre- sentation space Hoanez--ombbrelle et al.] 2018 Stanton et al.) (022). Similarly, in this paper investigate the optimization of sequences via RL directly in a latent representation space rather than in the protein sequence space. Actions, e.g., small perturbations in the latent vector, taken in this representation space can intuitively be understood as walking through a local functionality/fitnss. landscape.
The second cause is related to the models used as an oracle for optimization. These models are trained on experimental data obtained for specifie function, covering only small portion of the

---

Under review as a conference paper at ICLR 2023
State [Latent vecton)
Action I"ertur oettoo))
0-)
(-.)
Palkyy
Sequarce
Ercaddr
Dacoda
Rewarrd
Funchionaliy
Enooddr
Piackstor
Desooda
compressed
Low dimension
Sequeniee
Baseline (leamed mutation)
Previous approach (earned addition)
Our approach
State: Action: mutation
State: achoo
Action: addition
State: lntent vector
Action: perturaation
Longth (200-)
Length 1200-1
Sequancs
Sequeros
Policy
Amino acids (20)
Amnino acids
Add unino
Figure 1: The framework's s overview describes (top left) the encoder-decoder architecture trained to represent protein sequences in a latent space and (top right) the RL framework. The state is defined as the representation in the latent space, while the action is the perturbation in this representation. The perturbed representation is decoded back to a protein sequence using a sequence decoder. The reward is based the functionality predicted by the functionality predictor (oracle). Bottom row shows three RL-based state and action modeling options.
protein space, and their accuracy is consequently restricted to this small region. Later, we will demonstrate that even the most advanced model-based optimizations can assign high functionality values for randomly generated protein sequences. These random sequences, however, are unlikely to be functional. To reduce false positives, we suggest augmenting the experimental data with random sequences (i.e., negative examples) assigned to low functionality value. Such augmentation also helps set boundaries around the experimental data distribution which the oracle can be trusted.
We model protein sequence design as a Markov Decision Process (MDP) to optimize latent rep- resentation. Our method trains the optimization policy using model-based deep reinforcement learning (RL) framework (Fig. []. At each timestep, the policy updates the latent representation by small perturbations to maximize protein functionality or cellular fitness, i.e., walking uphill through the local landscape until the end of an episode. We evaluate method in two evalua- tion tasks, optimizing the functionality of the green fluorescence protein (GFP) and the cell fitness of
dehydratase (His3). Our results show that the proposed framework design sequences with higher protein functionality and cellular fitness than existing methods. Ab- lation studies show that, based on the evaluation of various model options for state and action for the RL framework, the proposed latent representation update can successfully optimize the protein and search the vast design space. We provide visual evidence that the trained policy can traverse the local functionality landscape efficiently for GFP and His3. Our method is general and can also be applied to representations learned from protein structures, such those presented in Zhang et al. (202): Eguchi et al. (2022).
RELATED WORKS
Protein representation learning Representation learning methods aim to learn compact and ex- pressive features describing data, Since a protein is composed of sequence distinct amino acids (N=20), it can be interpreted as a large word in which each character is an amino acid. Due to the similarity with natural language processing (NLP) tasks, methods such as BERT (Devlin al. 2018 and GPT-2 (Radford et al. 2019 have been used to train protein language models (Alley al. 2019
Brandes et al. [022] Lin et al. 2022 Ferruz et al. 2022 Rives et al. 2021). p. The protein language model trained with 250 million sequences by Rives et al. (2021) has shown learn representations
containing meaningful information about biological properties and reflecting protein structure.
was shown that the learned representation could be generalized across different applications achiev-
2

---

Under review as a conference paper at ICLR 2023
ing btateof-the-ane results for supervised prediction of mutational effects. Representation learning has also been applied to learn protein structures (Zhang et al. 2022). An mattiodd--ppecifii structure encoder was trained using variational autoencoders by Eguchi et al. (2022).
Protein functionality prediction Similar to protein representation learning, for functionality pre- diction, methods inspired by NLP continue achieve high performance. For the TAPE bench-
mark (Rao et al., 2019), long short-term memory (LSTM) (Hochreiter & Schmidhuber, 1997) and
transformers (Vaswani et al, 2017) are the top performing methods. [Notin et al.] (2022) proposed mnnsffrmer-inspiree architecture named Tranception for autoregressive functionality prediction that achieve state-of-the-e performance in 87 protein deep mutational scanning (DMS) datasets. Representing the protein as an image and using a convolutional neural network architecture (He) jt al., 20T6) are also investigated in (Rao et al.) 2019). Also proposed was the use of represen- tations learned from protein language model for zero-shot predictions (Alley et al. 2019: Meier cet al. 202T). Utilizing protein structures as inputs, graph neural networks have been used to predict functionality in (Gligorijevié al.] (021). Finally, combining a sequence-basee and structure- based representation for functionality prediction using language models and geometric vector perceptrons proposed in (Wang et al. 2022).
Biological sequence design Various machine learning methods have been used to design bio- logical sequences. DynaPPO is a model-based RL algorithm proposed specifically for biological sequence design (Angermueller et al.) 2019). which generates sequence from left to right by adding amino acids. Bayesian optimization was applied for cell fitness maximization in (Brookes et al. 2019 |elangge et al. 2019; Terayama et al. Swersky et al. 2020; Stanton et al. 2022). From these approaches, Sttnton et al. (022) used Bayesian optimization to optimize the sequence di- rectly in latent space by training denoising autoencoder that learns representations for corrupted sequences. It is important to note that our approach, compared to Stantoon et al.) (2022), does not use the corrupt-and-denoit learning. An evolutionary algorithm for efficient biological sequence optimization was proposed in (Sinai al., 2020). Methods using generative models to search and sample optimized sequences are investigated in Broooes et al. 2019 Kumar & Levine 2020
idea and defines optimization as episodic task to use reinforcement
Jain et al. 2022; Hoffman et al.,
2022 Melnyk al. 2021
3
METHODOLOGY
This section describes the primary components of the model-based RL framework, depicted in Details of network architectures are provided in Sec. 5 and Appendix
Fig. I
3.1 PROTEIN REPRESENTATION LEARNING
Sequence Encoder Let E X' be (protein) sequence of length L, where L denotes the number of amino acids in the sequence and X' is the sequence space. Each element I; of the sequence is a discrete number representing one of 20 amino acids, and is the i-th element from x. Given æ, the encoder network produces representation y of dimensions (R. ). The sequence encoder is divided into two parts: (i) a pre- trained protein language model encoder used obtain latent embeddings, (ii) a dimensionality reduction step is used to obtain the final representation.
In the first part, pre-trained protein language model is leveraged (see Fig. ([a)) We employ the
ESM-2 model (Lin et al. (022) trained on the Unireñ50 dataset (Suzek et al. 2015). This model uses a BERT (Devlin et al. 2018) encoder and trained with 150 million parameters an unsupervised fashion (see Appendix A.T. ESM-2 is utilized to map mutation effects onto latent space in our model. Given the sequence v as an input, ESM-2 outputs a matrix g E Q of dimensions (L+ E), where E is the dimension size of the embeddings, and Q is the embedding space. In ESM-2, CLS (classification) token is prepended to the sequence, and an EOS (end-of-sequencce) token is appended to the sequence when generating the embeddings. They serve as ggquence-beginnngg and eequence-ending characters, respectively.
In the second part, the output of the encoder with dimensions (L+2, E) is preprocessed before being passed to the policy. We use only the CLS token embedding to reduce the dimension from (L+2, E)
to (F.,). However, this could still be large for the action space given the embedding dimension E
3

---

Under review as a conference paper at ICLR 2023
(a) Sequence encoder-decoder
CLB. CLB,ED3 [L+2,
(.)
[E]]
JR)
(++2.))
ESM-2 (BERT)
CLS
Pooling
Repeal
[L+2,2E]
(L+2,E)
(L+2.E)
[L]
Concatenne
Sen Sect Seruennc nenee
Linear
Cropout
LM Head
Wiid-om
ESM-2 (BERTT)
(b) Functionality predictor
(....
Sequence decoder
[..]
(L+2,E]
[[L+))EE)
Reshace
Linesr
Lineer
(BERTT
GeLU
GeLU
Functionally
Latent recresentatin Arnino acid sequeence
Funclionality decoder
Figure 2: Model Architecture
of even the smallest ESM-2 model is 480. hence use 1-dimensional adaptive average pooling to further reduce the dimension to (R.) and obtain the final representation.
Sequence decoder Given a reduced representation y size (R,), we recover the amino acid sequence using the sequence decoder d. The representation size expanded (E.) using linear layer. To recover the original (L + 2, EE) dimensions of the embeddings obtained by ESM-2, the dimensions of the reduced representation are first expanded. Then, this matrix concatenated with the wild-type embeddings obtained using ESM-2 of the pre-trained encoder of size (L+ 2, 2E followed by linear layer. After passing dropout layer, the recovered representation is passed to the decoder's language model head, which maps the representation back to the sequence space, i... predicting each amino acid the sequence from their associated embeddings. The output the recovered sequence
3.2 PROTEIN FUNCTIONALITY PREDICTION
The functionality predictor network depicted in Fig. L is trained as downstream task of the pre-trained ESM-2 encoder network. Given an input sequence v, the encoder network produces representation Q of dimensions (L + 2, E). We flatten the representation to vector of size ((L+2) E, that is passed to an architecture with two lincar layers followed by GeLU (Hendrycks) & Gimpel 2016 non-linear activation. The final linear layer predicts value TR, where A is the functionality metric to be maximized and is the set of real numbers representing possible values for this metric. While the encoder can be shared since did not modify the architecture of the encoder nor the weights, the functionality decoder fxx is assumed to be trained independently for each task. To train the functionality predictor, negative examples are included when sampling the training batch to prevent the model from assigning high values to sequences outside the experimental data distribution. Negative examples are defined as random sequences with low functionality/fines value, and are crucial data enhancement technique for reducing false positives. The augmentation can help set boundaries around the experimental data region that can be trusted.
3.3 PROTEIN SEQUENCE DESIGN VIA MODEL-BASED REINFORCEMENT LEARNING
The final component trains a policy using an off-policy RL algorithm that models the reward function based on the functionality predictor.
System Modeling Consider a policy ा that interacts with a fully observable environment during each episode. An episode is defined by the number of timesteps in which the policy takes actions
4

---

Under review as a conference paper at ICLR 2023
to optimize the target sequence in the latent space. We sample an initial state so Q from the representation space (state space) at the beginning of each episode. At each timestep t, the agent observes state 8, Q and selects where A is the action space. The action a, is defined as a perturbation in which a continuous value at E [-c,d] c] is chosen for the j-th dimension, and is hyperparameter that be tuned based on the representation distribution of the experimental data. The hyperparameter controls how conservative the policy when traversing the functionality or fitness landscape. Both St and a have R dimensions.
action a, A according to a stochastic policy π + A,
We model the environment using a Markov Decision Process (MDP) and assume that the next state $1+1 is conditioned only on the current state and action a,. This function calculates the element- wise sum of 8t and a, so that 8t+1 = 8t + az As a result of performing the action at, the agent receives the reward ry- The reward is function
= r(st, az), where is the reward function r : + IR, and I is the set of real numbers representing reward values. This data later used to train the policy using off-policy reinforcement learning algorithm. The agent interacts with the environment until reaching a terminal state, which is one the followings: (i) when the last timestep T reached; (ii) when f(((s)) for the current state is greater than hyperparameter threshold; and (iii) when s, is in part of the representation space Q where fd((s)) cannot be trusted. Condition (iii) states how far 8, is from the experimental data. Each timestep's 's data is stored in replay buffer as a 5-tuple (8t, O2, 814 Tt, m²) where m, identifies the terminal state.
Off-Policy Reinforcement Learning We further optimize the policy using Soft Actor Critic (SAC) (Haarnoja et al.] 2018), an off-policy entropy-regularizee reinforcement learning (RL) al- gorithm for continuous action spaces. Off-policy is a type of RL learning that learns to optimize independently of the agent's 's actions. SAC is trained to maximize future rewards and entropy to promote randomness in action space exploration. An optimal policy x² following this objective is given as
π = arg max E 11
Σ ((((s)))) + (H(((()))))
(I)
where 7 E [0, 1] is a discount factor, a is a temperature hyperparameter and H is entropy. Details regarding SAC optimization and how actions are sampled by the policy are given in Appendix A.3
4 RESULTS
4.1 EXPERIMENT SETUP
Datasets We chose proteins of different lengths and functions for robustness in evaluation: the green fluorescent protein (GFP) that can be found naturally in jellyfish and imidazoleglycerol- phosphate dehydratase (His3) that is a key enzyme in our human body. The dataset proposed in (Sarkisyan al. 2016 is used to train the GFP encoder-decoder and its functionality predictor. The dataset contains 54,025 mutant sequences, with log- flgorsuoreeeecee intensity associated with each sequence. The length L of protein sequences is 237. For the His3 protein, we used the dataset in (Pokusaeva et al. 2019) that consists of mutant sequences of its evolutionarill--reevn segments and is associated with growth rate. We processed the data 50,000 sequences with the length of L = 30. This dataset is used to train the His3 encoder-decoder and fitness predictor. Datasets are randomly split into non-overlapping training and testing sets with a ratio of 90: 10. The encoder, decoder, and functionality predictor use the aforementioned training set. The test set is used for evaluating the trained models and selecting initial states for optimization.
We train two separate oracles for each dataset to prevent the circular use of the func- tionality predictor: one is exclusively used for optimization (which we call optimization oracle), and another is used to calculate the performance metric in experiments (which we call evaluation oracle). We follow the guideline presented in Kolli et al. (2022) that proposes the use of multiple oracles with different numbers of layers, activation functions, and hyperparameters to improve the reliability of the optimization process. We describe the experiment details and the comparison of evaluation and optimization oracles" performance in Appendices A.2 and A.7
Oracles
5

---

Under review as a conference paper at ICLR 2023
Implementation Details The encoder and decoder use pre-trained ESM-2 model (Lin et al, 2022) with 150 million trainable parameters. Optimization and evaluation oracles are trained using difterent ESM-2 models (Lin et al. 2022). The optimization oracle is trained using the model with 150 million parameters while the evaluation oracle uses the model with 35 million parameters. The latent representation space is set to R 8, which sets the size of state and action vectors. The magnitude of the perturbation applied to each element of the action vector is set to 0.1. The episode length T' is to 20 for GFP and 10 for His3. We set three alternative experimental rewards: (1) a dense reward, defined as
(s)(s)) (2) an absolute reward, defined as a binary value of rs 1, [()(s:) > 7th and if otherwise (values for Tth are listed in Table 12). and (3) binary reward, defined
binary value of = if / [(d(()) > f(d(s=-1) 1)) and 0 if otherwise. To set the initial state so of GFP, we sampled sequences with low (1.4 to 1.7) and high (3.2 to 3.4) functionality values from the test set. For His3, we sampled sequences with fitness values between 0.5 and 0.8. Appendix[A] describes other details regarding selecting initial states, network architectures, training hyperparameters, and optimization of our method and baseline methods.
Evaluation Metrics We use three evaluation metrics as reported in Jn,Jain et al. (0222 performance, novelty, and diversity. We also consider two additional metries for robustness: the originality of op- timized sequences, named as original, i.e. sequences not contained in the training dataset, and the distance between the optimized sequence and wild-type, named as dist( WT). The performance eval- uation metric is calculated the mean predicted functionality from the top K generated sequences. The predicted functionality is obtained by using the evaluation oracle presented Section 3.2 Let the generated sequences be contained in the following set G = " .. ). performance is de-
fined as K (. The novelty evaluation metric is defined to assess if the policy is generating sequences similar the ones contained in the experimental data. Defining P the experimental data set containing the wild-type protein sequence, novelty is given as follows:
Σ
Σ
dist(gt. .P)),
P
E0 P,EP
where dist is defined the number of different amino acids of two sequences. The diversity eval- uation metric is defined as the mean of the number of amino acids that are different among the optimized sequences and defined as:
Σ
Σ
dist(gT,g) 99))
(3)
K(K -
0,-9-
9700" {
The original metric is defined
K E, 1(0o & Pl) and the distance from wild-type (WT) metric
1
is given as functionality of GFP, we include the presence of the chromophore region (residues SYG in the wild-type protein) in the optimized sequence, these residues are related to the ability to emit fluorescence.
K Lo, EG- dist(w, (,), where w is the wild-type sequence. When testing the protein
4.2 EXPERIMENT
We generate 100 sequences for each method, and the 10 highest-performing sequences are evaluated.
GFP Sequence Design The results obtained for GFP sequence design are described in Table T We compare with four optimization methods: CbAS (adaptive sampling) (Brookes et | 2019)) BO (bayesian optimization) (Swersky et al.) 2020). GFlowNet (generative model) (Jain et al.] 2022). and DynaPPO (reinforcement learning) (Angermueller et al. (219).. The term directed evolution refers to the average functionality values of initial states used in the optimization process. We also include random mutations of initial states. is seen in Table Table the proposed method outperforms directed evolution. Only the proposed method and CbAS optimize GFP effectively, whereas BO, GFlowNet, and DynaPPO achieve low performance. Even though method limits the action to a small step in the latent space, it gets higher novelty and diversity when compared to CbAS, which also limits the search space when optimizing. It was interesting to observe that two of the

---

Under review as a conference paper at ICLR 2023
GFP sequences optimized by the proposed method achieved higher predicted functionality values, 3.9705 and 3.8059, when compared to the experimental wild-type functionality value, 3.72. It is interesting observe that methods that achieve the highest distance from the wild-type sequence are the ones that achieve the lowest performance.
Model
Performance
Novelty
Original
dist (T)
Diversity
Chromophore
Ours
3,491 + 0.352
8,451
100%
7.700
6,311
100%
Directed evolution
3.287 0.237
7.704
6,849
4.858
100%
CbAS
7.712
80%
6,900
1.956
100%
Random-1
2.104
6.611
80%
7.186
7.716
100%
Random-5
2,280 + 0.275
13.91
100%
9.950
12.37
90%
Random-P
1.511:0070)
14.71
100%
14.15
14.62
100%
BO
0.581 = 0.095
36.96
100%
36.70
6.867
100%
DynaPPO
0.004 + 0.003
218.9
100%
219.3
224.1
0%
GFlowNet
0.000 0.002
199.4
100%
200.1
12.53
0%
Table I: Results obtained for GFP sequence design. Random-1 and Random-5 indicates random
mutations in I and positions, respectively. Random-P indicates random perturbation in the latent representation. Standard deviation for Novelty and dist(WT) are shown in Table g
His3 Sequence Design Table 2 displays the results for His3, where our framework achieves the best performance. DynaPPO and BO were unable to optimize His3 effectively. Given that the length of the His3 protein only 30, interesting observe that single random mutation an effective strategy. Note, however, that even though the proposed method achieves higher overall performance, its novelty is lower than other methods. This indicates that the decoder is recovering similar sequences within the representation space. The proposed method achieves higher novelty and diversity when compared to CbAS.
Model
Performance
Novelty
Original
dist(WT)
Diversity
Ours
0.945 = 0.091
8.361
60%
10.95
3.521
CbAS
0.749 0.157
7.287
90%
4.700
2.356
Random-
0,858 = 0.058
7.372
80%
7.350
7.716
Random-5
0.678 0.096
9.777
100%
8.950
12.37
Directed evolution
0.616 = 0.110
6.889
6.710
6.942
DynaPPO
-0.201 # 0.142
27.41
100%
26.70
27.47
BO
-0.313 = 0.065
26.17
100%
27.50
4.756
Table 2: Results obtained for His3 sequence design. Random-1 and Random-5 indicates random mutations in I and positions, respectively. Standard deviation for Novelty and dist(WT) are shown in Table
4.3 ABLATION STUDIES
State and Action Modeling Table shows how state and action modeling influence sequence de- sign in GFP. For sequence modeling, we examine two types: a representation in which each amino acid is represented by one-hot encoding and a latent vector that the encoder's output. We vestigate three types of action modeling: (i) multi discrete sequence generation, (ii) conditional autoregressive addition of amino acids proposed in (Angermueller et al.] 2019). and (iii) the pro- posed perturbation in the latent vector. Only the latent vector as the state and the perturbation as the action can optimize GFP and His3 tasks. Compared to the latent space learned based on the embed- dings from ESM-2, using the one-hot encoded mutant sequence input makes the identification of structural changes related to the protein's 's functional site challenging. Also, modeling the action as the addition of amino acids is ineffective for traversing the vast and sparse protein space.
Representation Analysis
We analyzed the influence of the optimization method on the latent vector an input. Our method, which employs reinforcement learning for taking actions that update

---

Under review as a conference paper at ICLR 2023
State
Action
GFP
His3
Latent vector
Perturbation on latent vector
3,491 + 0.352
0.945 = 0.091
Directed evolution
3.287 = 0.237
0.616 + 0.110
Sequence
Generate sequence
0.006 # 0.004
-0.148 # 0.043
Latent vector
Generate sequence
0.005 = 0.003
-0.139 0. 144
Sequence
Amino acid addition
0.004 + 0.003
-0.201 0.142
Table Comparison on the performance by the state and action modeling. For rows with the action defined as mutation in the protein sequence, train the policy using Proximal Policy Optimization (PPO) (Schulman et al.) (017)) to handle the multiple discrete action space. More details regarding the state and action modeling for this ablation presented in Appendix A.3
latent vector, is compared random perturbation and to the aforementioned BO method proposed by (Swersky et al. (2020), both of which take a latent vector as input. First, Swersky et al. (2020) performs better when using the latent vector input (2.601 0.912) than when using the sequence as input (0.581 0.095). implies that the latent vector provides rich information in low-dimensional space that is easier to optimize than the sequence space. Even with the same representation, method outperforms Swersky et al. 2020) while achieving competitive novelty.
Model
Performance
Novelty
Diversity
Chromophore
Ours
3.491 0.352
8.451 # 2.05
6.311
100%
Directed evolution
3.287 0.237
7.704 # 2.66
4.858
100%
Swersky al. (2020) latent space
2.601 0.912
8.077 # 2.58
6.600
100%
Random perturbation
1.511 = 0.797
14.71 # 5.90
14.616
100%
Table 4: Comparison on our method and BO when using latent representation as an input.
Average: 0.000
Average:
Flucresosnce GFP
Data printsts (85)
40
Predicted flucrescnnce on random sequerczes
Predictor nagative evamples
Predictor traindd without negative examples
Fluorespance
Figure 3: The predicted fluorescence of random sequences tested using functionality predictor trained with or without negative examples.
Training with Negative Examples We trained the GFP functionality predictor with the original training dataset and the training dataset augmented with 40,000 negative examples (i.e., 82% of the size of the original training dataset). We tested both predictors with 10,000 negative examples that were not seen during training. Fig. 3 shows that the functionality predictor trained without negative examples incorrectly predicts high value for non-funeii functionai sequences (mean=4.002), whereas the one trained with negative examples accurately assigns zero functionality (mean=0.0) to random non-functional sequences.
8

---

Under review as a conference paper at ICLR 2023
(8)
3.5
14
3.0
23
2.5
a
os
2.0
1.5
Oplimization path
(c)
(d)
1.2
0.9
0.8
Fitress
0.3
Opttimized state
Optimization leads to higher fitness Optimization leads to lower fitness
0.0
Initial state
Figure 4: Visualization of the climbing process through local functionality and fitness landscape.
(a) and (b) Resulting optimization carried out by the trained policy during a single episode for the GFP task. The z-axis represents the intensity of log fluorescence. In (b) it is seen the policy travers- the proposed method. (c) Fitness landscape
ing local optima due to the RL-based formulation for the His3 dataset. We annotated initial states (red) and optimized states (blue). (d) Optimization performed by the trained policy for the His3 task. The representation space reduced to two using I-SNE (Van der Maaten & Hinton 2008). Appendix describes the methodology used gener- ate these landscapes. An HTMML interactive version of these graphs can be found in the supplementary materials.
DISCUSSION
How the trained policy traverses the functionality landscape We qualitatively evaluate our trained policy by analyzing its ability to traverse the local landscape. Figs. 4 (a) and (b) visual- ize the optimization trajectory during one episode of GFP, indicating continued improvement of the fluorescence level. This demonstrates that the model can successfully climb up the local landscape throughout the episode. The policy can also traverse local optima thanks to the MDP formulation (See Appendix I.III. For His3, the cell fitness landscape of the entire dataset is shown in Fig. [ which interestingly depicts a valley region that corresponds to part of the space leading to low fitness. Also, we can observe that the optimized states scattered throughout the landscape, implying that our policy can design diverse proteins. Fig. d) presents optimization steps taken by the trained policy, which now shows that large (optimistic) perturbations often lead to failure in optimization.
CONCLUSION
This paper addressed the problem of optimizing protein functionality and cellular fitness using data- driven methods. We investigated the limitations of the latest model-based biological sequence de- sign methods and proposed a novel optimization framework that can efficiently traverse the latent representation space, as opposed to optimizing the protein sequence by mutations in its amino acid sequence, Our framework trained a policy that continually updates the latent representation through- out an episode. We modeled the problem as MDP to maximize future rewards, which allowed the agent to explore the landscape while learning an optimal policy. This special mechanism was critical to efficiently traversing the local functionality and fitness landscape. The proposed framework out- performed other baseline methods in performance and is competitive in terms novelty. We have also demonstrated that trained policy navigate local landscape effectively and traverse local optima. In the future, we intend to filter the generated sequences and conduct wet lab experiments validate the proposed framework.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Ethan C Alley, Grigory Khimulya, Surojit Biswas, Mohammed AlQuraishi, and George M Church. Unified rational protein engineering with sequence-based deep representation learning. Nature methods, 16(12):1315-133 132221222 2019.
Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. In International conference on learning representations, 2019.
David Belanger, Suhani Vora, Zelda Mariet, Ramya Deshpande, David Dohan, Christof Anger- mueller, Kevin Murphy, Olivier Chapelle, and Lucy Colwell. Biological sequences design using batched bayesian optimization. 2019.
Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. Proteinbert: A uni- versal deep-learning model of protein sequence and function. Bioinformatics, 38(8):2102-2110 2022.
David Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. In International conference machine learning, pp. 773-782. PMLR, 2019.
Jacob Devlin, Ming- Heig-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04886 2018.
Raphael R Eguchi, Christian A Choe, and Po-Ssu Huang. Ig-vae: Generative modeling of protein structure by direct 3d coordinate generation. PLOS computational biology, 18(6)=e100021) 2022.
Noelia Ferruz, Steffen Schmidt, and Birte Höcker. Protgpt2 is a deep unsupervised language model for protein design. Nature communications, 13(1):1-10, 2022.
Karl Pearson F.R.S. Liii lines and planes of closest fit to systems of points in space. The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science, 2(11):559-57,, 1901. doi: 10. 01880001788-400094272207
Vladimir Gligorijević, P Douglas Renfrew, Tomasz Kosciolek, Julia Koehler Leman, Daniel Beren- berg, Tommi Vatanen, Chris Chandler, Bryn C Taylor, lan M Fisk, Hera Vlamakis, et al. Structure- based protein function prediction using graph convolutional networks. Nature communications, 12(1):1-14, 2021.
Rafael Gómez-Bombareil Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez- Lengeling Dennis Sheberla, Jorge eguilera-Iparragune Timothy Hirzel, Ryan P Adams, and Alán Aspuru-Guzik. Guzik. Automatic chemical design using data-driven contin- uous representation of molecules. ACS central science, 4(2):268 276,, 2018.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and appli- cations. arXiv preprint arXiv:1812.05905, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Dan Hendrycks and Kevin Gimpel. 1861011606.08415. 2016.
Gaussian error linear units (gelus). arXiv preprint
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Samuel C Hoffman, Vijil Chenthamarakshan, Kahini Wadhawan, Pin-Yu Yu Chen, and Payel Das. Opti- mizing molecules using efficient queries from property evaluations. Nature Machine Intelligence, 4(1::21-31, 2022
Po-Ssu Huang. Scott E Boyken, and David Baker. The coming of age of de novo protein design. Nature, 537(7620):3003327 2016.
10

---

Under review as a conference paper at ICLR 2023
Moksh Jain, Emmanuel Bengio, Alex Grrnandez-Gareiia Jarrid Rector- Dector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gflownets. In International Conference on Machine Learning, pp. 9786-9801. PMLR, 2022.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583--888 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1412.6900. 2014.
Sathvik Kolli, Amy Lu, Xinyang Geng, Aviral Kumar, and Sergey Levine. Data-driven optimiza- tion for protein design: Workflows, algorithms and metrics. In ICLR2022 Machine Learning for Drug Discovery, 2022.
Aviral Kumar and Sergey Levine. Model inversion networks for model-based optimization. Ad- vances in Neural Information Processing Systems, 33:5126-51333 5137, 2020.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos San- tos Costa, Maryam Fazel-Zarandi,, Tom Sercu, Sal Candido, al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. bioRxiv, 2022.
Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alex Rives. Language models enable zero-shot prediction the effects mutations on protein function. Advances in Neural Information Processing Systems, 34:29287-29303, 2021.
Igor Melnyk, Payel Das, Vijil Chenthamarakshan, and Aurelie Lozano. Benchmarking deep gener- ative models for diverse antibody sequence design. arXiv preprint arXiv:211 2021.
Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena Hurtado, Aidan Gomez, Debora Marks, and Yarin Gal. Tranception: protein fitness prediction with autoregressive transform- ers and inference-time retrieval. In International Conference on Machine Learning, pp. 16990- 17017. PMLR, 2022.
Victoria O Pokusaeva, Dinara R Usmanova, Ekaterina V Putintseva, Lorena Espinar, Karen S Sark- isyan, Alexander Mishin, Natalya Bogatyreva, Dmitry Ivankov, Arseniy Akopyan, Sergey Ya Avvakumov, et al. An experimental assay of the interactions of amino acids from orthologous sequences shaping a complex fitness landscape. PLoS genetics, 1000)00,, 2019.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Peter Chen, John Canny, Pieter Abbeel, and Yun Song, Evaluating protein transfer learning with tape. Advances in neural infor- mation processing systems, 32, 2019.
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma, et al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, 118(15)ee216633111 2021.
Karen S Sarkisyan, Dmitry A Bolotin, Margarita V Meer, Dinara R Usmanova, Alexander S Mishin, George V Sharonov, Dmitry N Ivankov, Nina Bozhanova, Mikhail Baranov, Onuralp Soyle- mez, et al. Local fitness landscape of the green fluorescent protein. Nature, 533(7603):397-401, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: /107006777.644 2017.
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 preprint wrXiv:2000.02144 2020.
11

---

Under review as a conference paper at ICLR 2023
Samuel Stanton, Wesley Maddox, Nate Gruver, Phillip Maffettone, Emily Delaney, Peyton Green- side, and Andrew Gordon Wilson. Accelerating bayesian optimization for biological sequence design with denoising autoencoders. arXiv preprint arXh::2203.12742, 2022.
Baris E Suzek, Yuqi Wang, Hongzhan Huang, Peter B McGarvey, Cathy H Wu, and UniProt Consor- tium. Uniref clusters: a comprehensive and scalable alternative for improving sequence similarity searches. Bioinformaties, 31(6):926 9320-93222202 2015.
Kevin Swersky, Yulia Rubanova, David Dohan, and Kevin Murphy. Amortized bayesian optimiza- tion over discrete spaces. In Conference Uncertainty in Artificial Intelligence, pp. 769-778. PMLR, 2020.
Kei Terayama, Masato Sumita, Ryo Tamura, and Koji Tsuda. Black-box optimization for automated discovery. Accounts of Chemical Research, 54(6):133-1146, 2021.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine
learning research, 9(11), 2008.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention you need. Advances in neural informa- fion processing systems, 30, 2017.
Zichen Wang, Steven A Combs, Ryan Brand, Miguel Romero Calvo, Panpan Xu, George Price, Nataliya Golovach, Emmanuel Salawu, Colby J Wise, Sri Priya Ponnapalli, et al. Lm-gyp: an extensible sequence and structure informed deep learning framework for protein property predic- tion. Scientific reports, 12(1):1-12, 2022.
Jian Wu, Matthias Poloczek, Andrew G Wilson, and Peter Frazier. Bayesian optimization with gradients. Advances in neural information processing systems, 30, 2017.
Kevin K Yang, Zachary Wu, and Frances Arnold. Machine-learning-guts directed evolution for protein engineering. Nature methods, 16(8):687-694) 2019.
Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining. arXiv preprint 0611:2 2022.
12

---

Under review as a conference paper at ICLR 2023
A
APPENDIX
A.1 SEQUENCE ENCODER-DECODER
Next, we give more information regarding the architecture and the training process of the sequence encoder-decoder. The model architecture is seen in [a)). The ESM-2 encoder comprises token embedding layer and 30 transformer layers. Each transformer layer consists of multi-head attention followed by layer normalization layer and two fully connected layers with GeLU (Hendrycks & Gimpel 2016) non-linear activation. Language model head comprises of linear layer followed with GeLU activation and the linear layer using the weight of the embedding layer in the encoder. The number of attention heads is set 20. We used the pre-trained weights provided in (Lin et al.) 2022). The ESM-2 encoder weights are not updated during the fine-tuning process. The sequence encoder-decoder fine-tuned using masked language model objective. We define the loss function as the mean of cross entropy loss on mutated positions and the mean cross-entropy loss on non- mutated positions to ensure that the decoder also focuses on recovering mutated positions. Since we provide the wild-type representation during sequence recovery and mutants in the GFP dataset have maximum of 15 mutations, such a loss function important to prevent the decoder from predicting every sequence the wild-type. We train the model to the train set of each dataset for 16 epochs using the Adam optimizer (Kingma & Ba, 2014). The initial learning rate of Adam is set to 8e-5, with weight decay set to le-5, The learning rate is reduced by 0.8x every epoch. Table shows the performance of a sequence decoder by the chosen embedding size of R.
Top-1 Accuracy
Dataset
Embedding size
Mutated positions Non-mutated positions
8
0.463
0.9958
GFP
32
0.4762
0.9903
8
0.6439
0.8407 0.9022
His3
32
0.8194
Table 5: Performance of the sequence decoder for GFP and His3.
A.2 FUNCTIONALITY PREDICTOR
Functionality and fitness predictors for GFP and His3 are also based on the pre-trained ESM-2 protein language model. The ESM-2 encoder has same architecture as described in Appendix A and the decoder is described in Sec. 3 We fine-tuned the predictor on the train set of each dataset for 16 epochs using the Adam (Kingma & Ba 2014) optimizer with the initial learning rate set to 8e-5 and weight decay set to le-5. The ESM-2 encoder weights are not updated during the fine-tuning process. The negative samples for GFP and His3 are associated with a functionality/titess. value of 0. The GFP value of denotes an extremely low fluorescence level (10(0-3.72) = 0.0002 times the brightness of the wild type protein). The His3 value of 0 denotes the minimum fitness value in the dataset, ranging from [0, 1.63.. We used the mean squared error between empirically obtained log- fluorescence intensity and predicted fluorescence
the loss function for GFP and the empirically obtained fitness and predicted fitness as the loss function for His3.
A.2.1 OPTIMIZATION ORACLE
In Table we report the test Spearman's 's p and mean squared error (MSE) of the optimization oracle used for model-based optimization.
A.2.2 EVALUATION ORACLE
In Table 7 we report the test Spearman's p and mean squared error (MSE) of the evaluation oracle used for based optimization. Compared to the optimization oracle, the evaluation oracle utilize a different pre- pra-tra ESM-2 model which also contains a different number of parameters. The evaluation oracle and optimization oracle are trained with similar methodology.
13

---

Under review as a conference paper at ICLR 2023
Test set
Negative test set
Task
Spearman'ss
MSE
MSE
GFP His3
0.8426 0.6635
0.1436 0.0080
2.930e-6
0.0164
Table 6: Performance of the optimization oracle for GFP and His3
Test set
Negative test set
Task
Spearman's p
MSE
MSE
GFP His3
0.8320 0.6820
0.4359 0.0110
6.247c-05 0.0288
Table 7: Performance of the evaluation oracle for GFP and His3
A.2.3 TRAINING CBAS AND DYNAPPO WITH AND WITHOUT NEGATIVE SAMPLES
We conducted an ablation study by training two baseline optimization methods, CbAS and Dy- naPPO, using both functionality predictors. The optimized sequences are then visualized using predicted structures obtained by AlphaFold2. As shown in Fig. 5 CbAS managed to propose valid fluorescent protein when guided by an oracle trained with negative examples but failed when guided by oracle trained without negative examples. Even though the structure predicted for the CbAS model for this case seems unstable, the oracle trained without negative examples assigned very high log- flg-fuurescence intensity value (4.014) for the sequence. Therefore, argue that bio- logical sequence design should be guided and evaluated by an oracle trained with an experimental dataset augmented with negative examples.
A.3 SOFT ACTOR CRITIC (SAC)
Recapitulating, an optimal policy x" following the SAC objective is given as
arg max En
('(r(s +
(4)
where 7 E [0, 1) is discount factor, is temperature hyperparameter and H is entropy. Following a policy и the Q function or action-value function is defined as
(s, = Ex
Σ 7r((s,, a.) + I (((((-sss) 1
(5)
Am1
where the entropy is added from the first timestep. The Bellman equation for Q" at timestep t is given by
Q" a.) E,, [r(st, a.) + (( ē(+1) + aH (((-(((4444))
(6)
which he state 8t+1 used from the replay buffer and ăe+1 is sampled from the current policy. During training, SAC models three networks: policy and two Q-functions Q and Q The
CbAS optimized GFP
DynaPPO optimized GFP
Guided by oracla trained Guideed oracle trained with negative examples
Guided oracle trained without
Guidedd by aracle trained with negative examples
without
Figure 5: The structure of optimized sequences without and with negative examples.
14

---

Under review as a conference paper at ICLR 2023
two Q functions are due to the clipped double-Q trick, more details can be found in Haarnoja et al. 2018 The loss functions to train Qøi and Qø2 in SAC are defined as
L( D) Be Jr S1.
[Q (St, ax) yr mm))
where the target y(Tes St+1, de). substituting the entropy function, given by
v(r 8c-, mi) = re 7(1 - m²)
min |=1,2
Qetarg, ($1+1. āx+1) - a log No (ro(@++11$++))
(8)
where 0tt1 value function V* is as follows
The policy no is trained to maximize the value function of state The
(*(+)) = Eu,sx (Q ae) - a log (+)]=)))]
The policy is reparameterized and then optimized following Haarnoja et al. (2018). Using a squash- ing Gaussian function, sampling the action is finally defined as
Öeg (5e. (an)((u(se) + (s() C E),
{~,N0,
(10)
A.4 IMPLEMENTATION DETAILS
A.4.1 POLICY
Soft Actor-Critic (SAC) We used three fully connected layers with a hidden size set to 256 for both the actor and critic networks of SAC. The Adam optimizer Kingma & Ba (2014) used for both actor and critic networks. The entropy temperature a is automatically tuned during training following Haarnoja al. (2018).
Proximal Policy Optimization (PPO) For the ablation studies in Table 3 using the mutation in the protein sequence output, we use Proximal Policy Optimization (PPO) Schulman et al. (2017) to handle the discrete action space setting. The PPO architecture is similar to the one used for the actor of SAC, using three fully connected layers with a hidden size set to 256. We define the action as multiple of discrete actions for the third and fifth row in Table
A.4.2 RANDOM PERTURBATION BASELINE
We added random value between [-1,1] to the initial state representation to test the effect of random perturbations on performance. Then, we decoded the sequence using the decoder presented in Sec. 3.1 and predicted the functionality using the functionality predictor presented in Sec. 3.2 Note that the baseline and random performance reported in the Tables and 2 are also calculated on the top K candidates.
A.4.3 PREVIOUS METHODS
For the comparison with previous baseline methods, we use the FLEXS implementation (Sinai et al., 2020) of DynaPPO, CbAS, and BO. All the models are trained for ten rounds. During training, observed the performance saturation after the first few rounds. The sequences proposed in the last round used evaluate the baseline models. For DynaPPO, its ensemble model is trained for ten rounds, and the environment batch size is set to 10. For CbAS, a CNN architecture with a hidden size equal 100 and number of channels equal 32 used. Also, the generator block of CbAS is trained for ten epochs.
A.5 INITIAL STATES OF THE OPTIMIZATION PROCESS
We selected initial states for the optimization process our framework according to two criteria: (i) sequences with room to improve, (ii) sequences sampled from regions that we can trust the oracle. Since the distribution for the ground truth functionality of GFP is bimodal, we sampled two sets of initial states from each mode, low initial states, and high initial states. For His3, only one set of initial states chosen. As shown in Fig. 6 we sample the sequences from pre-determined regions that do not include the most functional mutants. Additionally, the sequences sampled from examples in which the difference between their ground truth and predicted functionality or fitness is less than threshold set to 0.1.
15

---

Under review as a conference paper at ICLR 2023
High initial states
3.5
Initial stated
Peedlted Inness
2.0
Low intial
States
Prodicder
1.5
8.4
0.2
0.
0.0
1,0
1.5
2.5
4.0
Ground truth Tuorescinca
Gronnd truth finees
Figure 6: Ground truth and predicted fluorescence of GFP and fitness of His3.
A.6 EVALUATION ON GFP AND HIS3 DESIGN
In this section, along with the results presented in Tables and we report results on GFP and His3 on Novelty and (TT with standard deviation.
Model
Novelty
dist(WT)
Ours
8,451 # 2.05
7.700 + 0.78
Directed evolution
7.704 2.66
6.849 # 1.90
CbAS
7.712 2.05
6,900 = 0.83
Random-1 1
6.611 = 1.02
7.186 + 2.03
Random-5
13.91 1.22
9.950 + 0.87
Random-P
+
14.15 + 5.76
BO
36.96
36.70 + 5.31
DynaPPO
218.9 2.63
219.3 + 2.37
GFlowNet
199.4 + 2.00
200. + 1.87
Table 8: Results obtained for GFP sequence design. Random-1 I and Random-5 indicates random mutations in I and positions, respectively. Random-P indicates a random perturbation the latent representation.
Model
Novelty
dist(WTT)
Ours
8.361 + 2.01
10.95 + 1.32
CbAS
7.287 # 1.57
4.700 + 0.64
Random-1
7 1.56
7.350 + 1.39
Random-5
9.777 = 1.63
8.950 + 1.36
Directed evolution
6.889 + 1.57
6.710 + 1.57
DynaPPO
27.41 + 1.12
26.70 + 1.19
BO
26.17 1.03
27.50 + 0.50
Table 9: Results obtained for His3 sequence design. Random-I and Random-5 indicates random mutations in I and 5 positions, respectively.
A.7 COMPARISON OF EVALUATION AND OPTIMIZATION ORACLE
In Tables 10 and we compare results on GFP and His3 design using evaluation oracle and op-
timization oracle. It is shown that, for all methods, decrease in performance using the evaluation oracle observed. This decline was particularly pronounced for BO (Swersky et al. 2020). It is also observed that the standard deviation of the results increases when using the evaluation oracle.
16

---

Under review as a conference paper at ICLR 2023
Model
Evaluation oracle
Optimization oracle
Ours
3,491 = 0.352
3,531 0.06
Directed evolution
3,287 = 0.237
3.370 = 0.013
CbAS [Brookes et al. 2019
3.155 = 0.153
3.328 = 0.044
Random mutation (N=1)
2.824 = 0.100
3,410 = 0.094
Random mutations (N=5)
2.280 = 0.275
2.354 # 0.522
Random perturbation
1.511 = 0.797
1,973 = 0.832
BO (Swersky et al. 20200)
0,581 = 0.095
1.231 + 0.034
DynaPPO (Angermuelles al. 2019
0.004 = 0.003
0.014 + 0.001
GelowNet (Jain et al. 2022)
0.000 = 0.002
0.017 = 0.000
Table 10: Comparison of the results obtained for GFP sequence design using evaluation and opti-
mization oracle.
Model
Evaluation oracle
Optimization oracle
Ours
0,945 = 0.091
0.961 = 0.050
CbAS Brookes al. 2019
0.749 = 0.157
0.889 = 0.092
Random mutation (N=1)
0.858 = 0.058
0.856 + 0.070
Directed evolution
0.616 = 0.110
0.756 + 0.013
Random mutations (N=5)
0.678 = 0.096
0.518 + 0.184
DynaPPO (Angermueller et al. 2019
-0.201 0.142
-0.067 0.053
BO (Swersky et al. 2020)
-0.313 0.065
-0.089 0.029
Table 11: Comparison of the results obtained for His3 sequence design using evaluation and opti- mization oracle.
A.8 ABLATION STUDY ON INITIAL STATE AND REWARD MODELING
Table reports the effect of the initial state and reward. Note that, for this ablation study, the performance metric is computed using the optimization oracle. Starting from state of low or high functionality, our framework optimizes performance relative the directed evolution. The best performance is obtained when starting from state of high functionality and training with an absolute reward. Our framework can generate novel sequences even when beginning with a high functionality value. Nonetheless, beginning from state with low functionality, the trained policy does not produce designs with high functionality. This could mean that there may be gap between regions with low functionality and regions with high functionality in the representation space that requires additional time steps to be explored each episode.
Initial state
Reward
Performance*
Novelty
Diversity
Dense
1.856 + 0.631
70%
8.032
Absolute (rth = 1.8) Binary
1,855 + 0.601
50%
5.653
Low
1.546 + 0.271
100%
6.347
Directed evolution
1.492
5.426
Dense
3.448 + 0.094
100%
5,700
Absolute (rm = 3.3)
3.531 + 0.06
100%
6.311
High
Binary
3.452 + 0.081
95%
5.279
Directed evolution
3.370 + 0.013
4.858
Table 12: Ablation studies investigating the sampling of initial states and the reward modeling for the GFP dataset. The policy trained with absolute reward and starting from high initial states used for comparisons in Tables andd Performance(*) is computed using the optimization oracle.
17

---

Under review as a conference paper at ICLR 2023
Number of
Dimension
Figure
Task
Representations
Sequences
reduction
Fig.
(a)
GFP
1.55M
1403
t-SNE
Fig.1
(b)
GFP
1.17M
46
t-SNE
Fig.
(c)
His3
6.43K
6434
I-SNE
Fig.
GFP
1.17M
46
PCA
Table 13: Details of the complete functionality landscape.
Wild-type GFP
Our optimized GFP
DynaPPO optimized GFP
CbAs optimized GFP
Figure 7: Structures predicted by AlphaFold2 for sequences optimized by the proposed method, DynaPPO, and CbAs for GFP.
A.9 METHODOLOGY TO OBTAIN THE COMPLETE FUNCTIONALITY LANDSCAPE
The methodology plot the Fig. 4 and Fig. 8 is detailed next. First, range based the represen- tations obtained during the episode being plotted is defined. Then we decode the distinct sequences from the large number of representations using the decoder presented in Sec. 3.1 Table 13 shows the number of representations and sequences used to plot the complete functionality and fitness landscape. After this step. the reward is decoded for each sequence using the optimization oracle
presented in Sec. 3.2 To reduce the dimensionality, t-SNE (Van der Maaten & Hinton 2008) is
applied in Fig. and Principal Component Analysis (PCA) (FR.S. 1901) is applied Fig g with the two principal components kept to create the local landscape.
A.10
ALPHAFOLDD PREDICTIONS OF OPTIMIZED SEQUENCES
Fig. 7 visualizes the structure of the optimized sequence of GFP based on AlphaFold2 (Jumper et al, 2021p. The sequences produced by method and CbAS maintain the critical chromophore region that is known for emitting fluorescence in GFP and the beta sheets that secure the chromophore re- gion. In contrast, DynaPPO-optimized sequences failed preserve these essential structures of GFP. The different structures obtained by two reinforcement learning methods highlight the significance of state and action modeling the design biological sequences. Learning from a single-hot se- quence encoding makes it difficult for an algorithm to identify crucial structural information. On the other hand, the latent vector trained by extracting information from language model trained with millions of protein sequences can efficiently learn mutational effects and reflect protein structures, as shown in (Rives al. 202T)
18

---

Under review as a conference paper at ICLR 2023
A.11 POLICY ABILITY TO TRAVERSE LOCAL OPTIMA
We qualitatively evaluate our trained policy for the GFP by analyzing its ability to traverse local optima of the functionality landscape. The optimization process during episode is shown in Fig. g Until the end of the episode, at timestep maximize functionality efficiently. At timestep
10, the policy can traverse the landscape and 3, can escape local optima that do not trigger the conditions for the end of an episode. During timesteps t = 7, 8, 9, it can be seen that the policy still maximizing the functionality metric until the end of the episode. Additional details and the methodology used to obtain the functionality landscape in Fig. 8 are explained in A.9
3.5
1
2.5
3.0
0.2
0.0
0.6
0.8
02
2.5
Figure 8: Optimization process performed by the trained policy during one episode for the GFP task. The x-axis and y-axis are the two principal components of the representation space calculated using Principal Component Analysis (FR.S. 1901). The z-axis is the log- flgfeluoreeeenc intensity.
19


Reviews:
[{"review_id": "GgBynxdMQN9", "reviewer": "Reviewer_F4US", "summary": "", "questions": "", "limitations": "", "rating": 8, "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 manuscript presents a procedure for protein engineering using a model-based reinforcement approach building on the ESM2 language model. The authors demonstrate that the high dimensional ESM2 representation can be mapped to a lower dimensional representation space which is suitable for optimization, and that full amino acid sequences can be reconstructed from this reduced representation with meaningful accuracy. They then propose a off-policy reinforcement learning procedure to learn how to make updates in the representation space. Finally, the authors evaluate the method on two protein engineering datasets, and report convincing results.", "summary_of_the_review": "The paper proposes a new method for protein engineering. It is well-written, carefully documents its claims, and demonstrates convincing results. I have only minor suggestions for edits to the paper.", "strength_and_weaknesses": "The paper is well written, the method well described, and the results are convincing. Please see below for detailed comments to specific parts of the paper.", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "FNH7t-KIA", "reviewer": "Reviewer_KJ1Y", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "This paper aims to generate protein sequences with high functionality and cellular fitness. The authors apply a reinforcement learning framework to walk through the space of latent representations of a pre-trained protein language model, using a functionality predictor as a reward function. This approach is tested for its ability to design examples of proteins GFP and His3 with predicted functionality, novelty, and diversity. The paper also tests improving the functionality predictor for use in reinforcement learning by including negative examples in training.\n", "summary_of_the_review": "The paper has an interesting and somewhat novel approach but I am not convinced by the Performance and Novelty evaluation (as it is explained here) and I am uncertain about the AlphaFold result being informative as this would be the result of any single mutation even when known to destabilise the protein. This point about evaluation is non-trivial as papers like these should be setting the bar for future work; train test splits that are commensurate with the claims of the work and avoidance of circularity are really critical.  ", "strength_and_weaknesses": "Strengths:The approach combines existing models in a potentially interesting way : learning from the universe of natural proteins and labeled data from a large dataset of synthetic proteins with negative examples \n\nWeaknesses:  \n I believe the paper has a couple of  fundamental weaknesses as it stands that should be addressed before it could be accepted at ICLR. 1. While the evaluation is conceptually laudable  (performance, novelty, diversity), the metrics do not seem to capture the concepts well. Performance is evaluated irrespective of the distance from training.  Since ‘Novelty” is evaluated as binary- was in the training set or not. (I think - see equation in 3.1), the metric seems to have no measure of distance from training beyond a single mutation. (It is now well appreciated that predicting the effect of single mutations is relatively successful with even baseline methods such as conservation, Potts models or VAEs with alignments and transformers without alignments.) Therefore - for this piece of work to be evaluated I suggest it’s important to show sequence generation as a function of the distance from training data. A fundamental challenge in protein design is being able to generate sequences with a given function that have sequences different from natural or training examples. As one moves away from known sequences (in eg Hamming distance) - the harder it gets. For sequences that are only one  mutation away is relatively easy. ( many papers have shown this).  The performance results shown in Tables 1 indicate that their method is only 1% better than a  random single mutation for eg GFP, Table1, suggesting the metrics and/or the model is poor. Although the authors note this point , they do not follow up by addressing the reasons. \n2. The use of he Oracle twice is circular - therefore invalidates the claims of performance; there are some ways around this that the authors could try \nMore minor weaknesses: \n3. the reference used to justify the evaluation metrics is Hoffman et al 2022 - but this paper is about optimising small molecules - which are v different in \"seq distance to function\" relationships -  this is especially important in relation to the point about the Novelty measure above.  \n4. AlphaFold is not at all appropriate to support the claim of functional sequence optimisation - there are may mutations that will cause a protein to unfold that Alpha fold will predict as having almost exactly the same structure as it will align etc - therefore it proves nothing ( From their own FAQ page \"AlphaFold has not been validated for predicting the effect of mutations. In particular, AlphaFold is not expected to produce an unfolded protein structure given a sequence containing a destabilising point mutation.\" And there are papers writing about this  eg Pak et al 2021 \n", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "CrXXOVbJQ0e", "reviewer": "Reviewer_f1s7", "summary": "", "questions": "", "limitations": "", "rating": 3, "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 proposes a model-based reinforcement learning method for sequence design. The main innovation is that design is performed in the latent space, which should allow easier optimization given the combinatorial search space. Furthermore, the authors point out the flaw of using oracles trained on experimental data as the sole source of truth, as they give high functional scores to random sequences. They demonstrate that adding random sequences to the training set alleviates this issue.", "summary_of_the_review": "The paper proposes a theoretically sound procedure, but many small details are missing that make it difficult to assess the results. Clearer explanations of all aspects of training, model architecture, and experiments are needed. Additionally, it is possible the experimental setting is unfairly advantaging the proposed method by initializing from high-function sequences. Some attempts should be made to provide similar initialization to other methods, or it should be explained why this is not possible.", "strength_and_weaknesses": "Strength:\n* The paper is theoretically sound. The issues identified with current approaches are real, and the proposed solutions are sensible.\n\nWeaknesses:\n* The paper is not carefully written and it was hard for me to find information about architecture, training setup, and experimental results. Examples:\n    - The paper says the CLS token is used for an embedding of a sequence. Further it says average pooling is applied to this vector to reduce dimensionality. I am not familiar with average pooling being used for dimensionality reduction. The method should be explained or a citation provided.\n    - While the paper says CLS tokens are used for dimensionality reduction, this is not clear from Figure 5. \n    - The architecture of the MLP is not described in section A.1.\n    - In section A.1, the paper states that encoder weights are not updated during fine-tuning, but the next sentence states that the encoder and decoder are jointly finetuned.\n    - The decoder is not introduced in section A.1.\n    - \"Note that the information on mutated positions such as masks is not provided to the model.\" - I don't understand this statement. It would be nice to have a clearer explanation of the training procedure.\n    - Details about training set split are not provided.\n    - It was unclear how many total sequences the model was trained on and how many total sequences were seen in a given optimization round.\n    - In section 4.3, it states that \"only the latent vector [...] can simultaneously optimize GFP and His3 tasks.\" I'm not sure what this means -  how are the tasks being simultaneously optimized?\n* The proposed method is initialized from sequences with relatively high functionality. The method proposed is able to optimize further from there. They also show a random mutation baseline that is similarly initialized. However, it does not seem that other methods are provided this initialization. This should be clearly discussed and some effort should be made to make this initialization.\n* It's not clear what the highest functional value in the training set is. Is it higher than the values the model ultimately achieves?\n* Is the encoder/decoder trained only on the training set or on the whole dataset?\n* Figure 1 may benefit from providing a more detailed view of the process (what are the components trained on, how does sampling occur, etc.)\n\nMinor Comments:\n* The CLS token used for a protein-level vector but in ESM-2 it has no special meaning. Would averaging the final layers provide a better representation?\n* Abstract: \"GPF\" -> GFP\n* Page 5: \"CbAsoptimize\" -> \"CbAs optimize\"\n* In tables, I personally dislike the use of \"Baseline\" as an entry. Especially given that there are other methods you compare against, it would be nice to clearly define what this is in the table.\n", "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": "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": "65Kdi2LXqM", "reviewer": "Reviewer_62rv", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "This paper combines ideas from latent space optimization and reinforcement learning to propose a method for model-based optimization for biological sequences. They also introduce a nice idea for using random samples as negative sampling (data augmentation) which seems to significantly drive performance. They then apply their results to GFP and His3 and demonstrate improved performance when tested against a prediction oracle.", "summary_of_the_review": "This paper presents some nice ideas(RL, latent space optimization, negative sampling) for an exciting problem settings. They have some interesting evaluations(qualitative analysis of climbing the landscape, sanity check against AF2). However, the primary weaknesses of this paper are twofold: (1) the exposition surrounding the method is ambiguous(ie seemingly inconsistent language) and the experimental section largely deprioritizes the explaining the core details that give someone working in this space confidence that the method is effective (eg hyper-parameter optimization, how train/test splits are performed, inconsistent results from the literature) and (2) the experiments don't follow best practices in the field that allow the reader to interpret the results at face value. Luckily, I think both of these issues can be rectified and are not detrimental to the idea itself.", "strength_and_weaknesses": "Strengths:\n- There are a nice set of ideas presented in this work. Combining latent space optimization and reinforcement learning is an interesting set of ideas.\n- Negative sampling-based data augmentation is a great idea that is likely a huge driver in improvements of this method as well as can be applied to most other MBO methods\n- Sanity checking with AF2 is a nice \n\nWeaknesses:\n- As stated in Evaluation Metrics section, its problematic to evaluate performance with the same oracle that you use to do property prediction. I suggest following Kolli et al (2022)[1] to use an oracle trained on the full dataset for evaluation and using a smaller dataset for train/test split for functionality predictors.\n- Are the random mutations truly random? If so, a stronger baseline would be to compete against random mutations ranked by an oracle.\n- How were hyperparameters chosen on-line? Such as the action space, f(s_t) threshold, part of representation space that can't be trusted, the number of negative samples, the data distribution of negative samples, etc. Often methods in the Model-based optimization space can't be trusted if they have many hyperparameters without a methodology as to how to tune them off-line.\n- It seems that the performance of some of these baseline methods largely departs from other instances such as in Jain et al (2022) perhaps due to implementation details or differences in value K. Can this be commented on (either via an ablation over K or validating or commenting on the irreproducibility of experiments in prior work)? This would have a huge impact on ranking of these MBO algorithms seems rather poor despite using the same benchmark.\n- Why was the subset of two proteins chosen? We see in other works large variability in method ranking across methods and so they attempt to combat this by testing across 5+ benchmark tasks (occasionally in Notin et al (2022) [2] it has risen up to 87 DMS datasets for prediction tasks).\n- The analysis in Fig 3 would be more convincing if a held-out oracle was used rather than a similar evaluation vs prediction oracle.\n\n\n[1] Kolli et al (2022). \"Data-Driven Optimization for Protein Design: Workflows, Algorithms and Metrics\"\n[2] Notin et al (2022). \"Tranception: protein fitness prediction with autoregressive transformers and inference-time retrieval\"", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject