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

Active ML-guided design of therapeutic molecules typically relies on a surrogate model predicting the property to be optimized. The model predictions are used to determine which designs to evaluate in the lab, and the model is updated on the new lab measurements to inform the next cycle of decisions. A key challenge arises from distribution shifts, which occur when the experimental feedback from each cycle inspires changes in the candidate proposal or experimental protocol for the next cycle. To promote robustness to these shifts, we account for them explicitly in the model training. We apply domain generalization (DG) methods to classify the stability of interactions between an antibody and antigen across five domains defined by design cycles. Our results suggest that foundational models and ensembling improve predictive performance on out-of-distribution domains. We publicly release our codebase extending the DG benchmark "DomainBed" and the associated dataset of antibody sequences and structures emulating distribution shifts across design cycles.