Antibody DomainBed: Out-of-Distribution Generalization in Therapeutic Protein Design

We thank the reviewer for the thoughtful review. We address the concerns below and complement it with a revision of our paper.

[ prove the fidelity of the synthetic labels. More analysis should be incorporated to prove that the computed labels are well aligned with wet-lab validations.]

The change in Gibbs free energy, $\Delta G$, and the dissociation constant, $K_D$, can be shown to be theoretically equivalent up to a proportionality; we have $\Delta G = RT \ln K_D$, where $R$ is the gas constant, $1.98722 {\rm cal}/K\cdot {\rm mol}$, and $T$ is the absolute temperature [Jorgensen2008]. Free energy perturbation (FEP) has been applied to identify mutant antibodies with high affinity [Clark2017, Clark2019, Zhu2022], supporting the use of $\Delta \Delta G$ as a synthetic proxy for $K_D$.

That said, free energy is not directly computable and we have used Rosetta and FoldX scores as approximations. The Rosetta-computed $\Delta \Delta G$ is still capable of separating binders and non-binders from fluorescence-activated cell sorting (FACS) well and those from surface plasmon resonance (SPR) even better [Mason2021, Mahajan2022]. (FACS is a higher-throughput, lower-SNR method of identifying binders.) These points have been added to Section A.6 of the paper.

[validate them with wet-lab experiments.]

The validation on an experimental binding dataset was done in [Hummer2023], from which env 4 was curated. Models trained on FoldX-generated $\Delta \Delta G$ values are shown to achieve ROC AUC of 0.88 and average precision of 0.77 on the task of classifying binders from FACS [Mason2011].

[the causal features might be physico-chemical and geometric properties at the binding interface. However, no further analysis or interpretations on whether the tested DG algorithms have learnt this invariance are provided.]

The included Domain Generalization algorithms in general attempt to learn invariant, latent representations. Since we cannot directly relate those latent features to observable physico-chemical properties, we turned to saliency mapping, Figure 6 in the manuscript. We notice that, relative to SMA and IRM, ERM displays muted behavior in the regions known to interact with the antigen paratope.

[vanilla ERM beats most DG methods. Does this mean most of the DG methods can not learn the hypothesized invariance at all?]

Although this may be true on environments 0, 1, 2, 3, we observe that

1. in the zero shot setting (env 4) as presented in Table 4 in the Appendix, DG methods and ensembling improve over ERM when evaluated on a new target;
2. in Table 6 we notice significant advantage of some DG methods compared to ERM when using these models in a concept-shift setup, when the labeling function changes from $\Delta \Delta G$ to binding measurements (env 5).

We thank the reviewer for the thoughtful review. We address the concerns below and complement it with a revision of our paper.

[validation of the quality of these labels either experimentally or against other computational methods]

The score function we use for computing the $\Delta \Delta G$ predictions has been experimentally validated in the original publication [Park2015]. Our new encouraging results on env 5 with wet-lab binding measurements additionally verify the utility of these scores.

[The composition of the benchmark dataset seems to be somewhat arbitrary. [...] There is no justification behind the choice of generative model parameters and other environment metadata]

We set the environments to mimic the shift occuring between lab-in-the-loop iterations, Figure 2 in the manuscript, where the generative models change due to various improvements and the seeds and antigens are fixed. Arguably, there are multiple possible settings and users of our dataset can construct their dataset as they wish to measure specific domain shifts. The parameter in the WJS corresponds to the "smoothing scale," which controls the edit distance between the seed and the generated designs. Sequence similarity (or edit distance) is a common, widely-accepted metric used by scientists in the therapeutic antibody domain, as a way to categorize different 'subgroups' of antibodies [Miho2019]. We therefore rely on the edit distance as a metric to compare antibody sequences.

[demonstrate that better performance on the benchmark dataset leads to better active learning performance]

The development of well-performing, robust predictors immediately translates to the success of an active learning pipeline, as these ML predictors are used in the active-learning selection process before sending sequences to a wet-lab.

[Why are the results binarized?]

There is significant noise in the computed $\Delta \Delta G$ between -1 and 1 kcal/mol, because Rosetta and FoldX are less accurate at predicting $\Delta \Delta G$ for mutations with only a small effect on binding [Sirin2016, Hummer2023]. We therefore remove highly uncertain labels between -0.1 and 0.1 kcal/mol before attaching binary labels: label 1 if $\Delta\Delta G < -0.1$ (stabilizing) and label 0 if $\Delta\Delta G > 0.1$ (destabilizing).

Binarized labels also align with potential use cases in experimental binding prediction. While surface plasmon resonance (SPR) yields continuous affinity values, antibody screening is often carried out using noisier but higher-throughput measurement protocols like fluorescence-activated cell sorting (FACS) and enzyme-linked immunosorbent assay (ELISA) aimed at (binary) identification of binders.

