Data Distillation for extrapolative protein design through exact preference optimization

We appreciate the thoughtful comments and summary from reviewer qbi3, which have helped improve the quality of the manuscript.

W1: regarding preference optimization relying on a scorer function to define the triplets

We thank the reviewer for pointing out this important concern. We hope that clarifying the difference between scorer and evaluator will mitigate the concern. Scorer is a model that only sees the training region and therefore might perform poorly in out-of-distribution scenario as rightly suggested by the reviewer. Therefore, the current scorer used in our model is noisy particularly with respect to sequence in extrapolation region. On the other hand, the in-silico evaluator is a model which is trained on both training and extrapolating regions. Therefore, it is expected to perform better in the extrapolation region and can be thought as a noisy surrogate of wet-lab experiments. In section C (scorer vs evaluator), we have shown that scorer and in-silico evaluator do not correlate well in the extrapolation region as expected.

Q1: clarification regarding definition of hard triplet ranking wise for scorer distillation

We thank the reviewer for the clarification question. We have updated section 6.1 in main text to better define the hardness for scorer distillation. Particularly, for a given $x_{prompt}$, one can sample $K$ sequences $y_{1:K}=[y_1, y_2, .., y_K]$ from $P_{pair} (.|x_{prompt})$. We choose the triplewise rankings for all pairs based on $s_{i,j} = f(y_j) - f(y_i)$ where $f(y_j) < f(x_{prompt}) < f(y_i)$. By definition, $s_{ij} > 0$ can be considered as hard example since the pairwise model has wrongly generated $y_j$ with lower predicted fitness than prompt. We would like to encourage the pairwise model to generate more sequences similar to $y_i$ since its predicted fitness is higher than prompt sequence. This approach can be considered as scorer distillation since we are attempting to distill the scorer's knowledge in both training and extrapolation regions into the generative model.

Q2: re how important is the mask infilling step

We thank the reviewer for asking this detailed question. The mask infilling step is important while training the reference pairwise model. We followed the recommendation from [1] to do masking-infilling starting from the training sequences (seeds). Scorer function is utilized to assess whether the newly generated pair (seed, sequence) has a small but meaningful improvement in the desired direction. Since the focus of the paper was on triplet creation rather than pairs, we have not done ablation studies on importance of mask infilling for pairwise models. We followed the recommendation of [1].

[1] Padmakumar, Vishakh, et al. "Extrapolative controlled sequence generation via iterative refinement." ICML 2023.

We appreciate the thoughtful comments and summary from reviewer jZ4t, which have helped improve the quality of the manuscript.

W1-Q1: clarification question regarding the computational cost of scorer-distillation

We thank reviewer for highlighting this important concern. The reviewer has correctly pointed out that scorer-distillation is computationally expensive and that is the main reason we have not used this approach as our default. We have added table 7 in supplementary for computational cost of creating one triplet based on scorer distillation or the default one. We added a few lines in section 6.1 in main text to highlights this. As mentioned in the revised paper, scorer distillation is more than 200 times more expensive to run on p3 GPU (V100) machines in comparison to the default for GFP dataset (280 length protein).

W2-Q2: Regarding benchmarking the proposed methods against all baseline on a new dataset

We thank the reviewer for their excellent recommendation to further benchmarking on new datasets. While it is always better to show robustness of the proposed method on more datasets, due to the time limitation, we were not able to add another dataset. However, we hope that additional experimentations (consecutive vs non-consecutive bins, sequence similarity constrains) and thorough hyper-parameter tuning for all baselines help in showing the robustness of the proposed approach.

W3: evaluation through in-sillico evaluators that are trained on whole data

We thank to the reviewer for highlighting this important aspect. We acknowledge that there is a possible mismatch between benchmark outcomes utilizing in-sillico evaluators versus real world outcomes. However, we have tried to reduce the gap by carefully splitting the sequences between the “training” and “extrapolation” regions to ensure separation in both sequence and fitness spaces. The generative models and scorers have been trained on "training region" only and the in-sillico evaluator have been trained on whole dataset.

Q3: ablation studies for triplet construction

We thank the reviewer for their excellent suggestion. We have added additional experimentations regarding how to create triplet specifically on 1) being generated from consecutive vs non-consecutive bins 2) whether using sequence similarity constrain between prompt and generated sequences. We have reported the results in section 6.1 in main text and section D in supplementary. The results highlight that naively creating preference data from non-consecutive pairs would confuse the model's learning and worsen the performance on 3 out of 4 splits. We hypothesize that, in order to properly learn from non-consecutive bins, one needs to develop a more sophisticated prompt to incorporate the distance between bins as well. Similarly, the results suggest that max mutation 15 between seed and desired/undesired sequences would slightly improve the performance on hard splits of AAV and GFP datasets while worsen it for medium splits. However, calculating sequence similarity for large dataset is computationally expensive therefore we would not suggest it as our default.

