Protein Discovery with Discrete Walk-Jump Sampling

We thank the reviewer for their enthusiastic and encouraging review of our work, and the emphasis on the novelty of our approach and our real-world in vitro experiments. Below we have responded to each point raised by the reviewer:

1. The reviewer raises a good point that autoregressive LMs have very different capabilities and advantages compared to non-autoregressive models. The natural way to condition on metadata for dWJS is to add prefix tokens for metadata (e.g., species) and sample with those tokens included. It is also true that for wet-lab integration, sampling speed/computational cost is negligible compared to lab experiments. The advantages in sampling speed are more interesting for iteration and method development of new ML methods, where fine-tuning LMs may be insufficient and pre-training from scratch may be impractical. Nevertheless, we believe that pre-trained LLMs are an under-studied model type in ML for biology, hence our investigations into their baseline performance for sequence generation.

2. DC Score is primarily a way to formulate a single, aggregated scalar objective based on feature engineering, to detect if a generative model is producing samples that are “in-distribution” with respect to the training data. We find that DCS is a reliable measure to indicate if samples will express in the lab, and we used DCS as an in silico metric for developing our methods, which allowed us to achieve a high expression rate in the first round of lab experiments. However, we find that for DCS > 0.3, almost all proteins are expressed. We have added this additional context in the discussion of DCS, to more clearly motivate its usage and explain how it was used in silico before conducting wet-lab experiments. We have clarified these points in the discussion of DC score by moving some of the technical details to Appendix F to respect the page limits and adding the following text to Section 3.2: “No multiple sequence alignment or pre-processing of the sequences is required. For convenience and because we have small numbers of examples and low dimensions, we use kernel density estimation (KDE) to compute the joint density. However, DCS is completely general and can be combined with any density estimator…​​Empirically, we find that DCS is a useful in silico evaluation metric for developing generative methods and hyperparameter optimization, and that methods with DCS > 0.3 yield nearly 100% expressing proteins in the wet lab.”

3. We have not explored treating the modeling problem over a probability simplex and sampling, primarily because “greedy” decoding with argmax worked so well for achieving high sample quality and diversity. Sampling from the probabilities is certainly an interesting future direction for increasing the control over sampling.

4. The most important hyperparameter to set is the noise level, sigma. In section 3.2 we derive a “critical noise level” of ~0.5 for our data, which agrees with our empirical findings that 0.5 < sigma < 1.0 leads to good sampling. To provide some intuition, sigma should be large enough to smooth the data and make sampling easier, but not so large that the denoising network cannot recover “valid” samples. The “default” value of sigma in [0.5, 1.0] works well. We have included additional experiments and discussion in Appendix A.5 and Table 6 to illustrate this point.

5. The DC score is an effective method for detecting samples that are not in-distribution with respect to the training data. Empirically, we find that DCS > 0.3 leads to samples that nearly all express in the lab. A low DC score (e.g., 0.06 for ESM2 shown in Table 2) indicates protein samples that will not express. We have updated the manuscript to clarify this point.

6. It is interesting to think about how DC score could be generalized to also account for diversity. In our current setup, we use DCS to evaluate conformity to a reference set, and compute additional edit-distance-based metrics to evaluate diversity. To evaluate diversity, it might be possible to treat a set of generated samples as the reference set and do leave-one-out evaluation of each design’s conformity to the rest of the designs.

7. Our observation that dWJS easily samples many antibody germlines (Figure 1) in a single MCMC chain shows the fast mixing time of our method - this would be impossible with a traditional energy-based model, and is not a capability of autoregressive models. In future work, we intend to do a more theory-focused exploration of the mixing time of dWJS compared to traditional MCMC methods.

We appreciate the reviewer pointing out the uniqueness of our approach and the advantages over traditional energy- and score-based approaches. The reviewer’s two main concerns relate to the effect of the noise level, sigma, and the effect of outliers on the distributional conformity score, which we address below. We hope that the reviewer will consider raising their score following our response.

While the paper touts the simplification to a single hyperparameter choice (noise level, σ) as a strength, it might also be seen as a limitation since the entire model's performance could be sensitive to this single parameter.

1. We agree that the simplification to a single hyperparameter does indeed make the model performance sensitive to this hyperparameter - we appreciate the reviewer’s point that having multiple hyperparameters gives more control over model performance. We would also like to highlight that the decoupling of the “walk” and “jump” steps hopefully alleviates this concern; while the same model can be used for both walking and jumping, it is also possible to decouple them and train two independent models with different noise levels, different architectures, and different hyperparameters. In this way, it is possible to re-introduce complexity into dWJS if desired. It would also be a very interesting future direction to adapt our approach to use a pre-trained encoder, which would introduce many additional hyperparameters. Nevertheless, we believe that there is considerable value in having a simple, robust generative model with high sample quality, and we do not see the single hyperparameter as a fundamental limitation.

Set of properties includes both continuous, binary, and discrete values and estimating the distribution by a kernel density approach may not be very effective. Also DCS may be overly influenced by outliers and extreme data points in the dataset.

