Improving protein optimization with smoothed fitness landscapes
Authors: Andrew Kirjner
Jason Yim
Raman Samusevich
Shahar Bracha
Tommi S. Jaakkola
Regina Barzilay
Ila R Fiete
Venue: ICLR 2024 poster (2024)
Keywords: protein design
discrete optimization
protein engineering
markov chain monte carlo
graph signal processing
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: State-of-the-art protein optimization method by smoothing the protein fitness landscape.

Abstract: The ability to engineer novel proteins with higher fitness for a desired property would be revolutionary for biotechnology and medicine. Modeling the combinatorially large space of sequences is infeasible; prior methods often constrain optimization to a small mutational radius, but this drastically limits the design space. Instead of heuristics, we propose smoothing the fitness landscape to facilitate protein optimization. First, we formulate protein fitness as a graph signal then use Tikunov regularization to smooth the fitness landscape. We find optimizing in this smoothed landscape leads to improved performance across multiple methods in the GFP and AAV benchmarks. Second, we achieve state-of-the-art results utilizing discrete energy-based models and MCMC in the smoothed landscape. Our method, called Gibbs sampling with Graph-based Smoothing (GGS), demonstrates a unique ability to achieve 2.5 fold fitness improvement (with in-silico evaluation) over its training set. GGS demonstrates potential to optimize proteins in the limited data regime. Code: https://github.com/kirjner/GGS

Published as a conference paper at ICLR 2024
IMPROVING PROTEIN OPTIMIZATION WITH SMOOTHED FITNESS LANDSCAPES
Andrew Kirjner*
Jason Yim
Massachusetts Institute of Technology k kirjner@mit.. eduu
Massachusetts Institute of Technology jyim@mi t. edu
Raman Samusevich
Shahar Bracha
IOCB Prague, Czech Academy of Sciences, CIRCC, Czech Technical University in Prague, University of Chemistry and Technology, Prague
Massachusetts Institute of Technology shaharbr edu
raman. www.sssiiaassueooiiooooooa cas.cz
Tommi Jaakkola
Regina Barzilay
Massachusetts Institute of Technology
Massachusetts Institute of Technology
Commi@csale miriiimmmtttt.. mimtit...... edu
regina@csail .mit. edu
Ila Fiete
Massachusetts Institute of Technology fiet@@itt.ed
ABSTRACT
The ability to engineer novel proteins with higher fitness for a desired property would be revolutionary for biotechnology and medicine. Modeling the combi- natorially large space of sequences infeasible; prior methods often constrain optimization to small mutational radius, but this drastically limits the design space. Instead of hcuristics, we propose smoothing the fitness landscape to facilitate protein optimization. First, formulate protein fitness graph signal then Tikunov regularization to smooth the fitness landscape. We find optimizing in this smoothed landscape leads improved performance across multiple methods in the GFP and AAVV benchmarks. Second, we achieve state-of-the-ae results utilizing discrete energy-based models and MCMC in the smoothed landscape. Our method, called Gibbs sampling with Graph-basee Smoothing (GGS), demonstrates a unique ability achieve 2.5 fold fitness improvement (with in-silico evaluation) over its training set. GGS demonstrates potential to optimize proteins in the limited data regime. Code: https//gitttab.com/innnnerGGC
INTRODUCTION
In protein engineering, fitness can be defined as performance on a desired property or function. Examples of fitness include catalytic activity for enzymes (Anderson et al., 2021) and fluorescence for biomarkers (Remington, 2011). Protein optimization seeks to improve protein fitness by altering the underlying sequences of amino acids. However, the number of possible proteins increases exponentially with sequence length, rendering it infeasible to perform brute-force search to engineer novel functions, which often require multiple mutations from the starting sequence (i.e. at least (Ghafari & Weissman, 2019)). Directed evolution (Amold, 1998) has been successful in improving protein fitness, but it requires substantial labor and time.
"Contributed equally to this work. Authors agreed ordering can be changed for their respective interests. Advised equally this work.

---

