Improving protein optimization with smoothed fitness landscapes

We thank the reviewer for their time and detailed feedback. We appreciate the valid concerns and hope our response can address them. We have grouped together related points and itemized our response accordingly.

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. What I lack in the paper is perhaps more of a motivation of why particular modelling choices were made…The paper also provides no biological intuition about why we would expect the smoothness would help.

The core takeaway of our work is to demonstrate the benefits of regularization through smoothing. We argue our modeling choices are based on protein knowledge. As we say in our introduction, “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.” It is well-known that deep learning models perform poorly with noisy and small datasets, i.e. curse of dimensionality. Knowing this, it becomes crucial to remove the undesirable properties of protein data and apply regularization if we want good performance with deep learning. For the purpose of protein engineering, we care about finding the best sequences and not exactly modeling the true (un-smooth) fitness landscape.

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 purpose of Tikhunov regularization is to train a surrogate model of the fitness landscape that is synergistic with DL-based sampling algorithms. Our algorithm of choice is Gibbs With Gradients (GWG) which performs exceptionally well on (non-protein) discrete optimization benchmarks, but has not found success in protein sequence optimization. The theory of GWG explicitly states the convergence of the Markov chain and its performance are related to the smoothness of the model. We applied Tikunhov regularization to smooth the model and validated that GWG can achieve state-of-the-art performance if the correct regularization is applied.

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?

Does short length scales refer to short mutational gaps (i.e. requiring few mutations to improve over the wild type)? If so, many protein engineering problems can require the need for long length scales. For instance, Somermeyer et al [1] shows multiple GFP proteins found in nature with cgreGFP having the highest fluorescence but only 41% sequence similarity with the avGFP used in our work. A priori it can be unknown how many mutations are needed to improve over the starting sequences. Furthermore, even for short length scales, the data can be limited and noisy which is where smoothing can still help – we see this in our results where most fail to get non-zero fitness using a un-smoothed model.

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 Euclidean distance in one-hot space.

One of our baselines, BOqei, is a GP. The GP was implemented as part of a baseline in CoMs [2] and included in follow-up works (e.g. [3]) where it was found to perform poorly especially for long proteins like GFP where it achieved the worst performance. GPs in general suffer from scaling to high-dimensional problems such as protein sequences [4].

The paper would be much more satisfying for me if the smoothing procedure was motivated more clearly,

Thank you for the direct feedback! We hope our above answers have motivated smoothing. The introduction (third paragraph) mentions our motivation and Figure 1 provides geometric intuition of why smoothing helps with optimization. We have added a paragraph to the discussion about our modeling choices that depart from known fitness landscapes.

Page 2, "mutation is proposed renewed gradient computations." Something is wrong in this sentence

Thank you for this catch! We have corrected it.

Page 3, "in-silico oracles provides a accessible way for evaluation and is done in all prior works." This 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.

Indeed, Gruver et al does perform in-vitro evaluation for antibodies. We have corrected the statement, “most prior works.” Note Gruver et al makes use of in-silico screening to first filter their designs before running in-vitro experiments.

Page 4, eq (2) "H(x)" As far as I can see, H(x) has not been introduced(?)

We mention H(x) is the 1-Hamming ball in a foot note. We have now included a mathematical description in the text, "$H(x) = \{y \in V^M : d_(x, y) \leq 1 \}$ is the 1-ball around $x$ using Hamming distance."

Page 6. "we utilize a simpler CNN that achieves superior performance in terms of Spearman correlation and fewer false positives." Was 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.

Yes, we evaluated the measurement on the held out set provided by TAPE. Note TAPE used a ResNet which has far more parameters and is a more complicated model than the simple 1D CNN used in our work and in [6] where they observed a similar finding of the CNN outperforming TAPE.

Page 6. "Recall the protein optimization task is to use D" Perhaps help the reader by rephrasing to "Recall the protein optimization task is to use the starting set D"

Thank you for this improvement. We have applied the suggestion.

[1] https://elifesciences.org/articles/75842.pdf

[2] https://arxiv.org/pdf/2202.08450.pdf

[3] https://arxiv.org/abs/2203.04115

[4] https://arxiv.org/abs/1810.12283

[5] https://arxiv.org/abs/2203.12742

[6] https://www.biorxiv.org/content/10.1101/2021.11.09.467890v1

[7] https://www.nature.com/articles/s41586-021-03828-1

[8] https://www.ncbi.nlm.nih.gov/pmc/articles/PMC10019719

and perhaps even validated independent of the optimization procedure (I assume you would hope that the smoothed regressor would extrapolate better?)

