Source-linked AI summary
Learning Task-Specific Antibody Representations via Function-Aware Masking
Ayan Goel, Thomas A. Walton, Amirali Aghazadeh
TL;DR
Antibody MLMs rarely use masking as a functional inductive bias, despite antibodies having diverse task-relevant biological priors beyond CDRs. This paper introduces specialist and hybrid function-aware masking strategies that align mask placement with these priors. Specialist strategies improve aligned tasks, while hybrid strategies maintain broad generalization across functional objectives without added inference parameters.
Problem
Masking has rarely been leveraged as an inductive bias for antibody MLMs, while CDR-focused masking does not cover structural, developability, and other antibody functions.
Method
Function-aware masking uses sequence labels or model predictions to bias mask placement toward task-relevant residues, with specialist strategies targeting one prior and hybrid strategies mixing multiple priors.
Results
Specialist strategies improve aligned performance by up to 14% on structure-related tasks and up to 5.9× on CDR-related tasks, while hybrid strategies consistently outperform random masking across diverse functional tasks.
Takeaways & Limitations
Informed mask placement provides a flexible, parameter-free way to impose functional inductive biases during antibody language-model training.
Takeaways & Limitations
The study evaluates a medium-sized approximately 38M-parameter RoFormer, requires per-residue labels, and uses manually designed hybrid mixing strategies.
Abstract
from arXiv · showhide
Antibody-specific language models pretrained via masked language modeling (MLM) learn representations that are critical for downstream sequence design and property prediction tasks. Yet, the corruption process itself is rarely leveraged as a source of inductive bias during pretraining. While preferentially masking complementarity-determining regions (CDRs) improves binding-related predictions, antibodies possess diverse biological priors over a variety of functions. Herein, we introduce function-aware masking, a family of pretraining algorithms that align mask placement with specific functional priors (e.g., from IMGT annotations or structure predictions) to shape the learned representation space. We show that these specialist masking strategies significantly improve performance on their respective objectives, yielding up to a 14% gain on structure-related tasks and up to a 5.9x improvement on CDR-related tasks. To further improve performance across multiple functional axes, we develop hybrid masking strategies that integrate multiple priors, balancing reconstruction over binding, structural, and biophysical objectives. Our results demonstrate that informed mask placement provides a parameter-free mechanism for imposing functional inductive biases in antibody language model training.
1 Introduction
Antibody language models learn useful sequence representations, but uniform masking does not target the sparse regions most consequential to antibody function. Function-aware masking aligns mask placement with functional priors, improving aligned tasks and combining priors to support multiple objectives.
- Antibody-specific language models produce representations that generalize across sequence design and optimization tasks.
- Preferentially masking CDRs improves transfer on CDR infilling, binding affinity, and binding specificity tasks.
- Therapeutic antibody properties such as folding stability, aggregation resistance, and mutational robustness are governed partly by framework residues outside CDRs.
- Function-aware masking scores residues using sequence labels or model predictions and preferentially masks positions associated with a target functional task.The resulting models remain sequence-only at inference and incur no additional parameter cost.
- Specialist strategies improve aligned tasks by up to 14% on structure-related objectives and up to 5.9× on CDR-related objectives.
- Hybrid masking strategies combine functional priors to improve representation quality across multiple functional axes and mitigate degradation on unaligned tasks.
2 Background
Masked language modeling uses corruption and reconstruction to shape protein representations, while antibody function reflects multiple specialized biological regions and evolutionary processes. Prior antibody masking methods mainly target CDRs, motivating broader functional masking strategies.
- Random masking is commonly applied uniformly across sequence positions at a fixed probability, typically 15%.
- Structural mask placement has improved representation generalization on protein extrapolation tasks.
- Antigen recognition is primarily mediated by six CDRs, with paratopes forming direct target interfaces within these variable loops.
- Somatic hypermutation affects binding interfaces and corresponding framework regions during affinity maturation.
- Earlier antibody masking studies focused on CDR3 or combined CDR and framework masking, improving residue recovery and binding-related predictions.
3 Methods
The method uses biological priors to direct antibody MLM mask placement through specialist strategies targeting individual properties and hybrid strategies mixing multiple priors. Mask probabilities are calibrated to a fixed 15% budget, with hybrid schedules changing the sampling distribution during training.
- Function-aware masking comprises specialist strategies focused on one task and hybrid strategies that mix priors from multiple tasks.
- Per-residue labels and predictions, including CDR locations, paratope probabilities, germline mutation status, and structural topology, are precomputed for targeted sampling.
- 3.1 Specialist Strategies: Weight-based specialist strategies convert residue relevance weights into masking probabilities calibrated to an expected 15% budget.
- 3.1 Specialist Strategies: Masks are sampled independently for weight-based strategies, whereas SPAN, STRUCTURE, and STRUCTURE-LR place masks iteratively until 15% of positions are selected.
- 3.1 Specialist Strategies: The seven specialist strategies target CDRs, contiguous spans, paratopes, germline-mutated residues, affinity-maturation intersections, structural neighborhoods, or long-range contacts.STRUCTURE uses five nearest spatial neighbors, while STRUCTURE-LR restricts couplings to sequence separations greater than four indices.
- 3.2 Hybrid Strategies: Hybrid strategies dynamically sample among specialists using a time-varying mixing distribution that begins broadly and can favor better-performing specialists.
- 3.2 Hybrid Strategies: The evaluated hybrids vary scheduling, weighting, reversal, batching, and warm-start choices for combining masking strategies.
4 Experiments
The experiments pretrain identical medium-sized RoFormer models on a large antibody heavy-chain dataset, then freeze the encoders and evaluate representations with downstream probes. Metrics cover generative recovery, binding, and other distinct antibody functions.
- All strategies use a 38.1M-parameter RoFormer setup trained for 125,000 steps, except HYBRID-WARMSTART, which adds 50,000 steps.
- The dataset contains 497,309 filtered heavy-chain variable-domain sequences from the Observed Antibody Space.
- Evaluation freezes each pretrained encoder and extracts last-layer embeddings for downstream probes.
- The evaluation uses seven metrics spanning generative CDR3 recovery, paratope identification, and other distinct biological properties.CDR3 infilling is evaluated zero-shot, while the remaining metrics use linear or bilinear probing.
5 Results
Specialist masking strategies generally align antibody representations with their targeted biological functions, while hybrid strategies improve consistency across multiple functional axes. Mask placement and hybrid mixing choices both shape downstream performance.
- Specialist masking: Each specialist masking strategy achieved the highest score on at least one metric, except structure-based strategies.
- Specialist masking: CDR masking performed best on CDR3 infilling, while INTERFACE and GERMLINE performed best overall across multiple structure- and paratope-related tasks.
- Specialist masking: INTERSECTION masking improved paratope- and CDR-related tasks but did not retain the developability and structural gains of INTERFACE and GERMLINE.The passage attributes this gap to concentrating masks in fewer sequence locations.
- Specialist masking: SPAN masking generalized well across paratope, CDR3, and developability tasks but fell short on tasks specialized by INTERFACE and GERMLINE.This comparison indicates that mask placement mattered more than mask shape for those tasks.
- Specialist masking: With the exception of STRUCTURE masking, every specialist improved over RANDOM masking on multiple metrics, with INTERFACE and GERMLINE gaining across nearly all functional axes.
- Hybrid masking: Hybrid strategies combine multiple biological priors through stochastic mixtures and training curricula to improve robustness when specialist performance degrades on unaligned tasks.
- Hybrid masking: Hybrid strategies achieved higher average task rank, lower overall task variance, and broader consistency than narrow specialists, although they seldom ranked first on individual tasks.HYBRID-REVERSE ranked best on three of seven tasks, while HYBRID performed worst overall and some hybrids surpassed RANDOM controls.
6 Conclusion
The paper introduces function-aware masking to impose biological inductive biases during antibody language-model pretraining. Specialist strategies improve aligned tasks, while hybrid strategies preserve broader generalization across functional axes without adding inference-time parameters.
- Conclusion: Function-aware masking uses biological priors to impose targeted inductive biases during antibody language-model pretraining.
- Conclusion: Hybrid masking combines multiple functional priors to mitigate trade-offs from narrow specialization and maintain broad generalization across diverse functional tasks.
- Limitations: Function-aware masking requires per-residue labels, evaluations used a medium-sized RoFormer model of approximately 38M parameters, and hybrid strategies were manually designed.Larger-scale testing and systematic optimization of mixing distributions may yield additional gains.
A Additional Results
This section presents the full results table, combining experiments from Tables 1 and 2.
- The full results table combines experiments from Table 1 and Table 2.
B Implementation Details
The implementation details describe specialist and hybrid masking schemes and visualize empirical per-position mask probabilities for thirteen tested methods.
- Thirteen masking methods are visualized through empirical per-position mask probabilities, with CDRs highlighted for reference.
- STRUCTURE predicts antibody structures with IgFold, builds per-residue k-nearest-neighbor graphs using Euclidean Cα−Cα distances, and masks structural neighborhoods.STRUCTURE uses k = 5 nearest neighbors and samples only a few masks to preserve structural context for decoding.
B.2 Hybrid Masking Strategies
Hybrid masking strategies combine specialist strategies through timestep-dependent stochastic mixtures, with curricula varying how emphasis shifts during training.
- Hybrid (default): The default HYBRID curriculum starts with greater RANDOM and SPAN emphasis, shifts toward CDR and INTERFACE, then broadly covers strategies after step 40,000.
- Hybrid-stretched: HYBRID-STRETCHED delays the default curriculum’s transitions by extending them over more training steps while ending with the same distribution.
- Hybrid-reverse: HYBRID-REVERSE uses the HYBRID-STRETCHED strategy in reverse and stretches its schedule across more training steps.
- Hybrid-weighted: HYBRID-WEIGHTED assigns greater emphasis to the best-performing specialist, INTERFACE.
- Hybrid-perbatch: HYBRID-PERBATCH follows HYBRID’s probability distribution but samples one fixed specialist strategy for each batch, reducing computational overhead.
- Hybrid-warmstart: HYBRID-WARMSTART continues training an INTERFACE model for 50,000 additional steps using a specified timestep-dependent mixing distribution.
C Experimental Setup
Each from-scratch run takes approximately five hours on one RTX A6000, while warm-start training adds about two hours and IgFold prediction costs approximately 240 GPU-hours.
- ≈5 hours per from-scratch run on one RTX A6000, with HYBRID-WARMSTART adding ≈2 hours and IgFold predictions requiring ≈240 GPU-hours.The IgFold cost uses 8 GPUs in parallel.
C.1 Probe Details
The paper evaluates frozen antibody-language-model representations using task-specific linear probes, with fixed data splits and specialized probe configurations. Paratope labels use a separately trained teacher, while structural signals come from IgFold predictions.
- Probe Details: All Section 4 metrics except zero-shot CDR3 infilling use linear probes trained on frozen last-layer representations.Task-specific heads are trained without updating the pretrained encoder.
- Data Splits: Paratope and developability probes use TDC default splits, while structural probes split by PDB entry to prevent complex-chain leakage.Heavy and light chains from the same complex remain in the same split.
- Teacher Models: Paratope labels are generated by an AntiBERTa2-based per-token teacher trained only on the training split, with predictions restricted to pretraining sequences.No evaluation sequence is labeled by the teacher during teacher training.
- Teacher Models: Structural supervision uses IgFold predictions rather than a structure teacher trained in this work.IgFold neighborhood predictions were compared with SAbDab crystal structures for validation.
- Probe Configurations: Probe heads use task-specific architectures and metrics, including token classification for paratopes, bilinear residue-pair prediction for contacts, and a linear structural probe.The listed metrics include AUPRC and MCC for paratopes, AUROC and long-range precision-at-L for contacts, and Spearman ρ for structure.