Antibody DomainBed: Out-of-Distribution Generalization in Therapeutic Protein Design
This work addresses robustness issues in drug design for researchers, though it is incremental as it adapts existing domain generalization methods to a new application.
The study tackled the challenge of distribution shifts in machine learning-guided therapeutic protein design by applying domain generalization methods to predict antibody-antigen interaction stability across five design cycle domains, finding that foundational models and ensembling enhance out-of-distribution predictive performance.
Machine learning (ML) has demonstrated significant promise in accelerating drug design. Active ML-guided optimization of therapeutic molecules typically relies on a surrogate model predicting the target property of interest. The model predictions are used to determine which designs to evaluate in the lab, and the model is updated on the new measurements to inform the next cycle of decisions. A key challenge is that the experimental feedback from each cycle inspires changes in the candidate proposal or experimental protocol for the next cycle, which lead to distribution shifts. To promote robustness to these shifts, we must 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.