We have included regression results that validate with smoothing in table 10 and 11. The results can be found in our global response. We took our unsmoothed and smoothed models trained on the hard and medium splits, and evaluated on the held-out data for which we have the true experimental fitness values for. We see that, as expected, smoothed models have higher training MAE but lower test MAE than their unsmoothed counterparts across all benchmarks, indicative of improved extrapolation as a result of smoothing.

My 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.

This is a valid concern. However, we point out the prior protein engineering publications at ICLR, NeurIPS, and ICML have all used in-silico models to evaluate sequences. We have followed this practice. Some works have included experimental validation, but these authors have resources from biological collaborators or pharma companies [5].

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. Page 6. "Oracle" Since 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).

We state in our work, “to ensure a fair comparison, we use the same model architecture in all baselines.” We have ruled out the possibility of our method solely exploiting the model architecture.

Questions

Page 4, "Edges, E, are constructed with a k-nearest neighbor graph around each node based on the Levenshtein distance 3." In 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)?

Great point! A crucial component of our method is to augment the starting sequences by randomly mutating sequences to extrapolate the graph further than an edit distance of 1. With site saturation data, the starting sequences may have an arbitrary graph structure, but the more important sequences are the augmented ones that the model has not seen during training. These augmented sequences are >1 edit distance from the wild type and provide an approximation of how the model extrapolates. These predictions are smoothed if they are noise and we believe this procedure leads to improved sampling of better sequences.

Page 6, "4.1 Benchmark" It 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?

We apologize for the confusion. Our benchmark trains the oracle on the entire dataset while the difficulties (which we call tasks) limit the starting set each method (including ours) has access to. |D| arises out of the choices for range and gap. Our choices for range and gap were chosen after experimenting with the values where we saw all the baselines failed on the hard starting set. We have changed the formatting of Table 1, placing the |D| column on the right, rather than the middle

Page 1, "but high quality structures are not available in many cases" After AlphaFold, many would consider that high quality structures are now available in most cases.

While AlphaFold2 has made a revolution in the availability of structures, high quality structures are not readily available. The AlphaFold2 database paper [7] states only 58% of predicted residues have a confident prediction of which 36% have very high confidence. Furthermore, prediction of complexes lags behind the prediction of single chain. Lastly, an important point is AlphaFold2 has poor accuracy of point mutations [8].

We thank the reviewer for their time and feedback. Below we address the issues and comments.

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.”

We apologize for the poor choice in wording. dWJS is a concurrent work that performs smoothing by training their energy-based model with denoising. We meant to highlight the differences in smoothing approaches where we instead use graph-based smoothing in protein optimization. We have replaced the line with, “dWJS trains by denoising to smooth a energy-based model whereas we apply discrete regularization using graph-based smoothing techniques.”

Generally, 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.

We internally have debated over the right terminology. We noticed the protein engineering field has settled on using fitness landscape rather than function landscape [1,2,3]. We chose to stick with fitness to follow this convention. Our very first sentence mentions “fitness can be defined as performance on a desired property or function.” We would prefer to keep the wording to avoid major changes, but are happy to change all occurrences if the reviewer views this as a major issue.

Figure 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.

Thank you for the suggestion! We have changed the mention of oracle everywhere to “trained evaluator model.” We agree the term oracle places too much trust in the in-silico evaluations.

“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.

Thank you for pointing this out. We have changed this sentence to highlight the data availability, “these were chosen 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.”

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?

We have included error analysis of the train and test splits for predicting the true fitness. See the global response. We see there is an improvement in generalization error to the held out test set with smoothing. The effect is greater with GFP than AAV.

“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?

We do not use the noisy model, denoted $f_{\tilde{\theta}}$, because the predictions from the noisy model are inaccurate (cf. above table). We use the smoothed model $f_{\tilde{\theta}}$ instead for higher quality predictions. We find using the noisy model for our method leads to very poor performance (2nd to last row in tables 3 and 4).

“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?

This is a great point! By noisy, we mean the experimental noise. We briefly mention this in the introduction, “fitness measurements contain experimental noise.” We have clarified the sentence to say, “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.”

Generally, 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.

Thank you for the suggestion. We have cited Dallago and removed Figure 5.

[1] https://www.nature.com/articles/nrm2805

[2] https://www.biorxiv.org/content/10.1101/2021.11.09.467890v1

[3] https://www.nature.com/articles/s41592-019-0496-6

We thank the reviewer for their time and feedback. Below we address the issues and comments.

Weaknesses

The proposed method has many hyperparameters to tune. Given certain properties of protein fitness landscape, smoothing can hurt if not done properly.

We require four method-specific hyperparameters that we analyze in Section 4.3 and Appendix C.1 as well as providing sweeps for each of them. (We don’t consider CNN model and training hyperparameters since we use standard choices from prior works.) Please note that other works require equal or more method-specific hyperparameters than ours: COMs requires six, DynaPPO requires 4, GFlowNets require 4 hyperparameters. Our work is forthcoming about the hyperparameters and provides analysis into each of them:

