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

TL;DR

This paper applies domain generalization techniques to improve the robustness of antibody-antigen interaction predictions across different design cycles, addressing distribution shifts in therapeutic protein design.

Contribution

It introduces a DG benchmark for antibody design, evaluates foundational models and ensembling, and publicly releases related datasets and code.

Findings

Foundational models and ensembling enhance out-of-distribution prediction accuracy.

Explicitly modeling distribution shifts improves robustness in antibody design.

Public dataset and code facilitate further research in OOD generalization.

Abstract

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.