2. The DC score is completely general and can be combined with any density estimator, so this is not a fundamental limitation of our approach. We used kernel density estimation here for convenience, and because we have few examples and low dimensions. Density models that combine continuous and discrete values that work in the low sample regime is a promising research direction. It is possible for DCS to be influenced by outliers in a dataset (and we have noted this in Appendix F), but practically speaking we do not encounter this difficulty for protein datasets because there is usually significant sequence similarity between sequences, and “outliers” generally represent out-of-distribution samples that are not “valid” protein sequences, which is exactly what we are trying to detect with the DC score. We have clarified these points in the discussion of DC score by moving some of the technical details to Appendix F to respect the page limits and adding the following text to Section 3.2: “No multiple sequence alignment or pre-processing of the sequences is required. For convenience and because we have small numbers of examples and low dimensions, we use kernel density estimation (KDE) to compute the joint density. However, DCS is completely general and can be combined with any density estimator.”

Is the same value of σ=0.5 used for all reported experiments? It would be nice to see what effect lowering or increasing this value will have on some of the reported metrics for a better assessment of the sensitivity of the approach to this parameter.

3. The sigma value is held constant, but following the reviewer’s suggestion, we have included additional experiments in Appendix A.5 and Table 6, showing the effects of changing the noise level on sample quality. Briefly, under-smoothing by setting sigma too low produces poor quality samples, while over-smoothing by setting sigma too large decreases the diversity of the samples. These intuitive results show the dependence on sigma, while also highlighting that the tradeoff between sample quality, uniqueness, and diversity, is easily controlled with a single hyperparameter, sigma. We have partially copied Table 6 with the new experimental results here for convenience:

How is d obtained in this sentence? "... for the flattened sparse one-hot matrices with vocabulary size 21, d = 6237 ..." Not sure why d is so high if d=L.

4. We have clarified this in the draft: the dimension for flattened one-hot matrices, d = L * v = 297 * 21 = 6237, where L is the length of the sequence and v is the vocabulary size.

Please define the FID score (Fréchet inception distance (FID)) as well as the BLEU and add references.

5. We have defined the Fréchet inception distance and the BLEU score and added appropriate references.

We thank the reviewer for highlighting the intuitiveness and elegance of our approach, and the strength of our in silico and wet lab validation. We have addressed the reviewers’ questions in a point-by-point response below:

Only a single task (antibody design) is evaluated. Testing on other protein classes or discrete domains would be useful.

1. We agree with the reviewer, testing our approach on other protein classes or other discrete data settings would be interesting. We have focused on antibody design in order to present thorough and strong in silico and in vitro empirical evidence of our method’s capabilities. We have not made any assumptions unique to antibodies in developing our method, and it is a priority for future work to extend the empirical evaluation to other domains.

The distributional conformity score for evaluation is introduced late with little detail. More motivation and analysis would improve clarity. Certain details are unclear, like how sequences are aligned and handled.

2. We appreciate the reviewer’s comments related to the distributional conformity score - we have updated the text (moving some of the technical details to Appendix F to respect the page limits) to clarify the motivation for the development of DCS, sequence preparation, and intuition behind DCS. Text added to Section 3.2: “No multiple sequence alignment or pre-processing of the sequences is required. For convenience and because we have small numbers of examples and low dimensions, we use kernel density estimation (KDE) to compute the joint density. However, DCS is completely general and can be combined with any density estimator.”

Your proposed approach operates purely at the sequence level, focusing specifically on antibody sequences. What are your thoughts on the relative pros and cons of sequence-only approaches compared to structure-aware approaches that leverage protein 3D structure information (e.g. inverse folding methods)? Could structure information, either from experiments or structure prediction, be incorporated to potentially improve the performance of your model? For example, might recent powerful inverse folding techniques, where they can be viewed as structure-conditional sequence generative models, like ProteinMPNN [1], ESM-IF [2], PiFold [3], or LM-Design [4] be combined with dWJS to create even more advanced antibody design frameworks? I'm curious to hear your perspective on the value of adding structural awareness and how feasible it would be to integrate with your dWJS approach.

3. Structural data is certainly a rich source of information about proteins, and can be included into the dWJS framework in a number of ways; structure is implicitly incorporated into dWJS because antibodies are aligned according to the AHo numbering scheme. This alignment defines separate “framework” and “complementarity-determining regions” in the sequence, which have precise structural interpretations. We expect that multiple sequence alignment of most protein families already provides implicit structural information. Inverse folding models can be used to score sequence designs from dWJS if there is a known antibody structure of interest. Structure-awareness could also be more directly incorporated into our approach by adapting WJS for sampling from inverse folding models. These are promising directions for future research. Structure-aware approaches introduce additional engineering complexity and a source of experimental noise in the resolution of the crystal structure training data, which may actually degrade sequence sample quality. We also have orders of magnitude more sequence than structure data, and this is particularly acute for antibodies. For sequence design tasks, we find that sequence-only models are quite useful, but for antigen (protein target)-conditioned design where there are no known binding molecules, structure-awareness is a possible approach.