Lastly, algorithm development and benchmarking for domain generalization has been focused on classification tasks. Regression is an interesting future extension particularly for lab-in-the-loop design.

[environments consist of mixtures of antigens]

While the reviewer is correct in that lead identification and optimization efforts are often based on a particular target determined in the target discovery stage, the training data for the models need not be restricted to the single target of interest. The mechanistic (causal) stabilizing interaction between an antibody and antigen can be learned across targets. Data from multiple targets can arguably offer more information about factors governing the interaction between the epitope and paratope.

We would like to further clarify that a given target can have multiple "antigen variants" with varying sequences. We thus expose our model to various sequences on the antigen side for a given target.

In addition, target generalization is of interest in the de novo antibody design setting, where we initially do not have annotated binding data. Environment 4 emulates this by probing the ability of a model trained on HIV1 and SARS-CoV-2 data to generalize to an unseen target, HER2.

[Park2015] Park, Keunwan, et al. "Control of repeat-protein curvature by computational protein design." Nature structural & molecular biology 22.2 (2015): 167-174.

[Miho2019] Miho, E., Roškar, R., Greiff, V., & Reddy, S. T. (2019). Large-scale network analysis reveals the sequence space architecture of antibody repertoires. Nature communications, 10(1), 1321.

[Sirin2016] Sirin, S., Apgar, J. R., Bennett, E. M., & Keating, A. E. (2016). AB‐Bind: antibody binding mutational database for computational affinity predictions. Protein Science, 25(2), 393-409.

We thank the reviewer for the thoughtful review. We address the concerns below and complement it with a revision of our paper.

["foundational model"]

By foundation model, we meant a model pretrained on a large dataset, with good finetuning and zero-shot abilities. For more details, please see [Bommasani2021].
We chose ESM-2, a BERT-style model pretrained on general proteins as this is a standard model for representation learning tasks. There are antibodies and antigens in the pretraining dataset.

[One baseline to compare ESM2 against [...] using foundational models improves solving DG is tenuous]:

We added a graph neural network backbone, GearNet [Zhang2023], pretrained on AlphaFold2-predicted protein structures. Due to time constraint, we only ran some DG models (VREx, ERM and ENS). We observe that GearNet is the best performing model; one key reason for this boost in performance is that $\Delta\Delta G$ labels are directly computed from structure.

[The other conclusions, namely that ensembling helps baseline models and that choosing a good validation set is important are really not surprising and do not constitute a substantial contribution.]

Ensembling outperformed all DG models on our task, as is the case on image classification [Arpit2022]. We cannot conclude that invariant models do not outperform ERM when the hyperparameter ranges, model size, and validation set are carefully controlled. This was the main conclusion from DomainBed on image-based datasets [Gulrajani2007]. We recover this result in the real-world application of antibody design. Additionally, we now update the benchmark with results on unseen target (env 4) and real binding measurements (env 5) where we observe a performance boost with ensembling and DG approaches. Note that for new application domains and benchmarks, where there is no standard validation set, defining a validation close to the test set is key to controlling domain shifts; this is why we believe that emphasizing the importance of choosing a good validation set can help better perform model selection.

[Bommasani2021] Bommasani, Rishi, et al. "On the opportunities and risks of foundation models." arXiv preprint arXiv:2108.07258 (2021).

[Arpit2022] Arpit, Devansh, et al. "Ensemble of averages: Improving model selection and boosting performance in domain generalization." Advances in Neural Information Processing Systems 35 (2022): 8265-8277.

[Zhang2023] Protein Representation Learning by Geometric Structure Pretraining, ICLR 2023.

[Gulrajani2007] Gulrajani, I., & Lopez-Paz, D. (2020). In search of lost domain generalization. arXiv preprint arXiv:2007.01434.

[Could it be that the authors always create a perfect setting for ensemble models and simple moving averages? If showing this conclusion was the main purpose of the paper, are there simpler or additional datasets that could support it?]

The choice of this benchmark was done with the motivation of reproducing the lab-in-the-loop setting. The cost of wet-lab experiments leads to a low-data regime where we focus only on a few antigens and restricted number of designs. The low-data regime is particularly difficult; in the larger data regime (assuming enough diversity in data), models should perform better. We did not create a specific environment split to favor the ensembling models, we tried to reproduce two common domain shifts: covariate shift (variations due to the generative models in environments 0, 1, 2, 3) and concept drift (variation in the labeling function in environment 4 and 5) and note that, in these two settings, ensembles perform better.

[The particular 8M-parameter ESM2 model that the authors used sounds incredibly small. It is not clear if that is a necessity (and if so, why), or if it turns out to somehow be sufficient and as good as the next level ESM2 model.]