1. Sampling temperature $\tau$. Appendix C.1 provides a table of sweeping over different temperatures. Temperature is not specific to us since every MCMC method needs to tune the temperature for the respective need of balancing the sample quality and diversity.

2. GWG round $R$. Again the number of rounds is not specific to our method but is required for GWG. Figure 3 shows $R$ needs to be high enough for the Markov chain to mix and reach its stationary distribution.

3. Number of nodes in the protein graph $N_{nodes}$ is a hyperparameter specific to our method. Figure 3 shows higher $N_{nodes}$ to give better performance, but higher values lead to slower run time. Our choice of $N_{nodes}=250,000$ results in our method taking < 1 hour to sample and still achieves state-of-the-art result.

4. Smoothness weight $\gamma$ is the other hyperparameter specific to our method. Figure 3 shows the choice $\gamma$ is important for the final performance as too high can lead to lower performance. We explicitly state in Section 5 that determining $\gamma$ is a current limitation of our method and automatic ways of choosing $\gamma$ is a important future direction. Choosing $\gamma$ is non-trivial since protein landscapes are very difficult. Nevertheless, our work demonstrates the importance of careful smoothing is in protein optimization.

Questions

1. Please explain why after smoothing, the diversity and novelty of the final set of sequences decreases.

Our metrics reflect the tradeoff between exploitation and exploration where a method needs to maximize the reward (fitness in our case) while discovering the multiple modes of the distribution. Reward and diversity are often at odds since many diverse sequences can be sampled but each have low reward (and vice versa). Tables 3 and 4 show this general trend in methods without smoothing where higher diversity and novelty leads to the worse fitness (see GFN-AL and BOqei). Smoothing will reduce diversity by flattening the local and noisy variation in the fitness landscape (see the unsmoothed and smoothed landscapes in Figure 1). Our results show smoothing is effective to increase fitness across methods while avoiding low quality sequences that could artificially increase diversity.

We thank the reviewer for their time and feedback. We notice the reviewer gave low scores for Soundedness, Presentation, and Contribution. However, we are unable to find reasoning behind the poor Soundedness and Contribution. We kindly ask if the reviewer can elaborate on their low scores to help improve our work and address them in the discussion period. We believe the clarity has improved with the reviewer’s insightful comments. See our responses below.

Weaknesses and questions

The paper is hard to read and understand. I itemise areas for improvements.

We thank the reviewer for pointing out issues regarding clarity.

1. 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.

Indeed, Figure 1 is meant to convey the overall flow of logic. By using the same notation, $f_\theta$, in the training (top row) and sampling (bottom row), we hoped it explains training and inference with $f_\theta$. We have updated figure 1 to clearly indicate that sampling follows after training and the same model $f_\theta$ is used in both training and sampling.

2. construction of KNN graph could be described more clearly. (see Eq above Eq(1))

Thank you for the suggestion. We have added a line near Eq (1) to state: “The graph construction algorithm can be found in Algorithm 4.” Our graph construction pseudocode used to be in Algorithm 2 but we have placed it in its own Algorithm now for clarity. We are unsure what the reviewer meant by “(see Eq above Eq(1))” since this equation refers to the definition of Total variation.

3. 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.

We apologize for confusion and typos. We have clarified in the text that we use the one-hot representation, have rewritten Eq (2) to explicitly describe the matrix algebra, and made sure every symbol is described. Please see our updated writing in Section 3.3. We believe section 3.3 fully describes the GWG algorithm as introduced in Grathwohl et al. Note Eq (4), (5), and (6) in Grathwohl et al are given as Eq (2), (3) in our work while Algorithm 3 in the appendix describes the full GWG algorithm (Algorithm 1 in Grathwohl et al). Our original submission states, “We provide the GWG algorithm in Algorithm 3.” We have added in Section 3.3, “In this section, we describe the GWG procedure for protein sequences.” We are happy to answer any specific concerns or questions the reviewer has regarding GWG.

4. 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.

The acceptance rate differs based on the learned energy function which is non-trivial to analyze. We are not sure what the reviewer means by “Eq. (4) should give the acceptance rate” when there is no theoretical method to derive the acceptance rate when $f_\theta$ is a neural network (to the best of our knowledge). However, we have added in the appendix tables 10 and 11 which includes a column containing the acceptance rates for each task and difficulty including whether we use our smoothed method or not. We note that the acceptance rate is higher for the smoothed version of our method as further evidence smoothing is necessary for GWG to work.

GFP

AAV

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).

Thank you for this catch! This was a mistake and we have corrected the MH-step with q(x|x') vs q(x'|x).