A Variational Perspective on Generative Protein Fitness Optimization

We thank the reviewer for providing feedback on our manuscript and appreciate the comments. We will also discuss related work that you mentioned.

• The predictors and the oracle are sourced directly from previous works.

Indeed, we take this from recent work from ICLR’24 [1]. As mentioned to reviewer TQyg, we will further elaborate on this.

• Parameterization of the NN and time step sampling choices for flow matching.

Our experiments have shown to be very robust against the choice of time steps used for inference, as we use a flow-based approach that is much less sensitive to hyperparameter settings than standard time-discrete diffusion models. The choice of the architecture at hand is rather generic and used as such for many other related tasks.

The time steps are sampled uniformly between 0 and 1 in training, which is the standard for flow-based approaches, we saw no need to change it.

• Nowadays, transformer-based models, such as ESM, have shown superiority over CNNs.

Transformer-based models are standard in protein design, we show that guided flow matching performs well even with simpler architectures (although our network does include attention layers). More expressive models might yield further improvements. Note, previous work [1,2] has discussed that CNN-based models are competitive in limited-data settings.

As noted in Sec. 5, ESM embeddings are highly expressive, and we plan to explore their potential in future work.

• Regarding classifier guidance, it remains unclear why the objective is to optimize towards a specific fitness value rather than maximizing the fitness value.

Maximizing the fitness value is our ultimate goal; however, we intended to demonstrate the expressiveness of our approach by steering sequences toward specific fitness values. In practice, we compared both formulations which performed similarly. Given that our variational framework allows for fitness maximization in the likelihood term, we will explicitly incorporate this into Eq (3). If of interest, we can provide the results in the supplementary.

• The idea of applying flow matching on the latent space and guiding generation with a classifier has been proposed many times in different domains.

We agree that classifier guidance is widely used in generative AI. However, it cannot be trivially applied to discrete sequences, and its application in latent spaces for protein fitness optimization has not been explored before, to the best of our knowledge.

• All evaluations are conducted on relatively small benchmarks which casts significant doubt on the reliability of the evaluation.

We respectfully disagree with this concern. Our evaluation protocol and benchmarks align directly with established standards from previously published work [1]. Such dataset sizes reflect realistic scenarios. Moreover, the in-silico oracle was trained on the complete DMS datasets with 56,086 mutants for GFP and 44,156 for AAV.

• The dataset is split according to fitness values... a more appropriate split should be based on mutation depth.

Both tasks medium and hard indeed are based on fitness values as well as mutation gaps. This is discussed in [1] (independent original work) and in our manuscript in Sec. 4.1 and in Tab. 6.

• There is a lack of clarity regarding the dataset used to train the flow matching and VAE models.

For each of the four tasks, the flow matching and VAE models are trained exclusively on their respective limited datasets to emulate a realistic scenario. We will explicitly clarify this in the revised manuscript.

• The achievement of high fitness values comes at a substantial cost of sacrificing diversity and novelty.

We agree that there is an inherent trade-off between achieving high fitness values and maintaining diversity, which is also reflected in our approach. Given that our method achieves a substantial improvement in fitness compared to alternative approaches, we believe the benefit outweighs the marginal drop in diversity compared to competing methods.

• How the authors perform the grid search for hyperparameters is unclear.

The plots shown in the manuscript are generated by the oracle, which are an ablation and not the source of our hyperparameters, similar to [1]. We agree, this was not explained well enough. For our parameter estimation we combine grid searches over the diversity as well as the fitness given only by the predictor. We will add the diversity and predictor fitness plots (which resemble Fig. 5) in the appendix. Ultimately it is a heuristic tradeoff between fitness and diversity, which we will further elaborate upon.

• The authors should incorporate more realistic settings and additional baselines.

See response to reviewer jm6F, last point.

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

[2] Dallago, Christian, et al. "FLIP: Benchmark tasks in fitness landscape inference for proteins."

We sincerely thank the reviewer for their careful evaluation of our manuscript and for the positive assessment of our work. Below, we address each point individually:

• The authors may be interested in this very recent in silico benchmark, which introduces synthetic test functions that can be used for benchmarking black-box optimization methods like VLGPO.

Thank you for directing us toward this interesting recent work. We will examine it in more detail and consider its suitability for future applications, as synthetic test functions appear very promising for further assessing the generated sequences.