We do not want to measure the contribution of the backbone, but rather compare models at a fixed backbone. Running bigger models is extremely costly computationally, especially in our framework as we are required to run 60 experiments per algorithm. In the short time available for the rebuttal, we provide the results for a larger 35M ESM. The results included in Table 7 show that there is little benefit from increasing the model's size for our benchmark dataset. Indeed, we are in a low data regime and finetuning overparameterized models in this regime can lead to overfitting.

[The authors presented a closely related work at an ICML symposium this summer. Was the main change since that presentation the inclusion of additional analysis methods, including the ensembling, or was the dataset itself improved or modified and if so how?]

• We moved from an internal dataset to a fully public dataset as our goal was to release the benchmark to the wider community.
• We explored and extended the available DG benchmark to recent state-of-the-art models that rely on pretrained models and ensembling and observed that these models outperform invariant based models.
• We also conducted a more thorough experimentation of domain shifts, including a zero-shot generalization setting (env 4), to measure the performance of our models and a downstream evaluation on binding measurements (env 5).
• Finally, we added protein structures in the dataset and added new models: a pretrained graph neural network and a bigger, 35M-parameter ESM.

[Honegger2001] Honegger, Annemarie, and Andreas PluÈckthun. "Yet another numbering scheme for immunoglobulin variable domains: an automatic modeling and analysis tool." Journal of molecular biology 309.3 (2001): 657-670.

We thank the reviewer for the thoughtful review. We address the concerns below and complement it with a revision of our paper.

[the construction of the data, the characterization and quality of the dataset, as well as the ability to select models and generalize from it in a drug discovery setting, remain unclear.]

Please see our updated Section A.5, where we evaluate our dataset with regards to six antibody properties:

• naturalness and diamond score on both heavy and light chains (Figure 9) which measure the closeness of an antibody to OAS; the largest publicly available human antibody database.
• hydrophobicity, charge and aromaticity from the design sequences (Figure 9), antibody properties which contribute to binding.
• humanness (Figure 8), leveraging the scoring framework OASis provided by BioPhi; we achieve ranges expected for a synthetic design library.

Our analysis of these scores confirms that the synthetic designs in Antibody DomainBed conform with properties of natural antibodies, similarly to a point extensively explored in the original publication of the generative model we use (Frey et al. 2023).

Regarding the model selection and generalization, our objective was to

1. explore the ability of domain generalization algorithms to deliver as expected in biological datasets and
2. determine whether the results are sufficient and encourage further developments of DG algorithms that take into account drug discovery pipelines.

We observe encouraging results for generalizing to actual binding measurements as pointed in our response to all reviewers. We hope that by providing a complete simulation and evaluation framework we bring the two communities closer together which may initiate new progress to the benefit of both fields.

[ it is not clear if the selection process used here ends up being educational for applications in real world settings.] [Do the authors have unpublished evidence of the usefulness of the models selected by this approach in actual drug discovery settings compared to a baseline ERM approach carefully implemented?]

To verify the usefulness of $\Delta\Delta G$ predictions for better binding affinity prediction, we include a new environment (env 5), consisting of real wet-lab measurements for the HER2 target. A seqCNN classifier trained on wet-lab measurements data on other 4 internal targets (unrelated to HER2) achieves an accuracy of 64.0%, precision 49.0% and recall 51.0% which reconfirms the fundamental conundrum of drug design, namely, the ability to predict behaviour of new therapeutics on unseen targets. Moreover, training on $\Delta\Delta G$ predictions included in env 0-4 proves to be helpful in predicting the binding affinity for the new env 5. Namely, pure ERM training on $\Delta\Delta G$ scores, achieves accuracy: 76.0 +/- 5.7%, precision: 99.7 +/- 0.3%, recall: 83.2 +/- 4.9% on environment 5, and the DG algorithms follow, with CORAL acheiving up to 88% accuracy after ensembling. Please see Table 6 and 7 in the revision for a detailed overview of the results.

[It does not help that the presented characterization has jargon specialized to the field and terms like MMD (maximal mean discrepancy?)]

We detailed in our revised manuscript the meaning and interpretation of MMD (maximum mean discrepancy), a distance measure between distributions, and hope this will be clear to future readers.

[saliency visualization that are not clarified and have scales that are illegible in the axis labels of Fig 4 and 6.]

Appologies for the ambiguity. The x-axis in the saliency figure represents the position of an amino acid in the antibody sequence, after being aligned to a uniform numbering scheme, AHo [Honegger2001].

[antibodies vs small molecule applications]

We removed this statement from the manuscript. It was not our intention to compare the pipeline processes between small vs. large molecules in early-stage research, which are affected by historical reasons, resource allocation, and the speed benefits of animal immunization in the case of antibodies. We now refer only to proteins posing unique challenges, as suggested. We also believe that with the increased interest in therapeutic antibodies in biotech, there is an increasing gap between ML and this data modality as it is still underexplored by the community.