Paper: Improving protein optimization with smoothed fitness landscapes (ICLR-2024)

Reviewer comment (Theory):
What 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).

Author response:
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.

Reviewer comment (Theory):
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 me why that would be the case for proteins.

Author response:
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.

Reviewer comment (Experiments):
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?).

Author response:
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.

Reviewer comment (Evaluation):
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.

Author response:
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].

Reviewer comment (Experiments):
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.

Author response:
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.**

Reviewer comment (Experiments):
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)?

Author response:
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.

Reviewer comment (Experiments):
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*?

Author response:
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

Reviewer comment (Experiments):
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).

Author response:
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.**

Reviewer comment (Novelty):
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.

Author response:
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].

Reviewer comment (Writing):
Page 2, *"mutation is proposed renewed gradient computations."* — Something is wrong in this sentence.

Author response:
Thank you for this catch! We have corrected it.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Writing):
Page 4, eq (2) *"H(x)"* — As far as I can see, H(x) has not been introduced(?)

Author response:
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."

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Writing):
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".

Author response:
Thank you for this improvement. We have applied the suggestion.

Reviewer comment (Experiments):
The proposed method has many hyperparameters to tune.

Author response:
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.

Reviewer comment (Evaluation):
Please explain why after smoothing, the diversity and novelty of the final set of sequences decreases.

Author response:
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.

Reviewer comment (Experiments):
In 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?

Author response:
We perform extensive study in more mutation gap and starting range settings. Our results are in Appendix C.2 and summarized here.
We use a consistent mutation gap of 8 in the case of GFP, and 13 in the case of AAV, while decreasing the range percentile cutoffs. We do this to control how much smaller the starting dataset is getting as we increase the difficulty, as increasing the mutation gap can sharply decrease $|D|$.

In the manuscript, we stopped at the hard difficulty, since other baselines fail at this point. We find that for GFP, if we restrict the range to the bottom 10 percent of sequences and enforce a mutational gap of at least 8, then our smoothing method fails. However, for AAV, we find the model is able to still find signal and achieve 0.384 fitness despite the data being limited to the 10th percentile and mutation gap of 13 where $|D|=476$. Presumably, the difference arises as a result of the different contours of the fitness landscapes of the AAV and GFP

Finally, although we were not able to reproduce the results of an easy task for all baselines during the rebuttal period, we would agree with your assumption that as the mutational gap decreases, all methods should perform comparably.

**GFP Alternative Benchmarks (Mutation Gap = 8)**

| Range | D | Smoothed | Acceptance Rate | Median Fitness | Diversity | Novelty |
|-------|---|----------|-----------------|----------------|-----------|---------|
| 0-30% | 1129 | No | 0.01 | 0.00 | 22.86 | 17.0 |
|       |      | Yes | 0.43 | 0.67 | 4.45 | 9.12 |
| 0-20% | 792  | No | 0.01 | 0.00 | 22.22 | 16.48 |
|       |      | Yes | 0.59 | 0.60 | 5.42 | 9.82 |
| 0-10% | 397  | No | 0.01 | 0.00 | 23.02 | 16.82 |
|       |      | Yes | 0.01 | 0.00 | 15.73 | 21.17 |

**AAV Smaller Benchmarks (Mutation Gap = 13)**

| Range | D | Smoothed | Acceptance Rate | Median Fitness | Diversity | Novelty |
|-------|---|----------|-----------------|----------------|-----------|---------|
| 0-30% | 1157 | No | 0.001 | 0.30 | 0.53 | 6 |
|       |      | Yes | 0.473 | 0.31 | 13.80 | 14.69 |
| 0-20% | 920  | No | 0.001 | 0.39 | 6.43 | 9.13 |
|       |      | Yes | 0.546 | 0.27 | 15.98 | 19.41 |
| 0-10% | 476  | No | 0.001 | 0.25 | 3.08 | 5.63 |
|       |      | Yes | 0.64 | 0.38 | 7.05 | 9.48 |

Reviewer comment (Experiments):
As 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?

Author response:
Yes, we have explored different amounts of smoothing in Figure 3 where we show setting the smoothing weight $\gamma$ too high can be detrimental. Clearly too much smoothing destroys the signal and is detrimental. We note setting the optimal $\gamma$ as a limitation of our method in Section 5 (Discussion) and is an important future direction that requires non-trivial theoretical results. Our work demonstrates the importance of smoothing that other works have overlooked.

Reviewer comment (Writing):
The paper is hard to read and understand overall. Having one figure to show the overall flow of logic could help. Fig. 1 seems to do the job. There is some confusion between training and sampling. It appears that the authors first train f(x) and then use f(x) as a surrogate function for MCMC optimisation, but this point does not come out naturally.

Author response:
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.

Reviewer comment (Writing):
The construction of the KNN graph could be described more clearly (see equation above Eq. (1)).

Author response:
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.

Reviewer comment (Writing):
Symbols in Eq. (2) are ill-defined. The authors should provide in the appendix some details of GWG and reference the appendix in the main text.

Author response:
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.

Reviewer comment (Writing):
The clustered sampling section should be explained better.

Author response:
We apologize for lack of clarity with the clustered sampling. Due to limited space, we have added an additional figure in the appendix that illustrates clustered sampling. Section 3.3 now states “An illustration of clustering is provided in Figure 6.” We are happy to address any specic component of clustered sampling that was not clear.

Reviewer comment (Evaluation):
Is there a way to test that the surrogate function by itself is good enough? The authors look at the overall performance, which could infer correctness of the surrogate function.

Author response:
See our global response. We have analyzed the overall prediction error of our surrogate function with and without smoothing in each of our settings and found smoothness to be necessary in achieving improved generalization to unseen sequences.