• The biggest weakness is the reliance on GFP and AAV datasets as the benchmarks, which is a general limitation that affects the field. Have the authors considered other FLIP tasks, or other related tasks?

This is indeed an important point. While our method could, in principle, extend to other FLIP or ProteinGym tasks, we chose the GFP and AAV datasets as they currently are well-established benchmarks to allow for fair comparisons against our proposed approach. The main focus here was to investigate performance in limited initial datasets, reflecting common practical scenarios.

Nevertheless, exploring additional tasks beyond GFP and AAV is an exciting avenue for future research, and we fully agree on the importance of designing and adopting broader benchmarks in this context.

We would like to thank the reviewer for their valuable comments and feedback. Below, we address each point individually:

• Since the evaluation of the methods relies on the predictor, it would be helpful to provide details on the accuracy of $g_{\phi}$ as in-silico oracle.

Thank you for highlighting this important aspect. We reused the oracle and predictor from [1] entirely out-of-the-box, without any modifications, aiming to present a method in which alternative predictors or oracles could easily be used as drop-in replacements. We will comment on this in more detail in the revised version.

However, since [1] does not explicitly report the final performance of the in-silico oracle, we have now computed the Mean Squared Error (MSE) on a subset of 512 randomly selected samples from the ground truth dataset $\mathcal{S}^*$. The oracle’s predictions closely follow the target fitness values, resulting in MSE values of 0.012240 for GFP and 0.002758 for AAV.

• In Tables 2 and 3, the variability in diversity and novelty among different methods is significantly greater in GFP than in AAV, which is another observation worth discussing.

We appreciate the reviewer’s insightful observation. This aligns well with our findings, indicating that the GFP setting indeed appears more challenging compared to AAV. This discrepancy may arise due to GFP sequences being longer and thus representing sparser and higher-dimensional search spaces. We will discuss this observation in the revised version of the manuscript.

• It would be helpful if the author could also provide more analysis and statistics, such as diversity at different fitness, of the GFP and AAV dataset.

Thank you for this valuable suggestion. The diversity within the top-performing (99th percentile) sequences of the entire dataset $\mathcal{S}^*$ is 4.73 for GFP and 5.23 for AAV, as briefly discussed in Section 4.3. In contrast, the diversity of sequences in the four considered tasks is notably higher: for GFP (medium) it is 14.5, for GFP (hard) 16.3, for AAV (medium) 15.9, and for AAV (hard) 18.4. The training datasets naturally contain more diverse sequences, as it includes those with lower fitness levels. In contrast, the top-performing sequences exhibit lower diversity, aligning closely with the diversity observed in Tables 2 and 3. We will highlight this observation in the revised manuscript.

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

We would like to thank the reviewer for the constructive feedback and the insightful comments, which help to improve the quality of the manuscript.

We fully agree with the reviewer that assessing fitness with a computational model is less reliable than direct wet-lab validation. Experimental validation is the final goal, in-silico evaluation is an important step along the way. We also appreciate the comment regarding the potential overestimation of the performance gap compared to baseline methods. This is indeed a challenge for the community, but hard to quantify precisely.

Below, we address the reviewer’s more specific questions directly:

• Can you clarify the connections between the proposed approach and variational inference?

Thank you for this point. VLGPO uses variational inference through the initial embedding step to obtain continuous latent embeddings of the discrete protein sequence variants - the autoencoder is trained with variational inference. However, our approach extends this by introducing additional generative modeling components, namely flow matching and classifier-guided sampling, to further refine and control the generation process. We will clarify this connection explicitly in the revised manuscript.

• Can you comment on the poor performance of GFN-AL?

To ensure a fair evaluation, the results for all methods are taken from a comprehensive comparison done in recent published work [1]. Hence we can only speculate about the poor performance. GFN-AL might struggle in limited-data scenarios due to sparse reward signals from the few observed mutants, which could potentially limit effective exploration.

• Can you contrast your approach with classifier-free guidance?

Thank you for this question — we find it very interesting, as it highlights the advantages of using a separate classifier for guidance, which we have addressed in Fig. 3. There, we compare our method to a conditional model that directly learns the posterior $p(x|y)$. In that setup we additionally tested classifier free guidance, but it did not make a noticeable difference. As shown in both the figure and the ablation tables, classifier guidance allows for more effective steering of the generation process: in challenging scenarios like GFP (hard), classifier-free guidance struggles to produce higher-fitness sequences, while classifier guidance successfully targets high-reward regions.

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