Published as a conference paper at ICLR 2024
Step 1: ::raining
Goal
Graph-based Smoothing
Train starting data
T
fo(
data
Starting data
A. Noisy fitness landscape
B. Smoothed fitness landscape
C. Sequence to fitness model
Step 2: Sampling
End
Sample mutation
fol (x) Sequence
fe
gradilents
Start
D. Gibbs With Gradients
E. Mutation trajectory
Figure I: Overview. (A) Protein optimization is challenging due noisy fitness landscape where the starting dataset (unblurred) is fraction of the landscape with the highest fitness sequences hidden (blurred). (B) We develop Graph- based Smoothing (GS) estimate smoothed fitness landscape from the starting data. (C) A model is trained on the smoothed fitness landscape to infer the rest of landscape. (D) Gradients from the model are used Gibbs With Gradients (GWG) where on each step new mutation is proposed. (E) The goal of sampling is for each trajectory to gradually head towards higher fitness.
We aim computationally generate high-fitness proteins by optimizing a learned model of the fitness landscape, but face several challenges. Proteins can be notorious for highly non-smooth fitness landscapes': fitness can change dramatically with single mutations, fitness measurements contain experimental noise, and most protein sequences have zero fitness (Brookes et al., 2022). Furthermore, protein fitness datasets are scarce and difficult to generate due to their high costs (Dallago al., 2021). As result, machine learning (ML) methods are susceptible to predicting false positives and getting stuck local optima (Brookes et al., 2019). The 3D protein structure, available, can provide information in navigating the noisy fitness landscape such as identifying hot spot residues (Zerbe al., 2012), but high quality structures are not available in many cases.
One way to deal with noisy and limited data is to regularize the fitness landscape model?. Our work considers a smoothing regularizer which similar sequences (based on a distance measure) are predicted to have similar predicted fitness. While actual fitness lanscapes not smooth, smoothing can be an important tool the context of optimization, allowing gradient- gradient-based methods to reach higher peaks by avoiding local optima, especially in discrete optimization (Zanella, 2020). A few works have studied properties of protein fitness landscapes (Section 2), but none have directly applied smoothing with a graph framework during optimization.
We propose novel method for applying smoothing to protein sequence and fitness data together with an optimization technique that takes advantage the smoothing. First, we formulate sequences as a graph with fitness values as node attributes and apply Tikunov regularization to smooth the topological signal measured by the graph Laplacian. The smoothed data then fitted with neural network to be used as model for discrete optimization (Figure 1 top). Second, we sample over the energy function for high fitness sequences by using the model's gradients in a Gibbs With Gradients (GWG) procedure (Grathwohl et al., 2021). In GWG, discrete distribution is constructed based on model's gradients where mutations with improved fitness will correlate with higher probability. The process of taking gradients and sampling mutations is performed in iterative fashion where subsequent mutations will guide towards higher fitness (Figure bottom).
'Landscape refers to the sequence to fitness mapping. In the sequel, will "model" when referring the fitness landscape model.
2

---

Published as a conference paper at ICLR 2024
Figure shows overview of the method. We refer to the procedure of smoothing then sampling as Gibbs sampling with Graph-based Smoothing (GGS). To evaluate our method, we introduce a set of tasks using the well studied Green Fluorescent Proteins (GFP) (Sarkisyan et al., 2016) and Adeno-Associated Virus (AAV) (Bryant et al., 2021) proteins. We chose GFP and AAV because of their world importance and availability of large mutational data. We design a of tasks that emulate starting with noisy and limited data and evaluate with a trained model (as done in most prior works). We evaluate GGS and prior works on our proposed benchmarks show that GGS is oftte-of-theaart in GFP and AAV fitness optimization. Our contributions are summarized as follows:
We develop a novel sequence- sequence-based based protein optimization algorithm, GGS, which uses graph-based smoothing to train smoothed fitness model. The model is used as discrete energy function to progressively sample mutations towards higher-fitness sequences with GWG (Section 3).
We develop a set of tasks that measure a method's s ability to extrapolate towards higher fitness. We use publicly available GFP and AAV datasets to emulate difficult optimization scenarios of starting with limited and noisy data (Section 4.1).
Our benchmark shows prior methods fail to extrapolate towards higher fitness. However, we show graph-based smoothing can drastically improve their performance; in one baseline, the fitness jumps from 18% to 39% in GFP and 4% to 44% in AAV after smoothing (Section 4.2).
Our method GGS directly exploits smoothness to achieve state-of-the-an results with times higher fitness in GFP and 2 times higher in AAV compared to the next best method (Section 4.2).
2 RELATED WORK
Protein optimization and design. Approaches can broadly be categorized using sequence, structure or both. Sequence-basee methods have been explored through the lens of reinforcement learning (Angermueller et al., 2020), latent space optimization (Stanton et al., 2022: Lee et al.; Maus et al., 2022), generative models (Notin et al., 2022; Meier et al., 2021; Jain et al., 2022; Gruver et al., 2023), and model-based directed evolution (Sinai et al., 2020; Padmakumar et al., 2023; Ren et al., 2022). Together they face the issue of a noisy fitness landscape to optimize. We focus on sequence-base methods using Gibbs With Gradients (GWG) (Grathwohl al., 2021) which can perform state- of-the-art in discrete optimization but requires smooth energy function for strong performance. Concurrently, Emami et al. (2023) used GWG for protein optimization with product of experts distribution using a protein language model. However, they achieved subpar results.
Previous methods focused developing new sampling and optimization techniques. Our work is complimentary by addressing the need for improved regularization and smoothing. We show in our experiments that our smoothing technique can enhance the performance of prior methods.
Protein fitness regularization. The NK model was an early attempt to model smoothness of protein fitness through a statistical model of epistasis (Kauffman & Weinberger, 1989). Brookes al. (2022) proposed a framework to approximate the sparsity of protein fitness using generalized NK model (Buzas& & Dinitz, 2013). Concurrently, dWJS (Frey al., 2023) is most related to our work by utilizing Gaussian noise to regularize the discrete energy function during Langevin MCMC. dWJS trains by denoising to smooth a energy- based model whereas we apply discrete regularization using graph-based smoothing techniques.
Finally, we distinguish our smoothing method from traditional regularizers applied during training such dropout (Srivastava et al., 2014). Our goal is to smooth the fitness landscape in a way that is amenable for iterative optimization. We enforce similar sequences to have similar fitness which not guaranteed with dropout or similar regularizers applied in minibatch training. Evaluating multiple smoothing strategies is not the focus of our work, but rather to demonstrate their importance.
3
METHOD
The following describes our method. Section 3.1 details the problem formulation. Next section 3.2 describes the procedure for training a smoothed model. Lastly, section 3.3 provides background on Gibbs With Gradients (GWG) which is adapted for protein optimization. The full algorithm, Gibbs sampling with Graph-based Smoothing (GGS), is presented in Algorithm 1.
3

---

Published as a conference paper at ICLR 2024
3.1 PROBLEM FORMULATION
We denote the starting set of N proteins as D = (X,Y) where X x= (21,...,,00 .nn. C yaM are the sequences and Y (y),.... .UN are corresponding real-valued scalar fitness measurements. Each sequence E )Ad is composed of M residues from a vocabulary V of 20 amino acids. Subscripts refer different sequences. Note our method can be extended to other modalities, e.g. nucleic acids.
For in-silico evaluation, we denote the set of all known sequences and fitness measurements as D = (X", ,Y),. We assume there exists a unknown black-box function g V R such that g(z")= = 3"" In practice, 9 needs to be approximated by evaluator model, 0g- trained with weights to minimize prediction errоr on D Ie poses a limitation to evaluation since the true fitmess needs to be verified with biological experiments. Nevertheless, in-silico approximation provides accessible way for evaluation and is done in all prior works. The starting dataset a strict subset of the known dataset D C D to simulate fitness optimization scenarios. Given D, our task is to generate a set of sequences with higher fitness than the starting set.
3.2 GRAPH-BASED SMOOTHING ON PROTEINS
Our goal is develop a model of the sequence-tee to--tt-fffii to-fitness mapping that can be utilized when sampling higher fitness sequences. Unfortunately, the oigh-dimensional sequence space coupled with few data points and noisy labels can result noisy model that is prone to sampling false positives or getting stuck in local optima. To address this, we use smoothing techniques from graph signal processing.
The smoothing process is depicted in Figure First, we train noisy fitness model fø V.M RR
with weights 0 on the initial dataset 'D using Mean-Squared Error (MSE). ' is usually very small in real-world scenarios. We augment the dataset by using fa infer the fitness of neighboring sequences which we do not have labels for por-knon known as transductive inference. Neighboring sequences are generated by randomly applying point mutations to each sequence in X. The augmented and original sequences become nodes, V, in our graph while their fitness labels are node attributes. Edges, E, are constructed with a k-nearest neighbor (kNN) graph around each node based on the Levenshtein distance?. The graph construction algorithm can be found in Algorithm 4.
The following borrows techniques from Isufi et al. (2022). The smoothness the fitness variability in our protein graph is defined as the sum over the square of all local variability,
TV2(Y)
- LEV
E 2
Ay
Σ (NJ)EE
(yi - 33
V
TV refers to Total Variation and Ay, is the local variability of node i that measures local changes in fitness. Using TV2 as a regularizer, we solve the following optimization problem, known as Tikhunov regularization (Zhou & Schölkopf, 2004), for new set of smoothed fitness labels,
arg min ||Y r12 + TV2(Y). PERIVI
With abuse of notation, we represent Y as vector with each node's fitness. 7 is hyperparameter set to control the smoothness; too high can lead underfitting. We experiment with different y's in Section Since eq. (1) is quadratic convex problem, has closed form solution, y (I+>L)- where L is the graph Laplacian and Il is the identity matrix. The final step is to retrain the model on the sequences in the graph and their smoothed fitness labels. The result will be a model So with lower TV2 than before and thus improved smoothness. The smoothing algorithm is in Algorithm 2.
3.3 SAMPLING IMPROVED FITNESS WITH GIBBS
Equipped with model la from section 3.2, we apply it in a procedure to sample mutations that improve the starting sequences' fitness. fo can also be viewed as an energy-based model (EBM) that defines a Boltzmann distribution log p(x) fa(x) - log Z where Z is the normalization constant. Higher fitness sequences will be more likely under this distribution, while sampling will induce diversity and novelty. To sample from p(æ). we use Gibbs With Gradients (GWG) Grathwohl et al. (2021) which
Defined the minimum number of mutations between two sequences.
4

---

Published as a conference paper at ICLR 2024
0.8 AA
0.6
0.4
AB
Train noisy model
Transductive inference
AA
AB
Smoothing
AA
Train smooth model
Sø
fe
0.4
0.2
D
BB
BA
BB
Figure 2: Steps in graph-based smoothing on proteins illustrated with fictitious data of length sequences with vocabulary {A, B}. Above each node are corresponding fitness values. Solid nodes are those in our training set while dashed nodes are augmented via point mutations increase the smoothing effectiveness. See section 3.2 for description of each step.
has attracted significant interest due to its simplicity and state-of- theftt performance in discrete optimization. this section, we describe the GWG procedure for protein sequences. GWG uses Gibbs sampling with approximations of locally informed proposals (Zanella, 2020):
g(2'|x) (z) x exp (xtt) 1(x e H())), da(z); [ufo(=)) o (V
(2)
With slight abuse of notation, we use the one-hot sequence representation {0, 1 where I, {0, 1} VVl represents the ith index of the sequence with at its amino acid index and O elsewhere. {y E VM daamming (æ, y) < 1] is the 1-ball around x using Hamming distance. The core idea of GWG is use dø(z); as the first order approximation of a continuous gradient of the change in likelihood from mutating the ith index of x to different amino acid. The quality of the proposals eq. (2) rely on the smoothness of the energy (Theorem 1 in Grathwohl et al. (2021)). If the gradients, V fo. are noisy, then the proposal distributions are ineffective in sampling better sequences. Hence, smoothing fe is desirable (see section 4).
@ is the element wise product. ((x)
The choice of H(-) as the I-Hamming ball limits x' point mutations from and only requires O (M [/) compute to construct. Let the point mutation where x and x' differ defined by the residue location, glor [f......M' and amino acid substitution, joub {l..... | By limiting to point mutants junb), sampling q<>'(z) (x) is equivalent to sampling the following.
( do dob())
M,IVI
gloe.jouhy ~ a(-1 = Cat Softmax
(3)
where 7 is the sampling temperature and do()).. is the logits of mutating to (i.j). The proposal sequence x' is constructed by setting its gloe residue to jaub and equal to elsewhere. Each proposed sequence is accepted or rejected using Setropolis-Hasttng (MH),
min exp(fe(xx) - fo()
(((1xx) 1)
(xx (2)
provide the GWG algorithm Algorithm 3.
Clustered sampling. GWG requires starting sequence start mutating, A reasonable starting set are the sequences X used to train the model. On each round r, we use eq. (3) to propose Npoo mutations for each sequence. If accepted via eq. (4), then the mutated sequence will be added to the next round. However, this procedure can lead to an intractable number of sequences to consider.
To control compute bandwidth, quences in a round and take the sequence of each cluster with the highest predicted fitness using fa- Let C be the number of clusters which set based on amount of available compute. This procedure, known as Reduce, is,
perform hierarchical clustering (Müllner, 2011) all the
U (arg max fe(z) where { Je=1 = Cluster(X;C).
Reduce
(5)
= #eX=
Each round reduces the sequences from the previous round and performs GWG sampling.
x, Red X cwG(X,;0)
5

---

Published as a conference paper at ICLR 2024
To summarize, adapted GWG for protein optimization by developing smoothed model to satisfy GWG's smoothness assumptions and use clustering during sampling to reduce redundancy and compute. An illustration of clustered sampling is provided in Figure 5.
The full algorithm for smoothing and clustered sampling is provided Algorithm 1.
Algorithm 1 GGS: Gibbs sampling with Graph-based Smoothing
Require: Starting dataset:
D Y) (v {z)))
1: 0 - arg max; (ay,y)>
b Initial training
2: f < Smooth(D;
GS algorithm 2
3: for r = . 1 do
4: X, + Re;
5: X,+1 + ();
>> GWG algorithm
6: end for
7: Return Top)(Xn))
> Return Top-K best sequences based on predicted fitness fo
EXPERIMENTS
Our experiments demonstrate the benefits of smoothing in protein optimization. Section 4.1 presents a set of challenging tasks based on the GFP and AAV proteins that emulate starting with experimental noise and a sparsely sampled fitness landscape. Section 4.2 evaluates the performance of baselines and our method, GGS, on our benchmark. In addition, we find applying smoothing improves performance for two baselines. Section 4.3 provides sweeps over hyperparameters and analysis of GGS.
Baselines. We choose representative set of prior works that evaluated on GFP and AAV: GFlowNets (GFN-AL) (Jain er al., 2022), model-based adaptive sampling (CbAS) (Brookes et al., 2019), greedy search (AdaLead) (Sinai et al., 2020), bayesian optimization (BO-qei) (Wilson et al., 2017), conser- vative model-based optimization (CoMs) (Trabucco et al., 2021), and proximal exploration (PEX) (Ren et al., 2022). NOS (Gruver al., 2023) performs protein optimization with diffusion models. However, their framework is tailored to antibody optimization and requires non-trivial modifications for general proteins. We were unable to evaluate Song & Li (2023) due to unrunnable public code.
GGS implementation. We use a ID CNN (see Appendix B. for architecture and training) for model fo. To ensure fair comparison, we the same model architecture in baselines when possible. In graph-based smoothing (GS), we augment the graph until it has Nnodes 250, 000 nodes. We found larger graphs to not give improvements. Similarly, we use T 0.1, R 15 rounds and Nprop 100 proposals per round during GWG at which sequences would converge and more sampling did not give improvements. We choose the smoothing weight 7 1.0 through grid search. We study sensitivity to hyperparameters, especially y, in Section 4.3.
4.1 BENCHMARK
We develop set of tasks based on two well-studied protein systems: Green Fluoresent Protein (GFP) and Adeno-Associale due to their relatively large amount of measurements, 56,806 and 44,156 respectively, with sequence variability of up to 15 mutations from the wild-type. Other datasets are either too small or do not have enough sequence variability. GFP's fitness its fluorescence properties as biomarker while for AAV's is the ability package DNA payload, i.e. for gene delivery. We found GFP and AAV to suffice in demonstrating how prior methods fail to extrapolate.
Virus (AAV) (Sarkisyan et al., 2016: Bryant et al., 2021). These were chosen
One measure of difficulty is the number of mutations required achieve the highest known fitness; this assesses a method's exploration capability. We designate the set of optimal proteins, X-9hh any sequence in the 99th fitness percentile in the entire dataset Quantitatively, we compute the minimum number of mutations required from the training set to achieve the optimal fitness:
Gap Xu; X 99th) min (inn((istt/i (z, .)
X, 7 E x9th)).
(6)
"This may differ from the true optimal protein found in nature. Unfortunately, must work with existing datascts since every possible protein cannot be experimentally measured.

---

Published as a conference paper at ICLR 2024
A high mutational gap would require the method discovering many mutations in high dimensional space. A second measure of difficulty is the fitness range of the starting set of sequences. Starting with low range of fitness requires the method to learn from barely functional proteins and exploit limited knowledge to find mutations that confer higher fitness. Appendix A shows Gap and starting rate are necessary only requiring one mutation to achieve the optimal fitness.
found the previous GFP benchmark (Trabucco et al., 2022) as too "easy" by
Recall the protein optimization task is use the starting set D propose set of sequences with higher fitness. We design two difficulties, medium and hard, for GFP and AAV based on the properties of D. We restricted the range and the mutational gap modulate task difficulty. We found Gap= and Range 30% to suffice in finding where our baseline methods fail to discover better proteins. We use this setting as the hard difficulty and sought to develop GGS to solve it.
Table 1: GFP tasks
Table 2: AAV tasks
Difficulty
Range (%)
Gap
[D]
Difficulty
Range (%)
Gap
[D]
Medium
20th-40th
6
2828
Medium
20th-40th
6
2139
Hard
< 30th
7
2426
Hard
< 30th
7
3448
In-silico evaluation. We follow prior works in using a trained evaluator model as proxy for real- world experimental validation. A popular model choice is the TAPE transformer (Rao et al., 2019). However, we noticed poor performance of the transformer compared to simpler CNN that matches the findings of Dallago et al. (2021). We use CNN architecture for the evaluator due to its superior performance. Following Jain et al. (2022), each method generates 128 samples X = )2, whose approximated fitness predicted with the evaluator. We additionally report Diversity and Novelty that are also used in Jain et al. (2022). Descriptions of these metrics can be found in Appendix B.2 We emphasize that higher diversity and novelty are not equivalent to better performance, but provide insight into the exploration and exploitation trade-offs of different methods. For instance, a random algorithm would achieve maximum diversity and novelty.
4.2 RESULTS
We run 5 seeds and report the average metric across all seeds including the standard deviation in parentheses. We evaluate GGS and previously described baselines. To ensure fair comparison, the same CNN architecture the model across all methods - all our baselines (and GGS) perform model-based optimization. Since graph-based smoothing (GS) is a general technique, we sought to evaluate its effectiveness in each of our baselines. To incorporate GS, used the smoothed predictor as a replacement in each baseline which will be denoted with "+ GS". Table summarizes GFP results while table 4 summarizes AAV.
GGS substantially outperforms all unsmoothed baselines, consistently achieving a improvement in fitness from the starting range of fitness in each difficulty. However, the smoothed baselines (lines with + GS) demonstrated improvements in GFP where the sequence space is far larger than AAV suggesting the GFP fitness landscape is harder to optimize over.
to three fold improvement for CbAS, AdaLead. We find larger
The most difficult task clearly hard difficulty on GFP where all the baselines without smoothing cannot achieve fitness higher than the training set. With smoothing, GGS achieves the best fitness since the sampling procedure uses gradient-based proposals that benefit from a smooth model. Appendix C presents results additional difficulties to analyze GGS beyond hard..
We observe GGS is able to achieve the highest fitness while exhibiting respectable diversity and novelty. Notably, GGS's novelty falls within the range of the mutational gap in each difficulty, suggesting it is extrapolating an appropriate amount for each task. Our sampling procedure, GWG, fails to perform without smoothing which agrees with its theoretical requirements of requiring smooth model for good performance. We conclude smoothing is a beneficial technique not only for GGS but also for some baselines. GGS is able to achieve tate--of-the-ar results in our benchmark.
7

---

Published as a conference paper at ICLR 2024
Table 3: GFP optimization results. Bold indicates improvement with smoothing.
Medium difficulty
Hard difficulty
Method
Fitness
Diversity
Novelty
Fitness
Diversity
Novelty
GEN-AL
0.09 (0.1)
25.1 (0.5)
213 (2.2)
0.1 (0.2)
23.6 (1.0)
214 (4.2)
GFN-AL GS
0.15 (0.1)
16.3 (1.6)
213 (2.7)
0.16 (0.2)
22.2 (0.8)
215 (4.6)
CbAS
0.14 (0.0)
(1.1)
(0.4)
0.18 (0.0)
(1.3)
7.8 (0.4)
CbAS GS
(0.1)
(..)
(.0)
0.57 (0.0)
4.2
(0.6)
AdaLead
0.56
(..1)
(.00
0.18 (0.0)
(0.)
2.8
AdaLead + GS
0.59 (0.0)
(.3)
(.0
0.39 (0.0)
(.1)
2.00
BOqei
0.20
(0.0)
(.0
(.5)
(71)
(81)
BOqei GS
0.08
19.3 (0.0)
0.0))
(..0)
13.4
(.0)
CoMS
(.1)
133 (25)
(12)
0.1
144
201 (3.0)
CoMS + GS
(0.5)
129 (25)
128 (84)
0.0
114 (36)
187 (5.7)
PEX
0.47 (0.0)
(.0)
1.4 (0.2)
0.0
3.0
1.3 (0.3)
PEX GS
0.45 (0.0)
2.9 (0.0)
1.2 (0.3)
0.0
2.90)
1.2 (0.3)
GWG
0.1 (0.0)
(3.0
12.8 (0.4)
0.0
4.2 47.0)
7.6 (1.1)
GGS (ours)
0.76 (0.0)
3.7 (0.2)
5.0 (0.0)
0.74 (0.0)
3.6 (0.1)
8.0 (0.0)
Table 4: AAV optimization results. Bold indicates improvement with smoothing.
Medium difficulty
Hard difficulty
Method
Fitness
Diversity
Novelty
Fitness
Diversity
Novelty
GFN-AL
0.2 (0.1)
(1.2)
(1.1)
0.1 (0.1)
11.6 (1.4)
19.6 (1.1)
GFN-AL GS
0.18 (0.1)
(1.1)
20.6
0.1 (0.1)
9.5 (2.5)
(1.1)
CbAS
0.43
12.7
(0.4)
0.36
14.4
8.6
CbAS + GS
0.47 (0.1)
(.9)
5.3
0.4
(2.4)
(0.0
AdaLead
0..0)
(0.))
(0.4)
(0.0)
(0.1)
3..
AdaLead + GS
0..0
3.77
2.0
0.44
(.9.0.1)
2.0
BÔqei
(..8 (0.2)
(.0
0.32
17.9(0.)
0.0
BOqei GS
0.34 (2.3
0.0
0.02
(7.2
0.0
COMS
(0.1) (5.1
8.2
0.26
10.7
(2.0
CoMS + GS
(.11
(.6
8.67).7)
(.22
13.2
12.6
PEX
0.4
(..)
(0.))
0.3
(.8
1.3 (0.3)
PEX GS
0.4 (0.0)
(0.)
1.4 (0.2)
0.3 (0.0)
2.8 (0.0)
1.1 (0.2)
GWG
0.43 (0.1)
(6.3)
7.7 (0.8)
0.33 (0.0)
12.4)
12.2 (0.4)
GGS (ours)
0.51 (0.0)
(..2
5,4 (0.5)
0.60 (0.0)
4.5 (0.5)
7.0 (0.0)
4.3
ANALYSIS
We analyze the effect of varying the following hyperparameters: number of nodes Nustes in the protein graph, smoothness weight in eq. (1), and number of sampling rounds R during GWG sampling. For space, we leave the analysis of the sampling temperature 7 in appendix C. Figure presents the results of running GGS with different hyperparameters on the hard difficulty of GFP and AAV. Along the X-axis, we plot the median performance of the sequences during each round of GWG where 0 is initialization and 15 are the sequences and the end of GWG. The Y-axis shows the predicted fitness of the smoothed model in blue while the fitness scored with our is shown in red. Interestingly, find in the majority of cases the smoothed model's predictions are highly correlated
8

---

Published as a conference paper at ICLR 2024
with the evaluator along the sampling trajectory. This is despite the model being trained on 4% of the data with the hard filtering. Appendix C.2.2 shows the prediction eπror where we find smoothing greatly improves in predicting the fitness of unseen sequences despite having higher train emror
Graph size. We find Nooees = 250, 000 nodes to have the best performance over a smaller graph with 100,000 nodes. Larger graphs allow for better approximation of the fitness landscape. However, larger graphs require more compute. A future direction could be to determine optimal graph size with different node augmentations strategies than random mutations.
Smoothing. Too much smoothing 10.0 lead to worse performance in AAV while GFP is not sensitive. This suggests the optimal ?/ is dependent on the particular fitness landscape. Since real proteins landscapes are unknown, the biggest limitation of our method is determining the optimal 7. An important extension of GGS is to theoretically characterize landscapes (Buzas & Dinitz, 2013) and provide guidelines of selecting ?
Sampling convergence. We find set number of rounds are required for GWG sampling to converge when the landscape smooth enough (middle and right column). We find additional rounds are unnecessary; in practice, more rounds can be ran to ensure convergence. Results sweeping the temperature are in Appendix C.I where we see 0.1 clearly performs the best for GFP and AAV.
N nodes 100,000
N nodes 250,000
N nodes 250,000
y 1.0
7 = 1.0
y = 10.0
Oracle
Predicted
0.8
GFP fithess
0.6
0.6
0.6
0.4
0.0
0.0
3
1.0
1.0
AAV ritness
0.6
0.4
0.4
0.4
0.2
round
GWG
round
Figure 3: GGS hyperparameter analysis on GFP and AAV hard difficulty, See Section 4.3.
5 DISCUSSION
We present Gibbs sampling with Graph-based Smoothing (GGS) for protein optimization with smoothed fitness landscape. Our main contribution and insight a novel application of graph signal processing to protein optimization. We show smoothing is not only beneficial to method but also to our baselines. To evaluate, we designed a suite of tasks around two measure of difficulty: number of edits achieve the 99th percentile (mutational gap) and starting range of fitness. All baselines struggled to achieve good performance on our tasks. However, some baselines showed a three fold improvement with smoothing. GGS performed the best by combining Gibbs with gradients

---

Published as a conference paper at ICLR 2024
with smoothed model demonstrating the synergy of maddient-based sampling with smooth discrete energy-based model. Our results highlight the benefits of optimizing over a smooth landscape that may not be reflective of the true fitness landscape. We believe it's important to investigate how regularization can be used to transform protein fitness data to he compatible with modern optimization algorithms. Our goal is to not learn the excess biological noise, but find the signal in the data to discover the best protein. We conclude with limitations.
Evaluation limitations. The results demonstrate strong evidence of using smoothing given its improvement in multiple methods. Despite this, our evaluations follow prior works by utilizing trained model for ovaluation. This can be unreliable compared to testing out sequences with wet-lab validation. Unfortunately, wet-lab validation can be cost and time intensive. The ultimate test would be to use GGS in an active learning or experimental pipeline with wet-lab validation in the loop.
Method limitations. Our method utilizes several hyperparameters such as the graph size and smoothing parameter y. We demonstrated the effects of each hyperparameter in Section 4.3. Given the success smoothing, it is desirable to find systematic ways to determine optimal hyperparameters based on approximation of the underlying fitness landscape. We demonstrated our hyperparameter choices are not specific to either AAV or GFP, but this does not guarantee optimality for new landscapes. We believe the connections between spectral graph theory and protein optimization has more to give in advancing the important problem of protein optimization.
ACKNOWLEDGMENTS
The authors thank Hannes Stärk, Rachel Wu, Nathaniel Bennett, Sean Murphy, Jaedong Hwang. Josef Sivic, and Tomáš Pluskal for helpful discussion and feedback.
JY supported in part by NSF-GREP. JY, RB, and TJ acknowledge support from NSF Expe- ditions grant (award 1918839: Collaborative Research: Understanding the World Through Code), Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, the Abdul Latif Jameel Clinic for Machine Learning in Health, the D'RRA Discovery of Medical Countermeasures Against New and Emerging (DOMANE) threats program, the DARPA Accelerated Molecular Dis- covery program and the Sanofi Computational Antibody Design grant. IF is supported by the Office of Naval Research, the Howard Hughes Medical Institute (HHMI), and NIH RS was partly supported by the European Regional Development Fund under the project IMPACT (reg. no. 1...2.1.01/0.0/0... Czech Republic through the e-INFRA CZ (ID:90254), and the MISTI Global Seed Funds under the MIT-Czech Republic Seed Fund.
MIGHH-HH199055
003/0000468), the Ministry of Education. Youth and Sports of the
REFERENCES
Dave W Anderson, Florian Baier, Gloria Yang, and Nobuhiko Tokuriki. The adaptive landscape of a metallo-enzyme shaped by environment-dependeer epistasis. Nature Communications, 12(1): 3867, 2021.
Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. 2020.
Frances H Arnold. Design by directed evolution. Accounts of chemical research, 31(3):125-131, 1998.
David Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. In International conference on machine learning. pp. 773-782, PMLR, 2019.
David H Brookes, Amirali Aghazadeh, and Jennifer Listgarten. On the sparsity of fitness functions and implications for learning. Proceedings of the National Academy of Sciences, (19())e200944911 18, 2022.
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, (9(6):691-69) 2021.
10

---

Published as a conference paper at ICLR 2024
Jeffrey Buzas and Jeffrey Dinitz. An analysis of nk landscapes: Interaction structure, statistical properties, and expected number of local optima. IEEE Transactions on Evolutionary Computation, 18(6):807-818, 2013.
Christian Dallago, Jody Mou, Kadina E Johnston, Bruce J Wittmann, Nicholas Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K Yang. Flip: Benchmark tasks fitness landscape inference for proteins. bioRxiv, pp. 2021-11, 2021. 2021.
Patrick Emami, Aidan Perreault, Jeffrey Law, David Biagioni, and Peter St John. Plug & play directed evolution proteins with gradient-based discrete memc. Machine Learning: Science and Technology, 4(2):225014, 2023.
Nathan C Frey, Daniel Berenberg, Karina Zadorozhny. Joseph Kleinhenz, Julien Lafrance-Vanasse, Isidro Hotzel, Yan Wu, Stephen Ra, Richard Bonneau, Kyunghyun Cho, al. Protein discovery with discrete walk-jump sampling. arXiv preprint arXiv:2306. 12360, 2023.
Mahan Ghafari and Daniel B. Weissman. The expected time to cross extended fitness plateaus. Theo- retical Population Biology, 129:54-67, 67, 2019. ISSN 0040-5809. doi: https://oii.org/ wtt//....////........ 2019.03.008. URL https: /www. swwencedi reect com/ scisnceee article /etee//piiii 50040580918301010 Special issue in honor of Marcus Feldman's 's 75th birthday.
Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris Maddison. Oops took gradient: Scalable sampling for discrete distributions. In International Conference on Machine Learning, pp. 3831-3841. PMLR, 2021.
Nate Gruver, Samuel Stanton, Nathan C Frey, Tim GJ Rudner, Isidro Hotzel, Julien Lafrance-Vanasse, Arvind Rajpal, Kyunghyun Cho, and Andrew Gordon Wilson. Protein design with guided discrete diffusion. arXiv preprint arXiv:2305. 20009, 2023.
Elvin Isufi, Fernando Gama, David Shuman, and Santiago Segarra. Graph filters for signal processing and machine learning graphs. arXiv preprint arXiv:221 2022.
Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure F. P. Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio. Biological sequence design with GFlowNets. In Ka- malika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 9786-9801. PMLR, 17-23 Jul 2022. URI https: pestps://procee //proceedings..a.c mll. press v/lvvvl222//// jaio2/aaan2aaa. .tmm..
Stuart A Kauffman and Edward D Weinberger. The nk model of rugged fitness landscapes and its application maturation of the immune response. Journal of theoretical biology, 141(2):211-245, 1989.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6900 69800, 2014.
Minji Lee, Luiz Felipe Vecchietti. Hyunkyu Jung, Hyunjoo Ro. Meeyoung Cha, and Ho Min Kim. Protein sequence design in a latent space via model- reinforcement learning.
Natalie Maus, Haydn Jones, Juston Moore, Matt J Kusner, John Bradshaw, and Jacob Gardner. Local latent space bayesian optimization over structured inputs. Advances in Neural Information Processing Systems, 35:34505-34518, 2022.
Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alex Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. Advances in Neural Information Processing Systems, 34:29287-29303, 2021.
Daniel Müllner. Modern hierarchical, agglomerative clustering algorithms. arXiv preprint
arXiv: /00::///092239 109.2378, 2011.
11

---

Published as a conference paper at ICLR 2024
Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena Hurtado, Aidan N Gomez, Debora Marks, and Yarin Gal. Tranception: protein fitness prediction with autoregressive transformers and inference-time retrieval. In International Conference Machine Learning, pp. 16990-17017. PMIR, 2022.
Vishakh Padmakumar, Richard Yuanzhe Pang. He He, and Ankur P Parikh. Extrapolative controlled sequence generation via iterative refinement. arXiv preprimi arXiv:2303.04562, 2023.
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John Canny, Pieter Abbeel, and Yun S Song. Evaluating protein transfer learning with tape. In Advances in Neural Information Processing Systems, 2019.
S James Remington. Green fluorescent protein: perspective. Protein Science, 20(9):1009-1519 2011.
Zhizhou Ren, Jiahan Li, Fan Ding, Yuan Zhou, Jianzhu Ma, and Jian Peng. Proximal exploration for model-guided protein sequence design. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 18520-18536. PMLR, 17-23 Jul 2022. URL https: / http///ooceeeeeeeeee/ee/ee ren:22a html.
www.s.ppes//vvee vv62/
Karen S Sarkisyan, Dmitry A Bolotin. Margarita Meer, Dinara Usmanova, Alexander S Mishin, George V Sharonov, Dmitry N Ivankov, Nina G Bozhanova, Mikhail S Baranov, Onuralp Soylemez, et al. Local fitness landscape of the green fluorescent protein. Nature, 533(7603) 401, 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 preprint arXiv:2010.02 141, 2020.
Zhenqiao Song and Lei Li. Importance weighted inpectationmmaximizzane for protein sequence design. arXiv preprint arXiv:2305.00386 2023.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-199 1958, 2014.
Samuel Stanton, Wesley Maddox, Nate Gruver, Phillip Maffettone, Emily Delaney, Peyton Greenside, and Andrew Gordon Wilson. Accelerating bayesian optimization for biological sequence design with denoising autoencoders. arXiv preprint arXiv:2203. /2742, 2022.
Brandon Trabucco, Aviral Kumar, Xinyang Geng, and Sergey Levine. Conservative objective models for effective offline model-based optimization. Marina Meila and Tong Zhang (eds.), ), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 10358-10368. PMLR, 18-24 Jul 2021. URL https: https:////ppoceeeee h/tp://e/g/occeeigg..... mlr. wtl..press//e / v 1 /tttt//trtrctcttt.a// html.
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench Benchmarks for data-driven offline model-based optimization. CoRR, abs/2202.084500, 2022. URL https: //a rxiv. org/ abs/ 2202 08450.
James T Wilson, Riccardo Moriconi, Frank Hutter, and Mare Peter Deisenroth. The reparameterizzation trick for acquisition functions. arXiv preprint arXiv:1712.00424 2017.
Giacomo Zanella. Informed proposals for local meme in discrete spaces. Journal of the American Statistical Association, 115(530):852 865, 2020.
Brandon S Zerbe, David R Hall, Sandor Vajda, Adrian Whitty, and Dima Kozakov. Relationship between hot spot residues and ligand binding hot spots in protein eeeeeinnroteii interfaces. Journal of chemical information and modeling, 52(8):2236-2444 2244, 2012.
Dengyong Zhou and Bernhard Schölkopf. A regularization framework for learning from graph data. In ICML 2004 Workshop on Statistical Relational Learning and Its Connections to Other Fields (SRI. 2004), pp. 132,137, 2004.
12

---

Published as a conference paper at ICLR 2024
A ADDITIONAL GFP ANALYSIS
Design-bench difficulty. Prior works have used the GFP task introduced by design-bench (DB), suite of model-based reinforcement learning tasks (Trabucco al., 2022), which samples starting set of 5,000 sequences from the 50-60th percentile fitness range. However, we found this task to be easy in the sense only one mutation was required from sequences in the training set achieve the 99th percentile. We quantify this difficulty using the mutational gap described in eq. (6). Our proposed medium and hard difficulties (Section 4.1) require many more mutations reach the top fitness percentile, see Figure 4. Similar issues may be present in other benchmarks.
GFP task difficulty Medium (n=2828)
Easy 56=5
Hard (n=1636)
60th
50th
40th
Fitness
30th
20th
O
6 0
Mutation gap: edit distance from 99th quantile
Figure 4: Easy is taken from design-bench where sequences between the 50-60th percentile are used in training regardless of edit distance sequences in the 99th percentile. Data leakage is present due to multiple measurements that allows the wild-type and other top sequences to be included during training. Medium filters the training dataset have sequences in the 20-40th percentile and be 6 or more mutations away from anything in the top 99th percentile. Hard similarly filters for sequences at most the 30th percentile and 7 or more mutations away.
13

---

Published as a conference paper at ICLR 2024
B ADDITIONAL METHODS
B.I CNN ARCHITECTURE
We utilize ID convolutional neural network (CNN) architecture in model and oracle. The CNN takes a one-hot encoded sequence as input then applies a ID convolution with kernel width followed by max- max-pooling and dense layer to single node that outputs scalar value. It uses 256 channels throughout for total of 157,000 parameters. Despite its simplicity, we find the CNN to outperform Transformers. Indeed, this corroborates the results in Dallago al. (2021) that simple CNN can be effective in low data regimes.
Training performed with batch size 1024, ADAM optimizer (Kingma Ba, 2014) (with B1 0.9, B2 0.999), learning rate 0,0001, and 50 epochs, using single A6000 Nvidia GPU.
B.2 METRICS
We provide mathematical definitions of each metric. Note Io is the evaluator trained to predict the approximate fitness as proxy for experimental validation.
(Normalized) Fitness median((g(2;; 1 where E(ă;
9e (a( (a(())-min( max (Y )-mimny--)) is
the min-max normalized fitness based the lowest and highest known fitness in Y".
Diversity median median((dist(rr.)) x. &. E X, # x)) is the average sample similarity.
Novelty medi an ((n(3; X) (amatt) where n(;; X) min((dist(am) 5) 5 X - x)) is the minimum distance of sample x to any of the starting sequences X.
Algorithm 2 Smooth: Graph traph-based based Smoothing
Require: Sequences: X
Require: Noisy model weights: 0
1: V, E t CCreateGrap(())
D. Construct graph (Algorithm 4).
2: L + GraphLaplacian((. E))
P Compute graph Laplacian.
3: Y <- [f((z])...., Sā (
4: ÿ <- (I+
D. Compute smoothed fitness labels. Train smoothed dataset.
5: 0 < arg maxg W (1.0)) [10- fe(x))
6: Return 0
Algorithm 3 GWG: Gibbs With Gradients
Require: Parent sequences: X
Require: Model weights: 0
1: X'
2: for x E do
3:
fori= i prop do
Number of proposals per sequence.
4:
x' xx (gloe, (glee,anb) joub) q(-12)
5:
Sample index and token eq. (3) > Apply mutation
6:
E/w0e + -Vyub
7:
if acceppt using eq. (4) then X' X'xX'
8:
x' U (z')
9: end if
10: end for
11:
end for
12: Return X"
Return accepted sequences.
14

---

Published as a conference paper at ICLR 2024
Algorithm 4 CreateGraph
Require: Sequences: X
1: V < X
Construct nodes.
2: while [V] S Modes do
3:
I x~ (VV
4: x' +- (=PointMarationec)
Sample point mutation uniformly at random.
5: end while
6: E UUee kNN(2;; V)
Construct edges (Algorithm 5).
7: Return (V,
Algorithm KNN
Require: Current node:
Require: All nodes: V
D(x) U2EV//z di
Levenstein distance between every pair of sequences.
2: X' + Topk(D(x), V)
> Compute K closest sequences to a. > Construct neighborhood around
3: E(x) <- Unee (x,
4: Return E(x)
Top fitness sequences
Starting sequences
Sampled sequences
each cluster
GWG
Reduce
Xy+1
r+1
Cluster
Cluster 3
Cluster
Cluster
Figure Illustration of clustered sampling. V, is the starting set of sequences for sampling in round GWG (Algorithm 3) is generate many sample sequences, V To control computation, we hierarchically cluster all sampled sequences based on Levenshtein distance and take the top fitness sequence in each cluster, using trained fitness prediction model to score each sequence refer to this subroutine as Reduce (eq. (5)). The top sequences, V used for the next round.
C ADDITIONAL RESULTS
C. SAMPLING TEMPERATURE SWEEP
We determine the effect of different tmperatures y when running GGS the hard difficulty for GFP and AAV. All other hyperparameters follow those used in the main results, see Section 4.2. Table shows the results where clearly 7 0.1 performs the best for both AAV and GFP.
15

---

Published as a conference paper at ICLR 2024
Table 5: Temperature sweep.
GFP hard
AAV hard
Temperature (7)
Fitness
Diversity
Novelty
Fitness
Diversity
Novelty
0.01
0.65 (0.0)
5.3 (0.8)
7.4 (0.5)
0.45 (0.0)
(1.1)
9.0 (0.0)
0.1
0.74 (0.0)
3.6 (0.1)
8.0 (0.0)
0.6 (0.0)
4.5 (0.2)
7.0 (0.0)
0.0 (0.1)
28.2 (0.8)
11.4 (0.5)
0.45 (0.0)
11.9 (0.5)
8.0 (0.0)
2.0
0.0 (0.1)
36.1 (1.0)
13.0 (0.0)
0.33 (0.0)
16.7 (0.9)
8.5 (0.5)
C.2 SMOOTHING ANALYSIS
In this section, we provide further analyses into the effect of smoothing on performance of GGS, extrapolation unseen data, and acceptance rate the GWG sampling procedure. Throughout, we use the same parameters 001, = 1,r 15, Nodes 250, 000 as in the main text.
C.2.1 ADDITIONAL. BENCHMARKS
We first define additional benchmarks, one easier, and three harder, for each protein dataset.
Table 6: GFP extra tasks
Table 7: AAV extra tasks
Difficulty
Range (%)
Gap
(D
Difficulty
Range (%)
Gap
[D]
Easy
50th-60th
0
5609
Easy
50th-60th
0
4413
Harderl
< 30th
8
1129
Harderl
30th
13
1157
Harder2
20th
8
792
Harder2
20th
13
920
Harder3
< 10th
8
397
Harder3
10th
476
We note that the "easy" GFP task is equivalent to the design-bench baseline that is sometimes used as a benchmark protein engineering tasks. Due to experimental noise, protein variants are assayed multiple times, and can be assigned multiple fitness values, which means the fitness values of one sequence may occupy a large percentile range. In the case this task, multiple measurements of the wildtype GFP fitness found in the 50th-60th percentile range. Because WT GFP is also "top sequence," this task necessarily has mutational gap of O. Due to this leakage, we develop our own benchmarks in the main text, and extend those to AAV.
C.2.2 How SMOOTHING AFFECTS PERFORMANCE
The following two tables show how smoothed model outperforms its unsmoothed counterpart according to our evaluator across all GFP/AAV benchmarks except AAV Harder2 (see (+)), and GFP Harder3, where the smoothing not sufficient induce successful GWG sampling (see Table 10).
16

---

Published as a conference paper at ICLR 2024
Table 8: Smoothing improves GGS performance on GFP tasks
Difficulty
Smoothed
Median Fitness
Diversity
Novelty
No
0.05
24.83
13.36
Easy
Yes
0.84
5.45
3.51
No
0.51
10.5
15.4
Medium
Yes
0.76
3.7
5.0
No
0.10
23.02
16.8
Hard
Yes
0.74
3.6
8.0
No
0.00
22.86
17.0
Harderl
Yes
0.67
4.45
9.12
No
0.00
22.22
16.5
Harder2
Yes
0.60
5.42
9.82
No
0.00
23.02
16.8
Harder3
Yes
0.00
15.73
21.2
For the GFP task, our model fails (achieves 0 median fitness) when we restrict the data to the 10th percentile and mutation gap 8 for GFP where [D] = 397.
Table 9: Smoothing improves GGS performance AAV tasks
Difficulty
Smoothed
Median Fitness
Diversity
Novelty
No
0.47
2.69
7.81
Easy
Yes
0.49
9.18
7.99
No
0.37
6.60
6.62
Medium
Yes
0.48
4.66
5.59
No
0.33
12.32
13.8
Hard
Yes
0.60
4.5
7.0
No
0.30
0.53
6.00
Harderl
Yes
0.31
13.80
14.679
No
0.28*
4.46
11.93
Harder2
Yes
0.27
15.98
19.41
No
0.25
3.08
5.63
Harder3
Yes
0.38
7.05
9.486
(*):
The unsmoothed model only outperforms its smoothed counterpart when applying GWG the unsmoothed model generates only a few unique sequences nearby to starting set (as evidenced by low novelty for this benchmark)
For AAV, we find the model able to still find signal and achieve 0.384 evaluated fitness despite the data being limited to the 10th percentile and mutation gap of 13 where [D] = 476. It is notable, though, that the performance improvements gained from smoothing are smaller than in the case of GFP. Presumably, this is due to the vastly reduced dimension of the AAV sequence space in comparison to that of GFP, which may result in a neural network to learn smoother landscape without any regularization.
17

---

Published as a conference paper at ICLR 2024
C.2.3 How SMOOTHING AFFECTS EXTRAPOLATION SAMPLING
The following tables show the benefits of smoothing extrapolation to held ground truth experimental data, up to certain difficulty benchmark, as well as how smoothing vastly improves the acceptance rate for the GWG sampling procedure.
Table 10: Smoothing improves extrapolation and GWG sam- pling, up GFP Harder3
Difficulty
Smoothed
Train MAE
Holdout MAE
Acc. Rate
No
0.03
0.99
0.02
Easy
Yes
0.71
0.61
0.99
No
0.10
1.29
0.61
Medium
Yes
0.20
0.88
0.62
No
0.06
1.44
0.01
Hard
Yes
0.15
0.93
0.43
No
0.07
1.39
0.01
Harderl
Yes
0.15
0.94
0.43
No
0.01
1.41
0.01
Harder2
Yes
0.12
0.90
0.59
No
0.01
1.41
0.01
Harder3
Yes
0.01
1.42
0.01
Table Smoothing improves extrapolation up AAV Hard and GWG sampling on all AAV tasks
Difficulty
Smoothed
Train MAE
Holdout MAE
Acc. Rate
No
0.28
2.82
0.01
Easy
Yes
1.76
2.28
0.99
No
0.35
3.12
0.01
Medium
Yes
0.44
2.76
0.82
No
0.48
3.70
0.30
Hard Yes
0.55
3.09
0.78
No
0.66
3.99
0.01
Harderl
Yes
0.69
4.24
0.47
No
0.56
4.13
0.01
Harder2
Yes
0.58
4.37
0.55
No
0.47
4.58
0.01
Harder3
Yes
0.47
4.59
0.64
For each benchmark category, we evaluated the impact of smoothing on extrapolation abilities by analyzing the Mean Absolute Error (MAE) of the models on that benchmark's training and holdout datasets from the experimental ground truth. The effectiveness of smoothing was indicated by reduced MAE values on the holdout set. We also find that the MAE on the training set is lower for the unsmoothed models, as expected. In line with the results of the previous section, the effect of smoothing is reduced for AAV. As task difficulty increases, for both proteins, the effectiveness of smoothing on extrapolation decreases, which we expect as any signal leading from the training set to the fitter sequences gets obscured training size decreases.
18

---

Published as a conference paper at ICLR 2024
Finally, we note that in every case except two, smoothing dramatically increases accep- tance rate of the GWG sampling procedure, which aligns with the inversely proportional relationship between smoothness of the energy function and sampling efficiency. In the case the hardest GFP task, even the the smoothed model had overfit to the training set. As for the GFP medium task, suspect that this particular section of the experimental dataset allowed the unsmoothed model to learn smooth landscape initially.
19


Reviews:
[{"review_id": "EvNCz4aoc4", "reviewer": "Reviewer_PzPY", "summary": "The paper propose a smoothing method on fitness function given a protein sequence. Assume that the given original data set is small, authors proposed a sampling augmentation method and a TV smoothing regulariser. After which MCMC algorithm is use to further optimise the fit.", "questions": "Is there a way to test that the surrogate function by itself is good enough? The authors look at the overall performance that could infer to correctness of the surrogate function.\n\n\nsee above section on 'weakness'", "limitations": "", "rating": 3, "confidence": 3, "soundness": 1, "presentation": 1, "contribution": 2, "strengths": "Authors presented some good results on benchmark datasets.", "weaknesses": "The paper is hard to read and understand. I itemise areas for improvements.\n\n1. Having one figure to show overall flow of logic could help. Fig1 seems to do the job. There are some confusion between training and sampling. I understand that the author first train f(x) and then use f(x) as a surrogate function for MCMC optimisation. This point does not come out naturally.\n\n2. construction of KNN graph could be described more clearly. (see Eq above Eq(1))\n\n3. Symbols of Eq.(2) are ill defined. The authors should provide in the appendix some details of GWG and reference the appendix in the main text.\n\n4. Eq.(4) should give the acceptance rate. while q are the probability of trial moves. x and x' are two states for jumping in this one MC step. Usual notation is q(x|x') vs q(x'|x), notation of Eq.(4) certainly is not of this form. Instead i^loc and j^sub and being used. The same i^loc and j^sub cannot appear in both numerator and denominator of Eq.(4).\n\n5. Eq.(4) what is the temperature of this move? It seems the temperature is set to 1. Why is the temperature 1? Is there any annealing process?\n\n6. Clustered sampling section should be explained better.", "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": "", "extra_scores": {}}, {"review_id": "CnvJzdmAeS", "reviewer": "Reviewer_Ayis", "summary": "This study proposes to smooth the protein fitness landscape to facilitate protein fitness optimization using gradient based techniques. This is motivated by the ruggedness of protein fitness landscape which makes optimization challenging. A graph based smoothing technique for fitness landscape followed by Gibbs with Gradient sampling is used to perform protein fitness optimization. Evaluation of their method has been done on train sets designed from GFP and AAV with two degrees of difficulty defined by the mutational gap between the starting set and the optimum in the dataset (not included in the starting set).  Their method shows better performance than others in the proposed benchmark. The proposed graph smoothing technique has been shown to help with other methods as well.", "questions": "1)\tPlease explain why after smoothing, the diversity and novelty of the final set of sequences decreases.\n2)\tIn defining train sets with varying levels of difficulty only two medium (mutation gap 6) and hard (mutation gap 7) levels have been used. What happens if you make this harder (higher than 7)? Also, should we assume that for less mutational gap all methods perform comparably?\n3)\tAs stated in the paper, single mutations can dramatically change the fitness. In the smoothing performed, similar sequences are enforced to have similar fitness. Have you investigated where smoothing can be detrimental?\n4)\tHow is the number of proposals ($N_{\\text{prop}}$) per sequence set?\n5)\tHave you tried smaller sizes for the starting set? In real world problems the size of the starting set could be much smaller than 2000? \n6)\tWas the oracle only used at the end for performance evaluation? In AdaLead, did you use the oracle as the fitness landscape or $f_\\theta$?\n7)\tMention the augmented graph size (how does it change with the size of the sequence)\n8)\tMinor: In Eq 5, $X_0$ should be X.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Designing train sets with varying difficulties for the task of optimization.\nProposing a new method for smoothing the protein fitness landscape before optimization.", "weaknesses": "The proposed method has many hyperparameters to tune. \nGiven certain properties of protein fitness landscape, smoothing can hurt if not done properly.", "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": "", "extra_scores": {}}, {"review_id": "gcrzF0MrTh", "reviewer": "Reviewer_7UDm", "summary": "The authors introduces a method called Gibbs sampling with Graph-based Smoothing (GGS) that uses Tikunov regularization and graph signals to smooth the protein fitness landscape, improving the ability to create diverse, functional sequences.", "questions": "For the smoothing procedure, it’d be great to show the amount of error introduced into the labels of the sequences. For instances where either a reasonable oracle model exists, or sequences with large hamming distances have been measured, and this smoothing procedure is introduced, what is the correlation of function values before and after?\n\n“To control compute bandwidth, we perform hierarchical clustering (Mullner, 2011) on all the se- ¨ quences in a round and take the sequence of each cluster with the highest predicted fitness using fθ.” Why not use the “noisy model” for this, because it is the oracle for the true fitness of a sequence?\n\n“Section 4.1 presents a set of challenging tasks based on the GFP and AAV proteins that emulate starting optimization with a noisy and limited set of proteins.” I would like the authors to be clear by what the mean by “noisy”. Is it experimental noise? Is the landscape too sparsely sampled? Where is this noise coming from, and what relative distribution does it have?\n\nGenerally, I feel like Figure 5 is a distraction from the broader utility of the work. I’d just cite Dallago 2021 like you did for the use of CNNs.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Figure 1 is very helpful in the understanding of this approach.\n\nMy understanding of the section described in Section 3.2 is relatively clear.\n\nI think the Fitness, Diversity, and Novelty scores to be interpretable and helpful.\n\nI think it is encouraging that graph-based smoothing (GS) helps almost all other methods in Table 3. It’s also great that this is a relatively straightforward procedure.", "weaknesses": "“While dWJS is an alternative approach to fitness regularization, it was only demonstrated for antibody optimization. To the best of our knowledge, we are the first to apply discrete regularization using graph-based smoothing techniques for general protein optimization.” - This doesn’t seem justifiably novel. Proteins are proteins.\n\nGenerally, I wouldn’t use the term “fitness” when describing protein function. Rather, I would use phenotype or function, as fitness is a broad, poorly defined subset of fitness.\n\nFigure 5 is a reason why these function predictors should not be called “oracles”, because mapping the effect of mutation to function is difficult itself. I’d prefer “protein function approximator”, or something along those lines.\n\n“These were chosen due to their long lengths, 237 and 28 residues” What do you mean here? 28 isn’t that long. I realize it is in the context of a larger protein, but I’d be clear about that.", "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": "", "extra_scores": {}}, {"review_id": "04hqX2C9yJ", "reviewer": "Reviewer_zgX5", "summary": "The manuscript proposes a set of techniques for protein optimization. The first is a method for smoothing protein fitness landscapes. The second is a technique to optimizing in this landscape using the Gibbs With Gradients procedure, which has previously been shown to provide excellent results for discrete optimization. The authors also design two new optimization tasks based on the GFP and AAV datasets, which are designed to be more difficult than previous variants. Finally, the authors demonstrate empirically that their method performs competitively with the state-of-the-art.", "questions": "Page 4,\n*\"Edges, E, are constructed with a k-nearest neighbor graph around each node based on the Levenshtein distance 3.\"*\nIn real-world cases, the starting point is often a saturation mutagenesis experiment, where a lot of candidates will be generated with the same edit distance from the wild type (e.g. edit distance 1). In such cases, won’t the fixed k-out degree lead to an arbitrary graph structure (I mean, if the actual number of equidistant neighbors is much larger than k)?\n\nPage 6, *\"4.1 Benchmark\"*\nIt was difficult to follow exactly what \"develop a set of tasks\" implies. Since the benchmarks are built from existing datasets, the authors should make it clearer exactly what they are \"developing\": is it only the starting set, or do they also restrict themselves to a subsample of the entire set? In table 1 and 2, are both *Range*, *|D|*, and *Gap* specifically selected for, or does e.g. *|D|* arise out of a constraint on *Range* and *Gap*?\n\nPage 6. *\"Oracle\"*\nSince you are using a CNN both as your oracle, and as the basis for your smoothed landscape model, isn’t there a risk that your model is just particularly well suited for exploiting the extrapolation artifacts of the oracle? (repetition of concern stated above).\n\n### Minor comments:\nPage 1, *\"but high quality structures are not available in many cases\"*\nAfter AlphaFold, many would consider that high quality structures are now available in most cases.\n\nPage 2, *\"mutation is proposed renewed gradient computations.\"*\nSomething is wrong in this sentence\n\nPage 3, *\"in-silico oracles provides a accessible way for evaluation and is done in all prior works.\"*\nThis is not entirely accurate. People have optimized against actual experiments (e.g. Gruver, ..., Gordon-Wilson, 2023) - or optimized to find the optimal candidate in a fixed set of experimentally characterized proteins.\n\nPage 4, eq (2) *\"H(x)\"*\nAs far as I can see, H(x) has not been introduced(?)\n\nPage 6. *\"we utilize a simpler CNN that achieves superior performance in terms of Spearman correlation and fewer false positives.\"*\nWas this correlation measured on the GFP test set provided by TAPE after fitting on the training set?. If so, it's odd that the original TAPE paper did not find the CNN-based ResNet to outperform the transformer (actually, the transformer performance was dramatically higher). Please clarify.\n\nPage 6. *\"Recall the protein optimization task is to use D\"*\nPerhaps help the reader by rephrasing to \"Recall the protein optimization task is to use the starting set D\"", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "### Originality\nAlthough the GWG optimization procedure has been used in other contexts, the application to protein optimization is novel. To my knowledge, also the specific graph-based formulation of the regression problem itself is new.\n\n### Quality\nThe paper seems technically sound. Code was provided to ensure reproducibility, and the authors provide additional details about the method in the supporting material. \n\n### Significance\nThe paper does not give much insight into why the method outperform earlier approaches (see below for details), but the empirical results are convincing, which by itself could be sufficient to have impact on the growing subcommunity in ICML interested in protein modelling and design.", "weaknesses": "My main concern with the paper is that I - after reading it - do not feel much wiser about promising methodogical directions for protein modelling going forward.\nWhat I lack in the paper is perhaps more of a motivation of why particular modelling choices were made. For instance, why is the Tikhunov regularization a meaningful choice in the context of protein optimization? Intuitively to me, it seems like a fairly crude choice, ignoring much of what we know about proteins already (e.g. that certain amino acids are biochemically similar to others). The paper also provides no biological intuition about why we would expect the smoothness would help. Presumably, the idea must be that there are different length scales to the problem, and that we can ignore the short length scales and focus on the longer ones - but it is not obvious to be why that would be the case for proteins. Is part of the explanation that experimental data is typically quite noisy? But if that's the case, you would assume that you would get similar behavior by using a simple GP with observation noise - just using a kernel based on edit distance - or based on Eucledian distance in one-hot space. The paper would be much more satisfying for me if the smoothing procedure was motivated more clearly, and perhaps even validated independent of the optimization procedure (I assume you would hope that the smoothed regressor would extrapolate better?)\n\nMy other serious concern is about the empirical evaluation of the model. As far as I can see, when we evaluate an optimization model against an oracle, there is a risk that we end up optimizing against extrapolation artefacts of the oracle, in particular if we end up evaluating it far away from the data it was trained on. My concern is whether your method has an unfair advantage compared to the baselines, because it uses the same CNN architecture for both the model and the oracle - and could therefore be particularly well suited for exploiting these artefacts. To rule out this concern, it would be interesting to see how the model performs against an oracle trained using a completely different model architecture.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)