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

TL;DR

This work introduces Antibody DomainBed, the first large-scale domain generalization (DG) benchmark for therapeutic antibody design under distribution shifts encountered in active drug design. By treating design rounds as environments, it systematically evaluates invariance-based and ensemble DG methods, finding that foundational protein models and model ensembling yield the strongest out-of-distribution performance across most environments. The authors construct a realistic dataset and open-source toolkit aligned with DomainBed, including sequence- and structure-aware representations (SeqCNN, Finetuned ESM2, GearNet) and five design-round environments, with labels derived from DeltaDeltaG proxies. While DG methods generally improve robustness, results on a completely new antigen (Env4/Env5) remain challenging, highlighting ongoing questions about environment design and causal feature discovery. Overall, Antibody DomainBed shows DG’s promise for stabilizing protein-property predictions amid distribution shifts in real-world antibody design pipelines and provides a valuable resource for future method development in this domain.

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.

Figures (12)

• Figure 1: Stabilizing interactions between antibody and antigen. Antibody Onartuzumab (pink) binds to MET (green and blue), a lung cancer antigen, on the cell surface. The strength of binding is determined by the binding site of the antibody interacting with the antigen, boxed in white.
• Figure 2: Active ML-guided design of antibodies against specific antigens involves: (1) developing generative models to create novel antibodies from seed antibodies, (2) selecting promising designs with predictive models, (3) experimentally validating the designs, and (4) updating models with new data for the next cycle. Each cycle may vary in targets, generative models, and experimental assays.
• Figure 3: Antibody Domainbed environments. Left: edit distance (sequence similarity) between designs and seeds. Right: binding properties per generative model.
• Figure 4: MMD in the learned features of ESM between every pair of environments. DG algorithms result in features that are significantly more uniform across environments.
• Figure 5: Saliency visualizations for ERM, SMA, and IRM on the antibody residue positions for heavy and light chains. Colors of bars represent functional segments. Relative to ERM-SMA and IRM, ERM displays muted behavior in the regions known to interact with the antigen paratope.