Q4: clarification question regarding on how to balance between exploration and exploitation

We thank the reviewer for their clarification question regarding balancing between exploration and exploitation. Temperature ($\tau$) is one of the main hyper-parameters that creates balance between exploration and exploitation. Temperature ($\tau$) hyper-parameter is used to control the randomness in generation and as a way to balance between greedy search (token with max probability when $\tau \rightarrow 0$) and uniform sampling (when $\tau \rightarrow \infty$). As highlighted from the results of the hyper-parameter tuning in section F in supplementary, in order to increase the diversity when scorer is used in inference, it is recommended to increase the temperature from 0.7 to 1.0.

We appreciate the thoughtful comments and summary from reviewer P9DR, which have helped improve the quality of the manuscript.

W1: regarding the mismatch between the benchmark outcomes and real world outcomes

Thanks to the reviewer for highlighting this important aspect. We acknowledge that there is a possible mismatch between benchmark outcomes and real world outcomes. However, we have tried to reduce the gap by carefully splitting the sequences between the “training” and “extrapolation” regions to ensure separation in both sequence and fitness spaces.

W2: clarification regarding hyper-parameter tuning for proposed method

Thanks to the reviewer for highlighting this important clarification question. Based on the hyper-parameter tuning shown in section F in main text, we have chosen temperature 0.7 for non-scorer and 1.0 for when we are using scorer in inference. We have modified the section 5 in main text to clarify more regarding the hyper-parameter tuning for proposed method.

W3: clarification regarding hyper-parameter tuning for baselines

Thanks to the reviewer for highlighting this important concern. We have performed a thorough hyper-parameter tuning across all baselines and reported the results in section G. We have updated the results of the baseline methods according to their best hyper-parameters. Therefore, we have updated Figure 2 and Tables 2 and 5. The main conclusions after a thorough hyper-parameter tuning remain consistent with our previous findings. The main reason for the consistency is that reference papers [1] and [2] did thorough hyper-parameter tuning in their papers for the same datasets.

W4: suggestion regarding Figure 3

We thank the review for their constructive feedback. We have moved the figure 3 to supplementary and added more details regarding new experimentations (consecutive vs non-consecutive bins, sequence similarity constrains) and hyper-parameter tunings.

Q1: How is the sampling with scorer guidance performed?

We thank the reviewer for their clarification question. We have added more details in the inference section (section 3.6). Specifically, at iteration t, suppose there are M seeds available. We will generate N sequences per seed and choose the best one based on scorer as seed for next iteration.

Q2: In what order were $\beta$ and $\tau$ determined?

We thank the reviewer for their clarification question. We have tuned the $\beta$ first since it is the hyper-parameter used in fine-tuning and then we have tuned the temperature ($\tau$) since it is used in the inference. We have added these details in section 5 in main text.

Q3: What choices went into the selection of hyper-parameters for the baseline models?

We thank the reviewer for their clarification question. In the first submission, we didn’t perform hyper-parameter selection for baseline models since most of the state-of-the-art models ([1] and [2]) have been already reported on the same datasets/splits. Therefore, they have already done the hyper-parameter selection in their respective papers. We have utilized their default values reported in their papers or Github. However, in the revised manuscript, we performed hyper-parameter tuning as per reviewer's recommendation.

Q4: How does performance of the baseline models change when benchmarked with the same attention as the proposed model?

After thorough hyper-parameter tuning for all baselines methods with the same focus as given to the proposed methods, we can observe that the main conclusions remain consistent with our previous findings. We want to re-iterate that this finding is due to reference papers ([1] and [2]) having done their own thorough hyper-parameter tuning in their papers for the same datasets.

[1] Kirjner, Andrew, et al. "Improving protein optimization with smoothed fitness landscapes. ICLR, 2024.

[2] Lee, Minji, et al. “Robust Optimization in Protein Fitness Landscapes Using Reinforcement Learning in Latent Space” ICML 2024

Each of my points were thoughtfully addressed in the revised version of the manuscript. The work now contains thorough descriptions of the methods, how baselines were tuned to create fair evaluations, and all relevant results were updated in line with this. I believe that future researchers will be able to build on the perspective of this work and trust the findings. I raise my score to reflect the growth of the work.

Some minor aesthetic nitpicks: elements in Figure 1 aren't spaced correctly (e.g. the sub boxes in triplet curation aren't centered) and could benefit from a larger font for readability utilizing the negative space. Not a major contention, but would improve a camera ready version.

Q4: sequence similarity constrain for pair and triplet generation

