Paper: Protein Discovery with Discrete Walk-Jump Sampling (ICLR-2024)

Reviewer comment (Experiments):
Some of the comparisons to autoregressive language models (LMs) were a bit under-developed. For example, the paper mentions fast mixing sampling, but this is trivial for LMs.

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

Reviewer comment (Writing):
LMs provide a natural way to condition on metadata; how would the current model do that?

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

Reviewer comment (Evaluation):
Claims that LM sampling is expensive in terms of compute aren't that compelling, since for these antibody design applications the wet-lab experiments are orders of magnitude more expensive than the compute to generate the library.

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

Reviewer comment (Presentation):
The 'distributional conformity scores' (DCS) section was interesting and cool, but not central to the paper and felt like a way to add extra math into the paper.

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

Reviewer comment (Experiments):
Please clarify: was DCS used to filter any of the samples for the wet-lab experiments or was it only used to provide the rightmost column in Table 2?

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

Reviewer comment (Evaluation):
Given that your proposed model does not have strong DCS scores compared to some other baselines, how should I be interpreting these results?

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

Reviewer comment (Theory):
Overall, I would recommend removing the DCS content and allocating more space in the main paper to details/derivations of the model fitting/sampling.

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

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

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

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

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

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

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

| sigma              | W_property | Unique   | E_dist   | IntDiv   |
|--------------------|------------|----------|----------|----------|
| 0.1                | 0.378      | 1.0 | 120.6   | 60.0     |
| 3.0                | 0.130      | 0.995    | 44.2     | 30.0     |

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

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

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

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

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

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

Reviewer comment (Evaluation):
The distributional conformity score for evaluation is introduced late with little detail. More motivation and analysis would improve clarity.

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

Reviewer comment (Reproducibility):
Certain details are unclear, like how sequences are aligned and handled.

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

Reviewer comment (Novelty):
The proposed approach operates purely at the sequence level, focusing specifically on antibody sequences. The reviewer asks for 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).

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

Reviewer comment (Experiments):
The reviewer asks whether structure information, either from experiments or structure prediction, could be incorporated to potentially improve the performance of the model.

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

Reviewer comment (Novelty):
The reviewer inquires if recent powerful inverse folding techniques, such as ProteinMPNN, ESM-IF, PiFold, or LM-Design, might be combined with dWJS to create more advanced antibody design frameworks.

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