We thank the reviewer for their insightful question. We have not considered any sequence similarity constrain for pair creation. Since the focus of the paper was on triplet creation rather than pairs, we followed reviewer’s recommendation and did an ablation studies for sequence similarity in preference data creation. Results that are reported in section 6.1 of main text and section D.5 of supplementary suggest that max mutation 15 between seed and desired/undesired sequences would slightly improve the performance on hard splits of AAV and GFP datasets while worsen it for medium splits. However, calculating sequence similarity for large dataset is computationally expensive therefore we would not suggest it as our default.

Q5: Clarification regarding what $r_{\phi}$ (reward) stands for.

We thank the reviewer for their clarification question. Due to the page limitation, we have added their exact explanations in the section A of supplementary.

Q6: Clarification regarding top-p, top-k and role of temperature

We thank the reviewer for their clarification questions. We added few lines in the inference (section 3.6), to explain top-p , top-k and temperature.

Q7: the default operating mode of EXO with the “All” triplet creation (Table 2)

We thank the reviewer for their clarification question. No it is EXO with the “Mistakes” triplet creation. We have added clarification in the paper (Table 3).

Q8: updating citation "Improving protein optimization with smoothed fitness landscapes."

We thank the reviewer for their attention to details. We have modified it in the revised paper.

We sincerely appreciate reviewer jBZU for their thoughtful feedback, which has significantly improved the manuscript's quality.

W1: Addressing the "separation" concept introduced by [2] versus "extrapolation" concept introduced by [1,4] and our paper

We thank the reviewer for highlighting this concern. We were not aware of [2] introducing the “separation” concept which is closely related to “extrapolation” region. As pointed out by the reviewer, the exact definition of “extrapolation” and “separation” are different. For example, we have considered “extrapolation” in both fitness and sequence space. Following [1], all sequences in the training data have mutational gap of 6 or 7 meaning they are 6 or 7 mutations away from top 1% sequences in the extrapolation region. In addition, there is a clear separation in the fitness landscape as well where the scorer and generative models are trained only on “training region” with noisy understanding of the “extrapolation” region.

W2: Addressing "optimization in latent space" in related work

Following reviewer’s recommendation, we have added a category in the related work section for including “protein optimization in the latent space”. We included the recommended references ([2] and [3]) as well as one of the-state-of-the-art models we have compared with [4]. The latter reference has utilized latent space of protein large language model (ESM-2 with 650M parameters). They have introduced a novel reinforcement learning approach to navigate through the latent space. Our benchmark has shown that our model has outperformed theirs in 3 out of 4 splits.

Q1: clarification regarding how pairs are generated

We thank the reviewer for their clarification question. We have followed the recommendation from [5] for pair creation. We generated perturbed sequences by masking-infilling starting (5% masking which is around 2 mutations for AAV with length of 28 and 14 mutations for GFP with the length of 280) from the training sequences (seeds). Scorer function is utilized to assess whether the newly generated pair (seed, sequence) has a small but meaningful improvement (e.g. 0.5) in the desired direction. We have not performed ablation studies to assess the impact of mask-infilling approach for pair creation versus pairs created from consecutive bins.

Q2: triplets generated from consecutive vs non-consecutive bins

We thank the reviewer for their insightful questions that have significantly helped improved the revised paper. We have added an additional experimentation to addressing the consecutive vs non-consecutive bins. We have created preference data from non-consecutive bins in similar fashion as we have done in our default version (consecutive bins). We have reported the results in section 6.1 in main text and section D in supplementary. The results highlight that naively creating preference data from non-consecutive pairs would confuse the model's learning and worsen the performance on 3 out of 4 splits. We hypothesize that, in order to properly learn from non-consecutive bins, one needs to develop a more sophisticated prompt to incorporate the distance between bins as well. Due to the limited time we could not implement this but we think the model should benefit from it and can be considered as a future direction.

Q3: clarification question regarding "meaningful improvement" hyper-parameter in pair generation

We thank the reviewer for their clarification question. We have followed the recommendation from [5] to set the hyper-parameters (e.g. max fitness differences between pairs, or number of masking, etc) for pair creation and perturbations. We have added more details in the section 4.3 of main text.

[1] Kirjner, Andrew, et al. "Improving protein optimization with smoothed fitness landscapes. ICLR, 2024.

[2] Ghaffari, Saba, et al. "Robust Model-Based Optimization for Challenging Fitness Landscapes. ICLR, 2024.

[3] Gómez-Bombarelli, Rafael et al. “Automatic Chemical Design Using a Data-Driven Continuous Representation of Molecules.”

[4] Lee, Minji, et al. “Robust Optimization in Protein Fitness Landscapes Using Reinforcement Learning in Latent Space.” ICML 2024

[5] Padmakumar, Vishakh, et al. "Extrapolative controlled sequence generation via iterative refinement." ICML 2023.

We really appreciate the reviewer for helping us improving the manuscript quality. We would like to ask if our response has addressed your concerns? If not, we are very happy to discuss and further clarify.

Kind Regards, Authors